feat(loki.process): add regex field to json stage

Author: timonegk

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

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"fmt"
 	"reflect"
+	"regexp"
 
 	"github.com/go-kit/log"
 	"github.com/grafana/alloy/internal/runtime/logging/level"
@@ -13,32 +14,33 @@ import (
 
 // Config Errors
 const (
-	ErrExpressionsRequired  = "JMES expression is required"
-	ErrCouldNotCompileJMES  = "could not compile JMES expression"
-	ErrEmptyJSONStageConfig = "empty json stage configuration"
-	ErrEmptyJSONStageSource = "empty source"
-	ErrMalformedJSON        = "malformed json"
+	ErrExpressionsOrRegexRequired = "JMES expressions or regex is required"
+	ErrCouldNotCompileJMES       = "could not compile JMES expression"
+	ErrEmptyJSONStageConfig      = "empty json stage configuration"
+	ErrEmptyJSONStageSource      = "empty source"
+	ErrMalformedJSON             = "malformed json"
 )
 
 // JSONConfig represents a JSON Stage configuration
 type JSONConfig struct {
-	Expressions   map[string]string `alloy:"expressions,attr"`
+	Expressions   map[string]string `alloy:"expressions,attr,optional"`
+	Regex         string            `alloy:"regex,attr,optional"`
 	Source        *string           `alloy:"source,attr,optional"`
 	DropMalformed bool              `alloy:"drop_malformed,attr,optional"`
 }
 
 // validateJSONConfig validates a json config and returns a map of necessary jmespath expressions.
-func validateJSONConfig(c *JSONConfig) (map[string]jmespath.JMESPath, error) {
+func validateJSONConfig(c *JSONConfig) (map[string]jmespath.JMESPath, *regexp.Regexp, error) {
 	if c == nil {
-		return nil, errors.New(ErrEmptyJSONStageConfig)
+		return nil, nil, errors.New(ErrEmptyJSONStageConfig)
 	}
 
-	if len(c.Expressions) == 0 {
-		return nil, errors.New(ErrExpressionsRequired)
+	if len(c.Expressions) == 0 && len(c.Regex) == 0 {
+		return nil, nil, errors.New(ErrExpressionsOrRegexRequired)
 	}
 
 	if c.Source != nil && *c.Source == "" {
-		return nil, errors.New(ErrEmptyJSONStageSource)
+		return nil, nil, errors.New(ErrEmptyJSONStageSource)
 	}
 
 	expressions := map[string]jmespath.JMESPath{}
@@ -52,28 +54,36 @@ func validateJSONConfig(c *JSONConfig) (map[string]jmespath.JMESPath, error) {
 		}
 		expressions[n], err = jmespath.Compile(jmes)
 		if err != nil {
-			return nil, fmt.Errorf("%s: %w", ErrCouldNotCompileJMES, err)
+			return nil, nil, fmt.Errorf("%s: %w", ErrCouldNotCompileJMES, err)
 		}
 	}
-	return expressions, nil
+
+	re, err := regexp.Compile(c.Regex)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return expressions, re, nil
 }
 
 // jsonStage sets extracted data using JMESPath expressions
 type jsonStage struct {
 	cfg         *JSONConfig
 	expressions map[string]jmespath.JMESPath
+	regex       regexp.Regexp
 	logger      log.Logger
 }
 
 // newJSONStage creates a new json pipeline stage from a config.
 func newJSONStage(logger log.Logger, cfg JSONConfig) (Stage, error) {
-	expressions, err := validateJSONConfig(&cfg)
+	expressions, regex, err := validateJSONConfig(&cfg)
 	if err != nil {
 		return nil, err
 	}
 	return &jsonStage{
 		cfg:         &cfg,
 		expressions: expressions,
+		regex:       *regex,
 		logger:      log.With(logger, "component", "stage", "type", "json"),
 	}, nil
 }
@@ -164,6 +174,31 @@ func (j *jsonStage) processEntry(extracted map[string]interface{}, entry *string
 			extracted[n] = string(jm)
 		}
 	}
+	if j.regex.String() != "" {
+		for key, value := range data {
+			if j.regex.MatchString(key) {
+				switch value.(type) {
+				case float64:
+					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)
+				}
+			}
+		}
+	}
 	if Debug {
 		level.Debug(j.logger).Log("msg", "extracted data debug in json stage", "extracted_data", fmt.Sprintf("%v", extracted))
 	}

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

@@ -28,17 +28,36 @@ stage.json {
 
 stage.json {
     expressions = { "user" = "" }
-	source      = "extra"
+    source      = "extra"
 }`
 
+var testJSONAlloyRegex = `
+stage.json {
+  regex = "pod_.*"
+}
+`
+
+var testJSONAlloyRegexAll = `
+stage.json {
+  regex = ".*"
+}
+`
+
+var testJSONAlloyExpressionsAndRegex = `
+stage.json {
+  expressions = {"out" = "message", "app" = ""}
+  regex = "(app|duration)"
+}
+`
+
 var testJSONLogLine = `
 {
 	"time":"2012-11-01T22:08:41+00:00",
 	"app":"loki",
 	"component": ["parser","type"],
 	"level" : "WARN",
 	"nested" : {"child":"value"},
-    "duration" : 125,
+	"duration" : 125,
 	"message" : "this is a log line",
 	"extra": "{\"user\":\"marco\"}"
 }
@@ -72,6 +91,37 @@ func TestPipeline_JSON(t *testing.T) {
 				"user":  "marco",
 			},
 		},
+		"successfully extract regex values from json": {
+			testJSONAlloyRegex,
+      `{"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",
+			},
+		},
+		"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\"}",
+			},
+		},
+    "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 {
@@ -132,7 +182,7 @@ func TestJSONConfig_validate(t *testing.T) {
 		"no expressions": {
 			&JSONConfig{},
 			0,
-			errors.New(ErrExpressionsRequired),
+			errors.New(ErrExpressionsOrRegexRequired),
 		},
 		"invalid expression": {
 			&JSONConfig{
@@ -180,7 +230,7 @@ func TestJSONConfig_validate(t *testing.T) {
 	for tName, tt := range tests {
 		tt := tt
 		t.Run(tName, func(t *testing.T) {
-			got, err := validateJSONConfig(tt.config)
+			got, _, err := validateJSONConfig(tt.config)
 			if tt.err != nil {
 				assert.NotNil(t, err, "JSONConfig.validate() expected error = %v, but got nil", tt.err)
 			}