From f2feff86d9347c2e28d533ee62e807ba40f497c2 Mon Sep 17 00:00:00 2001 From: Luke Cawood Date: Tue, 20 Jan 2026 19:13:47 +1100 Subject: [PATCH] Add WithOperationNameSetter to allow control of the operation name sent to otel metrics --- config.go | 7 +++++++ option.go | 8 ++++++++ option_test.go | 15 +++++++++++++++ utils.go | 22 ++++++++++++++++------ 4 files changed, 46 insertions(+), 6 deletions(-) diff --git a/config.go b/config.go index be861810..6204d5cf 100644 --- a/config.go +++ b/config.go @@ -49,6 +49,9 @@ type InstrumentAttributesGetter func(ctx context.Context, method Method, query s // InstrumentErrorAttributesGetter provides additional error-related attributes while recording metrics to instruments. type InstrumentErrorAttributesGetter func(err error) []attribute.KeyValue +// OperationNameSetter allows controlling the operation name provided to metric instruments. +type OperationNameSetter func(ctx context.Context, method Method, query string) string + // SpanFilter is a function that determines whether a span should be created for a given SQL operation. // It returns true if the span should be created, or false to skip span creation. type SpanFilter func(ctx context.Context, method Method, query string, args []driver.NamedValue) bool @@ -92,6 +95,10 @@ type config struct { InstrumentErrorAttributesGetter InstrumentErrorAttributesGetter + // OperationNameSetter will be called to produce the operation name for metric instruments. + // If nil, the default behavior uses the method name. + OperationNameSetter OperationNameSetter + // DisableSkipErrMeasurement, if set to true, will suppress driver.ErrSkip as an error status in metrics. // The metric measurement will be recorded as status=ok. // Default is false diff --git a/option.go b/option.go index e3492564..8beb93bc 100644 --- a/option.go +++ b/option.go @@ -140,3 +140,11 @@ func WithInstrumentErrorAttributesGetter(instrumentErrorAttributesGetter Instrum cfg.InstrumentErrorAttributesGetter = instrumentErrorAttributesGetter }) } + +// WithOperationNameSetter takes OperationNameSetter that will be called to produce the operation name for metric instruments. +// If not set, the default behavior uses the method name. +func WithOperationNameSetter(operationNameSetter OperationNameSetter) Option { + return OptionFunc(func(cfg *config) { + cfg.OperationNameSetter = operationNameSetter + }) +} diff --git a/option_test.go b/option_test.go index e84ed5f2..b755ab2b 100644 --- a/option_test.go +++ b/option_test.go @@ -39,6 +39,10 @@ func TestOptions(t *testing.T) { return []attribute.KeyValue{attribute.String("errorKey", "errorVal")} } + dummyOperationNameSetter := func(_ context.Context, _ Method, _ string) string { + return "custom_operation" + } + testCases := []struct { name string options []Option @@ -109,6 +113,11 @@ func TestOptions(t *testing.T) { options: []Option{WithInstrumentErrorAttributesGetter(dummyErrorAttributesGetter)}, expectedConfig: config{InstrumentErrorAttributesGetter: dummyErrorAttributesGetter}, }, + { + name: "WithOperationNameSetter", + options: []Option{WithOperationNameSetter(dummyOperationNameSetter)}, + expectedConfig: config{OperationNameSetter: dummyOperationNameSetter}, + }, { name: "WithAttributes multiple calls should accumulate", options: []Option{ @@ -207,6 +216,12 @@ func TestOptions(t *testing.T) { tc.expectedConfig.InstrumentErrorAttributesGetter(assert.AnError), cfg.InstrumentErrorAttributesGetter(assert.AnError), ) + case tc.expectedConfig.OperationNameSetter != nil: + assert.Equal( + t, + tc.expectedConfig.OperationNameSetter(context.Background(), "", ""), + cfg.OperationNameSetter(context.Background(), "", ""), + ) default: assert.Equal(t, tc.expectedConfig, cfg) } diff --git a/utils.go b/utils.go index 08748d86..78af22a1 100644 --- a/utils.go +++ b/utils.go @@ -66,9 +66,14 @@ func recordLegacyLatency( duration time.Duration, attributes []attribute.KeyValue, method Method, + query string, err error, ) { - attributes = append(attributes, queryMethodKey.String(string(method))) + operationName := string(method) + if cfg.OperationNameSetter != nil { + operationName = cfg.OperationNameSetter(ctx, method, query) + } + attributes = append(attributes, queryMethodKey.String(operationName)) if err != nil { if cfg.DisableSkipErrMeasurement && errors.Is(err, driver.ErrSkip) { @@ -94,9 +99,14 @@ func recordDuration( duration time.Duration, attributes []attribute.KeyValue, method Method, + query string, err error, ) { - attributes = append(attributes, semconv.DBOperationName(string(method))) + operationName := string(method) + if cfg.OperationNameSetter != nil { + operationName = cfg.OperationNameSetter(ctx, method, query) + } + attributes = append(attributes, semconv.DBOperationName(operationName)) if err != nil && (!cfg.DisableSkipErrMeasurement || !errors.Is(err, driver.ErrSkip)) { attributes = append(attributes, internalsemconv.ErrorTypeAttributes(err)...) } @@ -147,13 +157,13 @@ func recordMetric( switch cfg.SemConvStabilityOptIn { case internalsemconv.OTelSemConvStabilityOptInStable: - recordDuration(ctx, instruments, cfg, duration, attributes, method, err) + recordDuration(ctx, instruments, cfg, duration, attributes, method, query, err) case internalsemconv.OTelSemConvStabilityOptInDup: // Intentionally emit both legacy and new metrics for backward compatibility. - recordLegacyLatency(ctx, instruments, cfg, duration, slices.Clone(attributes), method, err) - recordDuration(ctx, instruments, cfg, duration, attributes, method, err) + recordLegacyLatency(ctx, instruments, cfg, duration, slices.Clone(attributes), method, query, err) + recordDuration(ctx, instruments, cfg, duration, attributes, method, query, err) case internalsemconv.OTelSemConvStabilityOptInNone: - recordLegacyLatency(ctx, instruments, cfg, duration, attributes, method, err) + recordLegacyLatency(ctx, instruments, cfg, duration, attributes, method, query, err) } } }