traces_processor.go

package sqlparamstyleprocessor

import (
	"context"

	"go.opentelemetry.io/collector/component"
	"go.opentelemetry.io/collector/consumer"
	"go.opentelemetry.io/collector/pdata/pcommon"
	"go.opentelemetry.io/collector/pdata/ptrace"
	"go.opentelemetry.io/collector/processor"
	"go.uber.org/zap"
)

// tracesProcessor implements processor.Traces and normalizes PEP 249 paramstyle SQL in span attributes.
type tracesProcessor struct {
	config       *Config
	settings     processor.Settings
	nextConsumer consumer.Traces
	transformer  *ParamStyleTransformer
}

// newTracesProcessor creates a new traces processor instance.
func newTracesProcessor(
	config *Config,
	settings processor.Settings,
	nextConsumer consumer.Traces,
) (processor.Traces, error) {
	return &tracesProcessor{
		config:       config,
		settings:     settings,
		nextConsumer: nextConsumer,
		transformer:  NewParamStyleTransformer(config.NormalizeInClause),
	}, nil
}

// Capabilities returns the consumer capabilities of the processor.
func (p *tracesProcessor) Capabilities() consumer.Capabilities {
	return consumer.Capabilities{MutatesData: true}
}

// Start starts the processor.
func (p *tracesProcessor) Start(ctx context.Context, host component.Host) error {
	p.settings.Logger.Info("PEP 249 SQL paramstyle processor started",
		zap.Strings("attribute_keys", p.config.AttributeKeys),
		zap.Bool("normalize_in_clause", p.config.NormalizeInClause),
	)
	return nil
}

// Shutdown stops the processor.
func (p *tracesProcessor) Shutdown(ctx context.Context) error {
	p.settings.Logger.Info("PEP 249 SQL paramstyle processor stopped")
	return nil
}

// ConsumeTraces processes trace data, normalizing PEP 249 paramstyle SQL in configured span attributes.
func (p *tracesProcessor) ConsumeTraces(ctx context.Context, td ptrace.Traces) error {
	keySet := make(map[string]struct{}, len(p.config.AttributeKeys))
	for _, k := range p.config.AttributeKeys {
		keySet[k] = struct{}{}
	}

	resourceSpans := td.ResourceSpans()
	for i := 0; i < resourceSpans.Len(); i++ {
		scopeSpans := resourceSpans.At(i).ScopeSpans()
		for j := 0; j < scopeSpans.Len(); j++ {
			spans := scopeSpans.At(j).Spans()
			for k := 0; k < spans.Len(); k++ {
				p.processSpanAttributes(spans.At(k).Attributes(), keySet)
			}
		}
	}

	return p.nextConsumer.ConsumeTraces(ctx, td)
}

func (p *tracesProcessor) processSpanAttributes(attrs pcommon.Map, keySet map[string]struct{}) {
	attrs.Range(func(key string, value pcommon.Value) bool {
		if _, ok := keySet[key]; !ok {
			return true
		}
		if value.Type() != pcommon.ValueTypeStr {
			return true
		}
		sql := value.Str()
		if normalized, ok := p.transformer.Normalize(sql); ok {
			attrs.PutStr(key, normalized)
		}
		return true
	})
}