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

Author: timonegk

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

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"fmt"
 	"reflect"
+	"regexp"
 	"strings"
 	"time"
 
@@ -15,26 +16,27 @@ import (
 
 // Config Errors
 var (
-	ErrMappingRequired        = errors.New("logfmt mapping is required")
+	ErrMappingOrRegexRequired = errors.New("logfmt mapping or regex is required")
 	ErrEmptyLogfmtStageConfig = errors.New("empty logfmt stage configuration")
 )
 
 // LogfmtConfig represents a logfmt Stage configuration
 type LogfmtConfig struct {
-	Mapping map[string]string `alloy:"mapping,attr"`
+	Mapping map[string]string `alloy:"mapping,attr,optional"`
 	Source  string            `alloy:"source,attr,optional"`
+	Regex   string            `alloy:"regex,attr,optional"`
 }
 
 // validateLogfmtConfig validates a logfmt stage config and returns an inverse mapping of configured mapping.
 // Mapping inverse is done to make lookup easier. The key would be the key from parsed logfmt and
 // value would be the key with which the data in extracted map would be set.
-func validateLogfmtConfig(c *LogfmtConfig) (map[string]string, error) {
+func validateLogfmtConfig(c *LogfmtConfig) (map[string]string, *regexp.Regexp, error) {
 	if c == nil {
-		return nil, ErrEmptyLogfmtStageConfig
+		return nil, nil, ErrEmptyLogfmtStageConfig
 	}
 
-	if len(c.Mapping) == 0 {
-		return nil, ErrMappingRequired
+	if len(c.Mapping) == 0 && len(c.Regex) == 0 {
+		return nil, nil, ErrMappingOrRegexRequired
 	}
 
 	inverseMapping := make(map[string]string)
@@ -46,28 +48,35 @@ func validateLogfmtConfig(c *LogfmtConfig) (map[string]string, error) {
 		inverseMapping[v] = k
 	}
 
-	return inverseMapping, nil
+	re, err := regexp.Compile(c.Regex)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return inverseMapping, re, nil
 }
 
 // logfmtStage sets extracted data using logfmt parser
 type logfmtStage struct {
 	cfg            *LogfmtConfig
 	inverseMapping map[string]string
+	regex          regexp.Regexp
 	logger         log.Logger
 }
 
 // newLogfmtStage creates a new logfmt pipeline stage from a config.
 func newLogfmtStage(logger log.Logger, config LogfmtConfig) (Stage, error) {
 	// inverseMapping would hold the mapping in inverse which would make lookup easier.
 	// To explain it simply, the key would be the key from parsed logfmt and value would be the key with which the data in extracted map would be set.
-	inverseMapping, err := validateLogfmtConfig(&config)
+	inverseMapping, regex, err := validateLogfmtConfig(&config)
 	if err != nil {
 		return nil, err
 	}
 
 	return toStage(&logfmtStage{
 		cfg:            &config,
 		inverseMapping: inverseMapping,
+		regex:          *regex,
 		logger:         log.With(logger, "component", "stage", "type", "logfmt"),
 	}), nil
 }
@@ -98,13 +107,22 @@ func (j *logfmtStage) Process(labels model.LabelSet, extracted map[string]interf
 		return
 	}
 	decoder := logfmt.NewDecoder(strings.NewReader(*input))
-	extractedEntriesCount := 0
+	mappingExtractedEntriesCount := 0
+	regexExtractedEntriesCount := 0
 	for decoder.ScanRecord() {
 		for decoder.ScanKeyval() {
+			// handle "mapping"
 			mapKey, ok := j.inverseMapping[string(decoder.Key())]
 			if ok {
 				extracted[mapKey] = string(decoder.Value())
-				extractedEntriesCount++
+				mappingExtractedEntriesCount++
+			} else if j.regex.String() != "" {
+				// handle "regex"
+				fmt.Println(j.regex.String(), string(decoder.Key()))
+				if j.regex.MatchString(string(decoder.Key())) {
+					extracted[string(decoder.Key())] = string(decoder.Value())
+					regexExtractedEntriesCount++
+				}
 			}
 		}
 	}
@@ -114,8 +132,11 @@ func (j *logfmtStage) Process(labels model.LabelSet, extracted map[string]interf
 		return
 	}
 
-	if extractedEntriesCount != len(j.inverseMapping) {
-		level.Debug(j.logger).Log("msg", fmt.Sprintf("found only %d out of %d configured mappings in logfmt stage", extractedEntriesCount, len(j.inverseMapping)))
+	if mappingExtractedEntriesCount != len(j.inverseMapping) {
+		level.Debug(j.logger).Log("msg", fmt.Sprintf("found only %d out of %d configured mappings in logfmt stage", mappingExtractedEntriesCount, len(j.inverseMapping)))
+	}
+	if regexExtractedEntriesCount > 0 {
+		level.Debug(j.logger).Log("msg", fmt.Sprintf("found %d mappings via regex in logfmt stage", regexExtractedEntriesCount))
 	}
 	level.Debug(j.logger).Log("msg", "extracted data debug in logfmt stage", "extracted data", fmt.Sprintf("%v", extracted))
 }

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

@@ -26,6 +26,25 @@ stage.logfmt {
 		source  = "extra"
 }`
 
+var testLogfmtAlloyRegex = `
+stage.logfmt {
+    regex = "pod_.*"
+}
+`
+
+var testLogfmtAlloyRegexAll = `
+stage.logfmt {
+    regex = ".*"
+}
+`
+
+var testLogfmtAlloyRegexAndMapping = `
+stage.logfmt {
+    mapping = { "out" = "message", "app" = ""}
+    regex = "(app|duration)"
+}
+`
+
 func TestLogfmt(t *testing.T) {
 	var testLogfmtLogLine = `
 		time=2012-11-01T22:08:41+00:00 app=loki	level=WARN duration=125 message="this is a log line" extra="user=foo""
@@ -54,6 +73,35 @@ func TestLogfmt(t *testing.T) {
 				"user":  "foo",
 			},
 		},
+		"successfully extract regex values from logfmt": {
+			testLogfmtAlloyRegex,
+			`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 via regex from logfmt": {
+			testLogfmtAlloyRegexAll,
+			testLogfmtLogLine,
+			map[string]interface{}{
+				"time":     "2012-11-01T22:08:41+00:00",
+				"app":      "loki",
+				"level":    "WARN",
+				"duration": "125",
+				"message":  "this is a log line",
+				"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",
+      },
+    },
 	}
 
 	for testName, testData := range tests {
@@ -81,7 +129,7 @@ func TestLogfmtConfigValidation(t *testing.T) {
 		"no mapping": {
 			LogfmtConfig{},
 			0,
-			ErrMappingRequired,
+			ErrMappingOrRegexRequired,
 		},
 		"valid without source": {
 			LogfmtConfig{
@@ -108,7 +156,7 @@ func TestLogfmtConfigValidation(t *testing.T) {
 	for tName, tt := range tests {
 		tt := tt
 		t.Run(tName, func(t *testing.T) {
-			got, err := validateLogfmtConfig(&tt.config)
+			got, _, err := validateLogfmtConfig(&tt.config)
 			if tt.err != nil {
 				assert.EqualError(t, err, tt.err.Error())
 			} else {