feat(loki.process) Support structured metadata as source type of stage.labels for loki.process

Fixes #4717

Author: baurmatt

docs/sources/reference/components/loki/loki.process.md

@@ -582,15 +582,16 @@ stage.label_keep {
 
 ### `stage.labels`
 
-The `stage.labels` inner block configures a labels processing stage that can read data from the extracted values map and set new labels on incoming log entries.
+The `stage.labels` inner block configures a labels processing stage that can read data from the extracted values map or structured metadata and set new labels on incoming log entries.
 
 For labels that are static, refer to [`stage.static_labels`][stage.static_labels]
 
 The following arguments are supported:
 
-| Name     | Type          | Description                             | Default | Required |
-| -------- | ------------- | --------------------------------------- | ------- | -------- |
-| `values` | `map(string)` | Configures a `labels` processing stage. | `{}`    | no       |
+| Name          | Type          | Description                                                                                              | Default       | Required |
+| ------------- | ------------- | -------------------------------------------------------------------------------------------------------- | ------------- | -------- |
+| `values`      | `map(string)` | Configures a `labels` processing stage.                                                                  | `{}`          | no       |
+| `source_type` | `string`      | Where to retrieve the data from.  Allowed values are `"extracted"` (default) or `"structured_metadata"`. | `"extracted"` | no       |
 
 In a labels stage, the map's keys define the label to set and the values are how to look them up.
 If the value is empty, it's inferred to be the same as the key.
@@ -604,6 +605,16 @@ stage.labels {
 }
 ```
 
+```alloy
+stage.labels {
+    source_type = "structured_metadata"
+    values      = {
+      env  = "",         // Sets up an 'env' label, based on the 'env' structured metadata value.
+      user = "username", // Sets up a 'user' label, based on the 'username' structured metadata value.
+    }
+}
+```
+
 ### `stage.limit`
 
 The `stage.limit` inner block configures a rate-limiting stage that throttles logs based on several options.

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

@@ -4,92 +4,145 @@ import (
 	"errors"
 	"fmt"
 	"reflect"
-	"time"
 
 	"github.com/go-kit/log"
 	"github.com/prometheus/common/model"
 
 	"github.com/grafana/alloy/internal/runtime/logging/level"
+	"github.com/grafana/loki/pkg/push"
 )
 
 const (
 	ErrEmptyLabelStageConfig = "label stage config cannot be empty"
 	ErrInvalidLabelName      = "invalid label name: %s"
+	ErrInvalidSourceType     = "invalid labels source_type: %s. Can only be 'extracted' or 'structured_metadata'"
+
+	LabelsSourceStructuredMetadata string = "structured_metadata"
+	LabelsSourceExtractedMap       string = "extracted"
 )
 
 // LabelsConfig is a set of labels to be extracted
 type LabelsConfig struct {
-	Values map[string]*string `alloy:"values,attr"`
+	Values     map[string]*string `alloy:"values,attr"`
+	SourceType string             `alloy:"source_type,attr,optional"`
 }
 
 // validateLabelsConfig validates the Label stage configuration
-func validateLabelsConfig(c map[string]*string) (map[string]string, error) {
-	// We must not mutate the c.Values, create a copy with changes we need.
-	ret := map[string]string{}
-	if c == nil {
-		return nil, errors.New(ErrEmptyLabelStageConfig)
+func validateLabelsConfig(cfg *LabelsConfig) error {
+	if cfg.Values == nil {
+		return errors.New(ErrEmptyLabelStageConfig)
+	}
+
+	if cfg.SourceType == "" {
+		cfg.SourceType = LabelsSourceExtractedMap
 	}
-	for labelName, labelSrc := range c {
+
+	switch cfg.SourceType {
+	case LabelsSourceExtractedMap, LabelsSourceStructuredMetadata:
+	default:
+		return fmt.Errorf(ErrInvalidSourceType, cfg.SourceType)
+	}
+
+	// We must not mutate the c.Values, create a copy with changes we need.
+	returnValues := map[string]*string{}
+	for labelName, labelSrc := range cfg.Values {
 		// TODO: add support for different validation schemes.
 		//nolint:staticcheck
 		if !model.LabelName(labelName).IsValid() {
-			return nil, fmt.Errorf(ErrInvalidLabelName, labelName)
+			return fmt.Errorf(ErrInvalidLabelName, labelName)
 		}
 		// If no label source was specified, use the key name
 		if labelSrc == nil || *labelSrc == "" {
-			ret[labelName] = labelName
+			returnValues[labelName] = &labelName
 		} else {
-			ret[labelName] = *labelSrc
+			returnValues[labelName] = labelSrc
 		}
 	}
-	return ret, nil
+	cfg.Values = returnValues
+	return nil
 }
 
 // newLabelStage creates a new label stage to set labels from extracted data
 func newLabelStage(logger log.Logger, configs LabelsConfig) (Stage, error) {
-	labelsConfig, err := validateLabelsConfig(configs.Values)
+	err := validateLabelsConfig(&configs)
 	if err != nil {
 		return nil, err
 	}
-	return toStage(&labelStage{
-		labelsConfig: labelsConfig,
-		logger:       logger,
-	}), nil
+	return &labelStage{
+		cfg:    &configs,
+		logger: logger,
+	}, nil
 }
 
 // labelStage sets labels from extracted data
 type labelStage struct {
-	labelsConfig map[string]string
-	logger       log.Logger
+	cfg    *LabelsConfig
+	logger log.Logger
 }
 
-// Process implements Stage
-func (l *labelStage) Process(labels model.LabelSet, extracted map[string]any, _ *time.Time, _ *string) {
-	processLabelsConfigs(l.logger, extracted, l.labelsConfig, func(labelName model.LabelName, labelValue model.LabelValue) {
-		labels[labelName] = labelValue
-	})
+// Run implements Stage
+func (l *labelStage) Run(in chan Entry) chan Entry {
+	out := make(chan Entry)
+	go func() {
+		defer close(out)
+		for e := range in {
+			switch l.cfg.SourceType {
+			case LabelsSourceExtractedMap:
+				l.addLabelFromExtractedMap(e.Labels, e.Extracted)
+			case LabelsSourceStructuredMetadata:
+				l.addLabelsFromStructuredMetadata(e.Labels, e.StructuredMetadata)
+			}
+			out <- e
+		}
+	}()
+	return out
 }
 
-type labelsConsumer func(labelName model.LabelName, labelValue model.LabelValue)
-
-func processLabelsConfigs(logger log.Logger, extracted map[string]any, labelsConfig map[string]string, consumer labelsConsumer) {
-	for lName, lSrc := range labelsConfig {
-		if lValue, ok := extracted[lSrc]; ok {
+func (l *labelStage) addLabelFromExtractedMap(labels model.LabelSet, extracted map[string]any) {
+	for lName, lSrc := range l.cfg.Values {
+		if lValue, ok := extracted[*lSrc]; ok {
 			s, err := getString(lValue)
 			if err != nil {
 				if Debug {
-					level.Debug(logger).Log("msg", "failed to convert extracted label value to string", "err", err, "type", reflect.TypeOf(lValue))
+					level.Debug(l.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(logger).Log("msg", "invalid label value parsed", "value", labelValue)
+					level.Debug(l.logger).Log("msg", "invalid label value parsed", "value", labelValue)
 				}
 				continue
 			}
-			consumer(model.LabelName(lName), labelValue)
+
+			labels[model.LabelName(lName)] = labelValue
+		}
+	}
+}
+
+func (l *labelStage) addLabelsFromStructuredMetadata(labels model.LabelSet, metadata push.LabelsAdapter) {
+	for lName, lSrc := range l.cfg.Values {
+		for _, kv := range metadata {
+			if kv.Name != *lSrc {
+				continue
+			}
+
+			labelValue := model.LabelValue(kv.Value)
+			if !labelValue.IsValid() {
+				if Debug {
+					level.Debug(l.logger).Log("msg", "invalid structured metadata label value", "label", lName, "value", labelValue)
+				}
+				break
+			}
+
+			labels[model.LabelName(lName)] = labelValue
+			break
 		}
 	}
 }
+
+// Cleanup implements Stage.
+func (*labelStage) Cleanup() {
+	// no-op
+}