Feedbacks fixes

Signed-off-by: Ashmita Bohara <[email protected]>

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])

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)
 }