fix: integrate change requests from review

Signed-off-by: Timon Engelke <timon.engelke@inovex.de>

Author: timonegk

internal/component/loki/process/stages/json.go

@@ -15,10 +15,10 @@ import (
 // Config Errors
 const (
 	ErrExpressionsOrRegexRequired = "JMES expressions or regex is required"
-	ErrCouldNotCompileJMES       = "could not compile JMES expression"
-	ErrEmptyJSONStageConfig      = "empty json stage configuration"
-	ErrEmptyJSONStageSource      = "empty source"
-	ErrMalformedJSON             = "malformed json"
+	ErrCouldNotCompileJMES        = "could not compile JMES expression"
+	ErrEmptyJSONStageConfig       = "empty json stage configuration"
+	ErrEmptyJSONStageSource       = "empty source"
+	ErrMalformedJSON              = "malformed json"
 )
 
 // JSONConfig represents a JSON Stage configuration
@@ -143,58 +143,25 @@ func (j *jsonStage) processEntry(extracted map[string]any, entry *string) error
 		return errors.New(ErrMalformedJSON)
 	}
 
-	for n, e := range j.expressions {
-		r, err := e.Search(data)
+	for name, expr := range j.expressions {
+		rawResult, err := expr.Search(data)
 		if err != nil {
 			if Debug {
 				level.Debug(j.logger).Log("msg", "failed to search JMES expression", "err", err)
 			}
 			continue
 		}
-
-		switch r.(type) {
-		case float64:
-			// All numbers in JSON are unmarshaled to float64.
-			extracted[n] = r
-		case string:
-			extracted[n] = r
-		case bool:
-			extracted[n] = r
-		case nil:
-			extracted[n] = nil
-		default:
-			// If the value wasn't a string or a number, marshal it back to json
-			jm, err := json.Marshal(r)
-			if err != nil {
-				if Debug {
-					level.Debug(j.logger).Log("msg", "failed to marshal complex type back to string", "err", err)
-				}
-				continue
-			}
-			extracted[n] = string(jm)
+		value, err := j.simplifyType(rawResult)
+		if err == nil {
+			extracted[name] = value
 		}
 	}
 	if j.regex.String() != "" {
-		for key, value := range data {
+		for key, rawValue := range data {
 			if j.regex.MatchString(key) {
-				switch value.(type) {
-				case float64:
+				value, err := j.simplifyType(rawValue)
+				if err == nil {
 					extracted[key] = value
-				case string:
-					extracted[key] = value
-				case bool:
-					extracted[key] = value
-				case nil:
-					extracted[key] = nil
-				default:
-					jm, err := json.Marshal(value)
-					if err != nil {
-						if Debug {
-							level.Debug(j.logger).Log("msg", "failed to marshal complex type back to string", "err", err)
-						}
-						continue
-					}
-					extracted[key] = string(jm)
 				}
 			}
 		}
@@ -205,6 +172,31 @@ func (j *jsonStage) processEntry(extracted map[string]any, entry *string) error
 	return nil
 }
 
+// extractWithType returns the value if it's a simple type (string, number, bool),
+// otherwise, it returns it as a JSON string
+func (j *jsonStage) simplifyType(value any) (any, error) {
+	switch value.(type) {
+	case float64:
+		return value, nil
+	case string:
+		return value, nil
+	case bool:
+		return value, nil
+	case nil:
+		return nil, nil
+	default:
+		// If the value wasn't a string or a number, marshal it back to json
+		jm, err := json.Marshal(value)
+		if err != nil {
+			if Debug {
+				level.Debug(j.logger).Log("msg", "failed to marshal complex type back to string", "err", err)
+				return nil, err
+			}
+		}
+		return string(jm), nil
+	}
+}
+
 // Cleanup implements Stage.
 func (*jsonStage) Cleanup() {
 	// no-op

internal/component/loki/process/stages/json_test.go

@@ -93,35 +93,35 @@ func TestPipeline_JSON(t *testing.T) {
 		},
 		"successfully extract regex values from json": {
 			testJSONAlloyRegex,
-      `{"time":"2012-11-01T22:08:41+00:00", "pod_name": "my-pod-123", "pod_label": "my-label"}`,
+			`{"time":"2012-11-01T22:08:41+00:00", "pod_name": "my-pod-123", "pod_label": "my-label"}`,
 			map[string]interface{}{
-        "pod_name": "my-pod-123",
-        "pod_label": "my-label",
+				"pod_name":  "my-pod-123",
+				"pod_label": "my-label",
 			},
 		},
 		"successfully extract all values from json via regex": {
 			testJSONAlloyRegexAll,
 			testJSONLogLine,
 			map[string]interface{}{
-        "time": "2012-11-01T22:08:41+00:00",
-        "app": "loki",
-        "component": `["parser","type"]`,
-        "level": "WARN",
-        "nested": `{"child":"value"}`,
-        "duration": float64(125),
-        "message": "this is a log line",
-				"extra": "{\"user\":\"marco\"}",
+				"time":      "2012-11-01T22:08:41+00:00",
+				"app":       "loki",
+				"component": `["parser","type"]`,
+				"level":     "WARN",
+				"nested":    `{"child":"value"}`,
+				"duration":  float64(125),
+				"message":   "this is a log line",
+				"extra":     "{\"user\":\"marco\"}",
+			},
+		},
+		"successfully extract values with expressions and regex from json": {
+			testJSONAlloyExpressionsAndRegex,
+			testJSONLogLine,
+			map[string]interface{}{
+				"out":      "this is a log line",
+				"app":      "loki",
+				"duration": float64(125),
 			},
 		},
-    "successfully extract values with expressions and regex from json": {
-      testJSONAlloyExpressionsAndRegex,
-      testJSONLogLine,
-      map[string]interface{}{
-        "out": "this is a log line",
-        "app": "loki",
-        "duration": float64(125),
-      },
-    },
 	}
 
 	for testName, testData := range tests {

internal/component/loki/process/stages/logfmt.go

@@ -115,9 +115,9 @@ func (j *logfmtStage) Process(labels model.LabelSet, extracted map[string]any, t
 			if ok {
 				extracted[mapKey] = string(decoder.Value())
 				mappingExtractedEntriesCount++
-			} else if j.regex.String() != "" {
-				// handle "regex"
-				fmt.Println(j.regex.String(), string(decoder.Key()))
+			}
+			// handle "regex"
+			if j.regex.String() != "" {
 				if j.regex.MatchString(string(decoder.Key())) {
 					extracted[string(decoder.Key())] = string(decoder.Value())
 					regexExtractedEntriesCount++

internal/component/loki/process/stages/logfmt_test.go

@@ -93,15 +93,15 @@ func TestLogfmt(t *testing.T) {
 				"extra":    "user=foo",
 			},
 		},
-    "successfully extract values with expressions and regex from logfmt": {
-      testLogfmtAlloyRegexAndMapping,
-      testLogfmtLogLine,
-      map[string]interface{}{
-        "out": "this is a log line",
-        "app": "loki",
-        "duration": "125",
-      },
-    },
+		"successfully extract values with expressions and regex from logfmt": {
+			testLogfmtAlloyRegexAndMapping,
+			testLogfmtLogLine,
+			map[string]interface{}{
+				"out":      "this is a log line",
+				"app":      "loki",
+				"duration": "125",
+			},
+		},
 	}
 
 	for testName, testData := range tests {