From 416e01ca8e8f6e606135ea7b18f9db3393345ede Mon Sep 17 00:00:00 2001
From: Vladyslav Diachenko <82767850+vlad-diachenko@users.noreply.github.com>
Date: Fri, 21 Jul 2023 15:33:33 +0300
Subject: [PATCH] [promtail] non_indexed_labels stage (#9986)

**What this PR does / why we need it**:

added non_indexed_labels stage to copy values from extracted data
tonon-indexed-labels field of the Entry.


**Special notes for your reviewer**:

**Checklist**
- [ ] Reviewed the
[`CONTRIBUTING.md`](https://github.com/grafana/loki/blob/main/CONTRIBUTING.md)
guide (**required**)
- [ ] Documentation added
- [ ] Tests updated
- [ ] `CHANGELOG.md` updated
- [ ] If the change is worth mentioning in the release notes, add
`add-to-release-notes` label
- [ ] Changes that require user attention or interaction to upgrade are
documented in `docs/sources/upgrading/_index.md`
- [ ] For Helm chart changes bump the Helm chart version in
`production/helm/loki/Chart.yaml` and update
`production/helm/loki/CHANGELOG.md` and
`production/helm/loki/README.md`. [Example
PR](https://github.com/grafana/loki/commit/d10549e3ece02120974929894ee333d07755d213)

---------

Signed-off-by: Vladyslav Diachenko <vlad.diachenko@grafana.com>
---
 clients/pkg/logentry/stages/labels.go         |  16 +-
 .../pkg/logentry/stages/nonindexedlabels.go   |  43 +++
 .../logentry/stages/nonindexedlabels_test.go  | 113 +++++++
 clients/pkg/logentry/stages/stage.go          | 287 +++++++++---------
 4 files changed, 305 insertions(+), 154 deletions(-)
 create mode 100644 clients/pkg/logentry/stages/nonindexedlabels.go
 create mode 100644 clients/pkg/logentry/stages/nonindexedlabels_test.go

diff --git a/clients/pkg/logentry/stages/labels.go b/clients/pkg/logentry/stages/labels.go
index 7d33061f0e..3a3486977d 100644
--- a/clients/pkg/logentry/stages/labels.go
+++ b/clients/pkg/logentry/stages/labels.go
@@ -63,23 +63,31 @@ type labelStage struct {
 
 // Process implements Stage
 func (l *labelStage) Process(labels model.LabelSet, extracted map[string]interface{}, _ *time.Time, _ *string) {
-	for lName, lSrc := range l.cfgs {
+	processLabelsConfigs(l.logger, extracted, l.cfgs, func(labelName model.LabelName, labelValue model.LabelValue) {
+		labels[labelName] = labelValue
+	})
+}
+
+type labelsConsumer func(labelName model.LabelName, labelValue model.LabelValue)
+
+func processLabelsConfigs(logger log.Logger, extracted map[string]interface{}, configs LabelsConfig, consumer labelsConsumer) {
+	for lName, lSrc := range configs {
 		if lValue, ok := extracted[*lSrc]; ok {
 			s, err := getString(lValue)
 			if err != nil {
 				if Debug {
-					level.Debug(l.logger).Log("msg", "failed to convert extracted label value to string", "err", err, "type", reflect.TypeOf(lValue))
+					level.Debug(logger).Log("msg", "failed to convert extracted label value to string", "err", err, "type", reflect.TypeOf(lValue))
 				}
 				continue
 			}
 			labelValue := model.LabelValue(s)
 			if !labelValue.IsValid() {
 				if Debug {
-					level.Debug(l.logger).Log("msg", "invalid label value parsed", "value", labelValue)
+					level.Debug(logger).Log("msg", "invalid label value parsed", "value", labelValue)
 				}
 				continue
 			}
-			labels[model.LabelName(lName)] = labelValue
+			consumer(model.LabelName(lName), labelValue)
 		}
 	}
 }
diff --git a/clients/pkg/logentry/stages/nonindexedlabels.go b/clients/pkg/logentry/stages/nonindexedlabels.go
new file mode 100644
index 0000000000..06302309b6
--- /dev/null
+++ b/clients/pkg/logentry/stages/nonindexedlabels.go
@@ -0,0 +1,43 @@
+package stages
+
+import (
+	"github.com/go-kit/log"
+	"github.com/mitchellh/mapstructure"
+	"github.com/prometheus/common/model"
+
+	"github.com/grafana/loki/pkg/logproto"
+)
+
+func newNonIndexedLabelsStage(params StageCreationParams) (Stage, error) {
+	cfgs := &LabelsConfig{}
+	err := mapstructure.Decode(params.config, cfgs)
+	if err != nil {
+		return nil, err
+	}
+	err = validateLabelsConfig(*cfgs)
+	if err != nil {
+		return nil, err
+	}
+	return &nonIndexedLabelsStage{
+		cfgs:   *cfgs,
+		logger: params.logger,
+	}, nil
+}
+
+type nonIndexedLabelsStage struct {
+	cfgs   LabelsConfig
+	logger log.Logger
+}
+
+func (s *nonIndexedLabelsStage) Name() string {
+	return StageTypeNonIndexedLabels
+}
+
+func (s *nonIndexedLabelsStage) Run(in chan Entry) chan Entry {
+	return RunWith(in, func(e Entry) Entry {
+		processLabelsConfigs(s.logger, e.Extracted, s.cfgs, func(labelName model.LabelName, labelValue model.LabelValue) {
+			e.NonIndexedLabels = append(e.NonIndexedLabels, logproto.LabelAdapter{Name: string(labelName), Value: string(labelValue)})
+		})
+		return e
+	})
+}
diff --git a/clients/pkg/logentry/stages/nonindexedlabels_test.go b/clients/pkg/logentry/stages/nonindexedlabels_test.go
new file mode 100644
index 0000000000..4889210fc0
--- /dev/null
+++ b/clients/pkg/logentry/stages/nonindexedlabels_test.go
@@ -0,0 +1,113 @@
+package stages
+
+import (
+	"testing"
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/common/model"
+	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/loki/pkg/push"
+	util_log "github.com/grafana/loki/pkg/util/log"
+)
+
+var pipelineStagesNonIndexedLabelsFromLogfmt = `
+pipeline_stages:
+- logfmt:
+    mapping:
+      app:
+- non_indexed_labels:
+    app:
+`
+
+var pipelineStagesNonIndexedLabelsFromJSON = `
+pipeline_stages:
+- json:
+    expressions:
+      app:
+- non_indexed_labels:
+    app:
+`
+
+var pipelineStagesNonIndexedLabelsWithRegexParser = `
+pipeline_stages:
+- regex:
+    expression: "^(?s)(?P<time>\\S+?) (?P<stream>stdout|stderr) (?P<flags>\\S+?) (?P<content>.*)$"
+- non_indexed_labels:
+    stream:
+`
+
+var pipelineStagesNonIndexedLabelsFromJSONWithTemplate = `
+pipeline_stages:
+- json:
+    expressions:
+      app:
+- template:
+    source: app
+    template: '{{ ToUpper .Value }}'
+- non_indexed_labels:
+    app:
+`
+
+var pipelineStagesNonIndexedAndRegularLabelsFromJSON = `
+pipeline_stages:
+- json:
+    expressions:
+      app:
+      component:
+- non_indexed_labels:
+    app:
+- labels:
+    component: 
+`
+
+func Test_NonIndexedLabelsStage(t *testing.T) {
+	tests := map[string]struct {
+		pipelineStagesYaml       string
+		logLine                  string
+		expectedNonIndexedLabels push.LabelsAdapter
+		expectedLabels           model.LabelSet
+	}{
+		"expected non-indexed labels to be extracted with logfmt parser and to be added to entry": {
+			pipelineStagesYaml:       pipelineStagesNonIndexedLabelsFromLogfmt,
+			logLine:                  "app=loki component=ingester",
+			expectedNonIndexedLabels: push.LabelsAdapter{push.LabelAdapter{Name: "app", Value: "loki"}},
+		},
+		"expected non-indexed labels to be extracted with json parser and to be added to entry": {
+			pipelineStagesYaml:       pipelineStagesNonIndexedLabelsFromJSON,
+			logLine:                  `{"app":"loki" ,"component":"ingester"}`,
+			expectedNonIndexedLabels: push.LabelsAdapter{push.LabelAdapter{Name: "app", Value: "loki"}},
+		},
+		"expected non-indexed labels to be extracted with regexp parser and to be added to entry": {
+			pipelineStagesYaml:       pipelineStagesNonIndexedLabelsWithRegexParser,
+			logLine:                  `2019-01-01T01:00:00.000000001Z stderr P i'm a log message!`,
+			expectedNonIndexedLabels: push.LabelsAdapter{push.LabelAdapter{Name: "stream", Value: "stderr"}},
+		},
+		"expected non-indexed labels to be extracted with json parser and to be added to entry after rendering the template": {
+			pipelineStagesYaml:       pipelineStagesNonIndexedLabelsFromJSONWithTemplate,
+			logLine:                  `{"app":"loki" ,"component":"ingester"}`,
+			expectedNonIndexedLabels: push.LabelsAdapter{push.LabelAdapter{Name: "app", Value: "LOKI"}},
+		},
+		"expected non-indexed and regular labels to be extracted with json parser and to be added to entry": {
+			pipelineStagesYaml:       pipelineStagesNonIndexedAndRegularLabelsFromJSON,
+			logLine:                  `{"app":"loki" ,"component":"ingester"}`,
+			expectedNonIndexedLabels: push.LabelsAdapter{push.LabelAdapter{Name: "app", Value: "loki"}},
+			expectedLabels:           model.LabelSet{model.LabelName("component"): model.LabelValue("ingester")},
+		},
+	}
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			pl, err := NewPipeline(util_log.Logger, loadConfig(test.pipelineStagesYaml), nil, prometheus.DefaultRegisterer)
+			require.NoError(t, err)
+
+			result := processEntries(pl, newEntry(nil, nil, test.logLine, time.Now()))[0]
+			require.Equal(t, test.expectedNonIndexedLabels, result.NonIndexedLabels)
+			if test.expectedLabels != nil {
+				require.Equal(t, test.expectedLabels, result.Labels)
+			} else {
+				require.Empty(t, result.Labels)
+			}
+		})
+	}
+}
diff --git a/clients/pkg/logentry/stages/stage.go b/clients/pkg/logentry/stages/stage.go
index 84c5ea7beb..f6568fb3ef 100644
--- a/clients/pkg/logentry/stages/stage.go
+++ b/clients/pkg/logentry/stages/stage.go
@@ -3,6 +3,7 @@ package stages
 import (
 	"os"
 	"runtime"
+	"sync"
 	"time"
 
 	"github.com/go-kit/log"
@@ -15,31 +16,32 @@ import (
 )
 
 const (
-	StageTypeJSON            = "json"
-	StageTypeLogfmt          = "logfmt"
-	StageTypeRegex           = "regex"
-	StageTypeReplace         = "replace"
-	StageTypeMetric          = "metrics"
-	StageTypeLabel           = "labels"
-	StageTypeLabelDrop       = "labeldrop"
-	StageTypeTimestamp       = "timestamp"
-	StageTypeOutput          = "output"
-	StageTypeDocker          = "docker"
-	StageTypeCRI             = "cri"
-	StageTypeMatch           = "match"
-	StageTypeTemplate        = "template"
-	StageTypePipeline        = "pipeline"
-	StageTypeTenant          = "tenant"
-	StageTypeDrop            = "drop"
-	StageTypeSampling        = "sampling"
-	StageTypeLimit           = "limit"
-	StageTypeMultiline       = "multiline"
-	StageTypePack            = "pack"
-	StageTypeLabelAllow      = "labelallow"
-	StageTypeStaticLabels    = "static_labels"
-	StageTypeDecolorize      = "decolorize"
-	StageTypeEventLogMessage = "eventlogmessage"
-	StageTypeGeoIP           = "geoip"
+	StageTypeJSON             = "json"
+	StageTypeLogfmt           = "logfmt"
+	StageTypeRegex            = "regex"
+	StageTypeReplace          = "replace"
+	StageTypeMetric           = "metrics"
+	StageTypeLabel            = "labels"
+	StageTypeLabelDrop        = "labeldrop"
+	StageTypeTimestamp        = "timestamp"
+	StageTypeOutput           = "output"
+	StageTypeDocker           = "docker"
+	StageTypeCRI              = "cri"
+	StageTypeMatch            = "match"
+	StageTypeTemplate         = "template"
+	StageTypePipeline         = "pipeline"
+	StageTypeTenant           = "tenant"
+	StageTypeDrop             = "drop"
+	StageTypeSampling         = "sampling"
+	StageTypeLimit            = "limit"
+	StageTypeMultiline        = "multiline"
+	StageTypePack             = "pack"
+	StageTypeLabelAllow       = "labelallow"
+	StageTypeStaticLabels     = "static_labels"
+	StageTypeDecolorize       = "decolorize"
+	StageTypeEventLogMessage  = "eventlogmessage"
+	StageTypeGeoIP            = "geoip"
+	StageTypeNonIndexedLabels = "non_indexed_labels"
 )
 
 // Processor takes an existing set of labels, timestamp and log entry and returns either a possibly mutated
@@ -107,134 +109,119 @@ func toStage(p Processor) Stage {
 	}
 }
 
+type StageCreationParams struct {
+	logger     log.Logger
+	config     interface{}
+	registerer prometheus.Registerer
+	jobName    *string
+}
+
+type stageCreator func(StageCreationParams) (Stage, error)
+
+var stageCreators map[string]stageCreator
+
+var stageCreatorsInitLock sync.Mutex
+
+// initCreators uses lazyLoading to resolve circular dependencies issue.
+func initCreators() {
+	if stageCreators != nil {
+		return
+	}
+	stageCreatorsInitLock.Lock()
+	defer stageCreatorsInitLock.Unlock()
+	if stageCreators != nil {
+		return
+	}
+	stageCreators = map[string]stageCreator{
+		StageTypeDocker: func(params StageCreationParams) (Stage, error) {
+			return NewDocker(params.logger, params.registerer)
+		},
+		StageTypeCRI: func(params StageCreationParams) (Stage, error) {
+			return NewCRI(params.logger, params.config, params.registerer)
+		},
+		StageTypeJSON: func(params StageCreationParams) (Stage, error) {
+			return newJSONStage(params.logger, params.config)
+		},
+		StageTypeLogfmt: func(params StageCreationParams) (Stage, error) {
+			return newLogfmtStage(params.logger, params.config)
+		},
+		StageTypeRegex: func(params StageCreationParams) (Stage, error) {
+			return newRegexStage(params.logger, params.config)
+		},
+		StageTypeMetric: func(params StageCreationParams) (Stage, error) {
+			return newMetricStage(params.logger, params.config, params.registerer)
+		},
+		StageTypeLabel: func(params StageCreationParams) (Stage, error) {
+			return newLabelStage(params.logger, params.config)
+		},
+		StageTypeLabelDrop: func(params StageCreationParams) (Stage, error) {
+			return newLabelDropStage(params.config)
+		},
+		StageTypeTimestamp: func(params StageCreationParams) (Stage, error) {
+			return newTimestampStage(params.logger, params.config)
+		},
+		StageTypeOutput: func(params StageCreationParams) (Stage, error) {
+			return newOutputStage(params.logger, params.config)
+		},
+		StageTypeMatch: func(params StageCreationParams) (Stage, error) {
+			return newMatcherStage(params.logger, params.jobName, params.config, params.registerer)
+		},
+		StageTypeTemplate: func(params StageCreationParams) (Stage, error) {
+			return newTemplateStage(params.logger, params.config)
+		},
+		StageTypeTenant: func(params StageCreationParams) (Stage, error) {
+			return newTenantStage(params.logger, params.config)
+		},
+		StageTypeReplace: func(params StageCreationParams) (Stage, error) {
+			return newReplaceStage(params.logger, params.config)
+		},
+		StageTypeDrop: func(params StageCreationParams) (Stage, error) {
+			return newDropStage(params.logger, params.config, params.registerer)
+		},
+		StageTypeSampling: func(params StageCreationParams) (Stage, error) {
+			return newSamplingStage(params.logger, params.config, params.registerer)
+		},
+		StageTypeLimit: func(params StageCreationParams) (Stage, error) {
+			return newLimitStage(params.logger, params.config, params.registerer)
+		},
+		StageTypeMultiline: func(params StageCreationParams) (Stage, error) {
+			return newMultilineStage(params.logger, params.config)
+		},
+		StageTypePack: func(params StageCreationParams) (Stage, error) {
+			return newPackStage(params.logger, params.config, params.registerer)
+		},
+		StageTypeLabelAllow: func(params StageCreationParams) (Stage, error) {
+			return newLabelAllowStage(params.config)
+		},
+		StageTypeStaticLabels: func(params StageCreationParams) (Stage, error) {
+			return newStaticLabelsStage(params.logger, params.config)
+		},
+		StageTypeDecolorize: func(params StageCreationParams) (Stage, error) {
+			return newDecolorizeStage(params.config)
+		},
+		StageTypeEventLogMessage: func(params StageCreationParams) (Stage, error) {
+			return newEventLogMessageStage(params.logger, params.config)
+		},
+		StageTypeGeoIP: func(params StageCreationParams) (Stage, error) {
+			return newGeoIPStage(params.logger, params.config)
+		},
+		StageTypeNonIndexedLabels: newNonIndexedLabelsStage,
+	}
+}
+
 // New creates a new stage for the given type and configuration.
 func New(logger log.Logger, jobName *string, stageType string,
 	cfg interface{}, registerer prometheus.Registerer) (Stage, error) {
-	var s Stage
-	var err error
-	switch stageType {
-	case StageTypeDocker:
-		s, err = NewDocker(logger, registerer)
-		if err != nil {
-			return nil, err
-		}
-	case StageTypeCRI:
-		s, err = NewCRI(logger, cfg, registerer)
-		if err != nil {
-			return nil, err
-		}
-	case StageTypeJSON:
-		s, err = newJSONStage(logger, cfg)
-		if err != nil {
-			return nil, err
-		}
-	case StageTypeLogfmt:
-		s, err = newLogfmtStage(logger, cfg)
-		if err != nil {
-			return nil, err
-		}
-	case StageTypeRegex:
-		s, err = newRegexStage(logger, cfg)
-		if err != nil {
-			return nil, err
-		}
-	case StageTypeMetric:
-		s, err = newMetricStage(logger, cfg, registerer)
-		if err != nil {
-			return nil, err
-		}
-	case StageTypeLabel:
-		s, err = newLabelStage(logger, cfg)
-		if err != nil {
-			return nil, err
-		}
-	case StageTypeLabelDrop:
-		s, err = newLabelDropStage(cfg)
-		if err != nil {
-			return nil, err
-		}
-	case StageTypeTimestamp:
-		s, err = newTimestampStage(logger, cfg)
-		if err != nil {
-			return nil, err
-		}
-	case StageTypeOutput:
-		s, err = newOutputStage(logger, cfg)
-		if err != nil {
-			return nil, err
-		}
-	case StageTypeMatch:
-		s, err = newMatcherStage(logger, jobName, cfg, registerer)
-		if err != nil {
-			return nil, err
-		}
-	case StageTypeTemplate:
-		s, err = newTemplateStage(logger, cfg)
-		if err != nil {
-			return nil, err
-		}
-	case StageTypeTenant:
-		s, err = newTenantStage(logger, cfg)
-		if err != nil {
-			return nil, err
-		}
-	case StageTypeReplace:
-		s, err = newReplaceStage(logger, cfg)
-		if err != nil {
-			return nil, err
-		}
-	case StageTypeDrop:
-		s, err = newDropStage(logger, cfg, registerer)
-		if err != nil {
-			return nil, err
-		}
-	case StageTypeSampling:
-		s, err = newSamplingStage(logger, cfg, registerer)
-		if err != nil {
-			return nil, err
-		}
-	case StageTypeLimit:
-		s, err = newLimitStage(logger, cfg, registerer)
-		if err != nil {
-			return nil, err
-		}
-	case StageTypeMultiline:
-		s, err = newMultilineStage(logger, cfg)
-		if err != nil {
-			return nil, err
-		}
-	case StageTypePack:
-		s, err = newPackStage(logger, cfg, registerer)
-		if err != nil {
-			return nil, err
-		}
-	case StageTypeLabelAllow:
-		s, err = newLabelAllowStage(cfg)
-		if err != nil {
-			return nil, err
-		}
-	case StageTypeStaticLabels:
-		s, err = newStaticLabelsStage(logger, cfg)
-		if err != nil {
-			return nil, err
-		}
-	case StageTypeDecolorize:
-		s, err = newDecolorizeStage(cfg)
-		if err != nil {
-			return nil, err
-		}
-	case StageTypeEventLogMessage:
-		s, err = newEventLogMessageStage(logger, cfg)
-		if err != nil {
-			return nil, err
-		}
-	case StageTypeGeoIP:
-		s, err = newGeoIPStage(logger, cfg)
-		if err != nil {
-			return nil, err
-		}
-	default:
+	initCreators()
+	creator, ok := stageCreators[stageType]
+	if !ok {
 		return nil, errors.Errorf("Unknown stage type: %s", stageType)
 	}
-	return s, nil
+	params := StageCreationParams{
+		logger:     logger,
+		config:     cfg,
+		registerer: registerer,
+		jobName:    jobName,
+	}
+	return creator(params)
 }