diff --git a/cmd/anonymizer/app/query/query.go b/cmd/anonymizer/app/query/query.go index c6b261bde2e..ffbf2126d9e 100644 --- a/cmd/anonymizer/app/query/query.go +++ b/cmd/anonymizer/app/query/query.go @@ -50,14 +50,14 @@ func New(addr string) (*Query, error) { }, nil } -// isNotFoundErr is a helper function called when the trace isn't found -func isNotFoundErr(err error) ([]model.Span, error) { +// unwrapNotFoundErr is a conversion function +func unwrapNotFoundErr(err error) error { if s, _ := status.FromError(err); s != nil { if s.Message() == spanstore.ErrTraceNotFound.Error() { - return nil, spanstore.ErrTraceNotFound + return spanstore.ErrTraceNotFound } } - return nil, fmt.Errorf("failed to fetch the provided trace id: %w", err) + return err } // QueryTrace queries for a trace and returns all spans inside it @@ -71,13 +71,13 @@ func (q *Query) QueryTrace(traceID string) ([]model.Span, error) { TraceID: mTraceID, }) if err != nil { - isNotFoundErr(err) + return nil, unwrapNotFoundErr(err) } var spans []model.Span for received, err := stream.Recv(); err != io.EOF; received, err = stream.Recv() { if err != nil { - isNotFoundErr(err) + return nil, unwrapNotFoundErr(err) } for i := range received.Spans { spans = append(spans, received.Spans[i]) diff --git a/cmd/anonymizer/app/writer/writer.go b/cmd/anonymizer/app/writer/writer.go index b8f1ee338d5..5d04e61df53 100644 --- a/cmd/anonymizer/app/writer/writer.go +++ b/cmd/anonymizer/app/writer/writer.go @@ -34,7 +34,7 @@ type Config struct { CapturedFile string `yaml:"captured_file" name:"captured_file"` AnonymizedFile string `yaml:"anonymized_file" name:"anonymized_file"` MappingFile string `yaml:"mapping_file" name:"mapping_file"` - AnonymizerOpts anonymizer.Options `yaml:"anonymizer_opts" name:"anonymizer_opts"` + AnonymizerOpts anonymizer.Options `yaml:"anonymizer" name:"anonymizer"` } // Writer is a span Writer that obfuscates the span and writes it to a JSON file. @@ -129,6 +129,8 @@ func (w *Writer) WriteSpan(msg *model.Span) error { if w.spanCount >= w.config.MaxSpansCount { w.logger.Info("Saved enough spans, exiting...") + w.Close() + os.Exit(0) } return nil @@ -141,5 +143,4 @@ func (w *Writer) Close() { w.anonymizedFile.WriteString("\n]\n") w.anonymizedFile.Close() w.anonymizer.SaveMapping() - os.Exit(0) }