Use the same SourceType for labels and truncate stage

Author: baurmatt

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

@@ -24,7 +24,7 @@ const (
 // LabelsConfig is a set of labels to be extracted
 type LabelsConfig struct {
 	Values     map[string]*string `alloy:"values,attr"`
-	SourceType string             `alloy:"source_type,attr,optional"`
+	SourceType SourceType         `alloy:"source_type,attr,optional"`
 }
 
 // validateLabelsConfig validates the Label stage configuration
@@ -34,11 +34,11 @@ func validateLabelsConfig(cfg *LabelsConfig) error {
 	}
 
 	if cfg.SourceType == "" {
-		cfg.SourceType = LabelsSourceExtractedMap
+		cfg.SourceType = SourceTypeExtractedMap
 	}
 
 	switch cfg.SourceType {
-	case LabelsSourceExtractedMap, LabelsSourceStructuredMetadata:
+	case SourceTypeExtractedMap, SourceTypeStructuredMetadata:
 	default:
 		return fmt.Errorf(ErrInvalidSourceType, cfg.SourceType)
 	}
@@ -87,9 +87,9 @@ func (l *labelStage) Run(in chan Entry) chan Entry {
 		defer close(out)
 		for e := range in {
 			switch l.cfg.SourceType {
-			case LabelsSourceExtractedMap:
+			case SourceTypeExtractedMap:
 				l.addLabelFromExtractedMap(e.Labels, e.Extracted)
-			case LabelsSourceStructuredMetadata:
+			case SourceTypeStructuredMetadata:
 				l.addLabelsFromStructuredMetadata(e.Labels, e.StructuredMetadata)
 			}
 			out <- e

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

@@ -142,7 +142,7 @@ func TestLabels(t *testing.T) {
 		},
 		"label value is set from name for extracted": {
 			config: LabelsConfig{
-				SourceType: LabelsSourceExtractedMap,
+				SourceType: SourceTypeExtractedMap,
 				Values: map[string]*string{
 					"l1": &lv1,
 					"l2": nil,
@@ -157,7 +157,7 @@ func TestLabels(t *testing.T) {
 		},
 		"label value is set from name for structured_metadata": {
 			config: LabelsConfig{
-				SourceType: LabelsSourceStructuredMetadata,
+				SourceType: SourceTypeStructuredMetadata,
 				Values: map[string]*string{
 					"l1": &lv1,
 					"l2": nil,
@@ -215,7 +215,7 @@ func TestLabelsStage_Process(t *testing.T) {
 		},
 		"extract_success_structured_metadata": {
 			LabelsConfig{
-				SourceType: LabelsSourceStructuredMetadata,
+				SourceType: SourceTypeStructuredMetadata,
 				Values: map[string]*string{
 					"testLabel": ptr("testStrucuturedMetadata"),
 				}},
@@ -243,7 +243,7 @@ func TestLabelsStage_Process(t *testing.T) {
 		},
 		"different_source_name_structured_metadata": {
 			LabelsConfig{
-				SourceType: LabelsSourceStructuredMetadata,
+				SourceType: SourceTypeStructuredMetadata,
 				Values: map[string]*string{
 					"testLabel": &sourceName,
 				}},
@@ -267,7 +267,7 @@ func TestLabelsStage_Process(t *testing.T) {
 		},
 		"empty_structured_metadata": {
 			LabelsConfig{
-				SourceType: LabelsSourceStructuredMetadata,
+				SourceType: SourceTypeStructuredMetadata,
 				Values: map[string]*string{
 					"testLabel": &sourceName,
 				}},

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

@@ -1,9 +1,7 @@
 package stages
 
 import (
-	"encoding"
 	"errors"
-	"fmt"
 	"maps"
 	"slices"
 	"strings"
@@ -36,46 +34,14 @@ type TruncateConfig struct {
 }
 
 type RuleConfig struct {
-	Limit      units.Base2Bytes   `alloy:"limit,attr"`
-	Suffix     string             `alloy:"suffix,attr,optional"`
-	Sources    []string           `alloy:"sources,attr,optional"`
-	SourceType TruncateSourceType `alloy:"source_type,attr,optional"`
+	Limit      units.Base2Bytes `alloy:"limit,attr"`
+	Suffix     string           `alloy:"suffix,attr,optional"`
+	Sources    []string         `alloy:"sources,attr,optional"`
+	SourceType SourceType       `alloy:"source_type,attr,optional"`
 
 	effectiveLimit units.Base2Bytes
 }
 
-var (
-	_ encoding.TextMarshaler   = TruncateSourceType("")
-	_ encoding.TextUnmarshaler = (*TruncateSourceType)(nil)
-)
-
-type TruncateSourceType string
-
-// UnmarshalText implements encoding.TextUnmarshaler.
-func (t *TruncateSourceType) UnmarshalText(text []byte) error {
-	str := string(text)
-	switch str {
-	case string(TruncateSourceLine), string(TruncateSourceLabel), string(TruncateSourceStructuredMetadata), string(TruncateSourceExtractedMap):
-		*t = TruncateSourceType(str)
-	default:
-		return fmt.Errorf("unknown source_type: %s", str)
-	}
-
-	return nil
-}
-
-// MarshalText implements encoding.TextMarshaler.
-func (t TruncateSourceType) MarshalText() (text []byte, err error) {
-	return []byte(t), nil
-}
-
-const (
-	TruncateSourceLine               TruncateSourceType = "line"
-	TruncateSourceLabel              TruncateSourceType = "label"
-	TruncateSourceStructuredMetadata TruncateSourceType = "structured_metadata"
-	TruncateSourceExtractedMap       TruncateSourceType = "extracted"
-)
-
 // validateTruncateConfig validates the TruncateConfig for the truncateStage
 func validateTruncateConfig(cfg *TruncateConfig) error {
 	if len(cfg.Rules) == 0 {
@@ -90,10 +56,10 @@ func validateTruncateConfig(cfg *TruncateConfig) error {
 		}
 
 		if r.SourceType == "" {
-			r.SourceType = TruncateSourceLine
+			r.SourceType = SourceTypeLine
 		}
 
-		if r.SourceType == TruncateSourceLine && len(r.Sources) > 0 {
+		if r.SourceType == SourceTypeLine && len(r.Sources) > 0 {
 			return errors.New(errSourcesForLine)
 		}
 
@@ -138,7 +104,7 @@ func (m *truncateStage) Run(in chan Entry) chan Entry {
 			truncated := map[string]struct{}{}
 			for _, r := range m.cfg.Rules {
 				switch r.SourceType {
-				case TruncateSourceLine:
+				case SourceTypeLine:
 					if len(e.Line) > int(r.effectiveLimit) {
 						e.Line = e.Line[:r.effectiveLimit] + r.Suffix
 						markTruncated(m.truncatedCount, truncated, truncateLineField)
@@ -147,7 +113,7 @@ func (m *truncateStage) Run(in chan Entry) chan Entry {
 							level.Debug(m.logger).Log("msg", "line has been truncated", "limit", r.effectiveLimit, "truncated_line", e.Line)
 						}
 					}
-				case TruncateSourceLabel:
+				case SourceTypeLabel:
 					if len(r.Sources) > 0 {
 						for _, source := range r.Sources {
 							name := model.LabelName(source)
@@ -160,7 +126,7 @@ func (m *truncateStage) Run(in chan Entry) chan Entry {
 							m.tryTruncateLabel(r, e.Labels, k, v, truncated)
 						}
 					}
-				case TruncateSourceStructuredMetadata:
+				case SourceTypeStructuredMetadata:
 					if len(r.Sources) > 0 {
 						for i, v := range e.StructuredMetadata {
 							if slices.Contains(r.Sources, v.Name) {
@@ -174,7 +140,7 @@ func (m *truncateStage) Run(in chan Entry) chan Entry {
 							e.StructuredMetadata[i] = m.tryTruncateStructuredMetadata(r, v, truncated)
 						}
 					}
-				case TruncateSourceExtractedMap:
+				case SourceTypeExtractedMap:
 					if len(r.Sources) > 0 {
 						for _, source := range r.Sources {
 							if v, ok := e.Extracted[source]; ok {

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

@@ -80,7 +80,7 @@ func Test_TruncateStage_Process(t *testing.T) {
 			config: []*RuleConfig{
 				{
 					Limit:      15,
-					SourceType: TruncateSourceLabel,
+					SourceType: SourceTypeLabel,
 					Suffix:     "[truncated]",
 				},
 			},
@@ -96,7 +96,7 @@ func Test_TruncateStage_Process(t *testing.T) {
 			config: []*RuleConfig{
 				{
 					Limit:      15,
-					SourceType: TruncateSourceLabel,
+					SourceType: SourceTypeLabel,
 					Suffix:     "[truncated]",
 					Sources:    []string{"app"},
 				},
@@ -113,7 +113,7 @@ func Test_TruncateStage_Process(t *testing.T) {
 			config: []*RuleConfig{
 				{
 					Limit:      15,
-					SourceType: TruncateSourceStructuredMetadata,
+					SourceType: SourceTypeStructuredMetadata,
 					Suffix:     "<trunc>",
 				},
 			},
@@ -132,7 +132,7 @@ func Test_TruncateStage_Process(t *testing.T) {
 			config: []*RuleConfig{
 				{
 					Limit:      15,
-					SourceType: TruncateSourceStructuredMetadata,
+					SourceType: SourceTypeStructuredMetadata,
 					Suffix:     "<trunc>",
 					Sources:    []string{"meta1"},
 				},
@@ -155,18 +155,18 @@ func Test_TruncateStage_Process(t *testing.T) {
 				},
 				{
 					Limit:      15,
-					SourceType: TruncateSourceLabel,
+					SourceType: SourceTypeLabel,
 					Suffix:     "[truncated]",
 					Sources:    []string{"app"},
 				},
 				{
 					Limit:      15,
-					SourceType: TruncateSourceStructuredMetadata,
+					SourceType: SourceTypeStructuredMetadata,
 					Suffix:     "<trunc>",
 				},
 				{
 					Limit:      8,
-					SourceType: TruncateSourceExtractedMap,
+					SourceType: SourceTypeExtractedMap,
 					Sources:    []string{"field2"},
 				},
 			},
@@ -298,7 +298,7 @@ func Test_ValidateTruncateConfig(t *testing.T) {
 			config: &TruncateConfig{
 				Rules: []*RuleConfig{{
 					Limit:      10,
-					SourceType: TruncateSourceLabel,
+					SourceType: SourceTypeLabel,
 					Suffix:     "...",
 				}},
 			},
@@ -309,7 +309,7 @@ func Test_ValidateTruncateConfig(t *testing.T) {
 			config: &TruncateConfig{
 				Rules: []*RuleConfig{{
 					Limit:      10,
-					SourceType: TruncateSourceStructuredMetadata,
+					SourceType: SourceTypeStructuredMetadata,
 					Suffix:     "...",
 				}},
 			},
@@ -321,7 +321,7 @@ func Test_ValidateTruncateConfig(t *testing.T) {
 				Rules: []*RuleConfig{
 					{
 						Limit:      10,
-						SourceType: TruncateSourceLabel,
+						SourceType: SourceTypeLabel,
 						Sources:    []string{"app"},
 						Suffix:     "...",
 					},

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

@@ -0,0 +1,38 @@
+package stages
+
+import (
+	"encoding"
+	"fmt"
+)
+
+type SourceType string
+
+const (
+	SourceTypeLine               SourceType = "line"
+	SourceTypeLabel              SourceType = "label"
+	SourceTypeStructuredMetadata SourceType = "structured_metadata"
+	SourceTypeExtractedMap       SourceType = "extracted"
+)
+
+var (
+	_ encoding.TextMarshaler   = SourceType("")
+	_ encoding.TextUnmarshaler = (*SourceType)(nil)
+)
+
+// UnmarshalText implements encoding.TextUnmarshaler.
+func (t *SourceType) UnmarshalText(text []byte) error {
+	str := string(text)
+	switch str {
+	case string(SourceTypeLine), string(SourceTypeLabel), string(SourceTypeStructuredMetadata), string(SourceTypeExtractedMap):
+		*t = SourceType(str)
+	default:
+		return fmt.Errorf("unknown source_type: %s", str)
+	}
+
+	return nil
+}
+
+// MarshalText implements encoding.TextMarshaler.
+func (t SourceType) MarshalText() (text []byte, err error) {
+	return []byte(t), nil
+}