Support custom metrics and trace labels provided from the context

README.md

@@ -49,6 +49,7 @@ Compatibility
 
 |          `otelsql`          | `go.opentelemetry.io/otel/trace` | `go.opentelemetry.io/otel/metric` |
 |:---------------------------:|:--------------------------------:|:---------------------------------:|
+| `v0.14.*` <sup>&nbsp;</sup> |      `v1.25.0` ~> `latest`       |       `v1.25.0` ~> `latest`       |
 | `v0.13.*` <sup>&nbsp;</sup> |      `v1.24.0` ~> `latest`       |       `v1.24.0` ~> `latest`       |
 | `v0.12.*` <sup>&nbsp;</sup> |      `v1.17.0` ~> `latest`       |       `v0.40.0` ~> `latest`       |
 | `v0.11.*` <sup>&nbsp;</sup> |      `v1.16.0` ~> `latest`       |       `v0.39.0` ~> `latest`       |

begin_internal_test.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/metric/noop"
 	semconv "go.opentelemetry.io/otel/semconv/v1.20.0"
 
@@ -138,6 +139,7 @@ func TestBeginStats(t *testing.T) {
 	testCases := []struct {
 		scenario string
 		begin    beginFunc
+		ctxLabel []attribute.KeyValue
 		expected string
 	}{
 		{
@@ -172,6 +174,24 @@ func TestBeginStats(t *testing.T) {
 				}
 			]`,
 		},
+		{
+			scenario: "extra labels",
+			begin:    nopBegin,
+			ctxLabel: []attribute.KeyValue{
+				attribute.String("extra", "label"),
+			},
+			expected: `[
+				{
+					"Name": "db.sql.client.calls{service.name=otelsql,instrumentation.name=begin_test,db.instance=test,db.operation=go.sql.begin,db.sql.status=OK,db.system=other_sql,extra=label}",
+					"Sum": 1
+				},
+				{
+					"Name": "db.sql.client.latency{service.name=otelsql,instrumentation.name=begin_test,db.instance=test,db.operation=go.sql.begin,db.sql.status=OK,db.system=other_sql,extra=label}",
+					"Sum": "<ignore-diff>",
+					"Count": 1
+				}
+			]`,
+		},
 	}
 
 	for _, tc := range testCases {
@@ -198,7 +218,9 @@ func TestBeginStats(t *testing.T) {
 						beginStats(r),
 					}, tc.begin)
 
-					_, _ = begin(context.Background(), driver.TxOptions{}) // nolint: errcheck
+					ctx := ContextWithMetricsLabels(context.Background(), tc.ctxLabel...)
+
+					_, _ = begin(ctx, driver.TxOptions{}) // nolint: errcheck
 				})
 		})
 	}

exec_internal_test.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/metric/noop"
 	semconv "go.opentelemetry.io/otel/semconv/v1.20.0"
 
@@ -112,6 +113,7 @@ func TestExecStats(t *testing.T) {
 	testCases := []struct {
 		scenario string
 		execer   execContextFunc
+		ctxLabel []attribute.KeyValue
 		expected string
 	}{
 		{
@@ -146,6 +148,24 @@ func TestExecStats(t *testing.T) {
 				}
 			]`,
 		},
+		{
+			scenario: "extra labels",
+			execer:   nopExecContext,
+			ctxLabel: []attribute.KeyValue{
+				attribute.String("extra", "label"),
+			},
+			expected: `[
+				{
+					"Name": "db.sql.client.calls{service.name=otelsql,instrumentation.name=exec_test,db.instance=test,db.operation=go.sql.exec,db.sql.status=OK,db.system=other_sql,extra=label}",
+					"Sum": 1
+				},
+				{
+					"Name": "db.sql.client.latency{service.name=otelsql,instrumentation.name=exec_test,db.instance=test,db.operation=go.sql.exec,db.sql.status=OK,db.system=other_sql,extra=label}",
+					"Sum": "<ignore-diff>",
+					"Count": 1
+				}
+			]`,
+		},
 	}
 
 	for _, tc := range testCases {
@@ -172,7 +192,9 @@ func TestExecStats(t *testing.T) {
 						execStats(r, metricMethodExec),
 					}, tc.execer)
 
-					_, _ = exec(context.Background(), "", nil) // nolint: errcheck
+					ctx := ContextWithMetricsLabels(context.Background(), tc.ctxLabel...)
+
+					_, _ = exec(ctx, "", nil) // nolint: errcheck
 				})
 		})
 	}

go.mod

@@ -32,7 +32,7 @@ require (
 	github.com/yudai/gojsondiff v1.0.0 // indirect
 	github.com/yudai/golcs v0.0.0-20170316035057-ecda9a501e82 // indirect
 	golang.org/x/net v0.22.0 // indirect
-	golang.org/x/sys v0.18.0 // indirect
+	golang.org/x/sys v0.19.0 // indirect
 	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -72,8 +72,8 @@ golang.org/x/net v0.22.0 h1:9sGLhx7iRIHEiX0oAJ3MRZMUCElJgy7Br1nO+AMN3Tc=
 golang.org/x/net v0.22.0/go.mod h1:JKghWKKOSdJwpW2GEx0Ja7fmaKnMsbu+MWVZTokSYmg=
 golang.org/x/sys v0.0.0-20220908164124-27713097b956/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.18.0 h1:DBdB3niSjOA/O0blCZBqDefyWNYveAYMNF1Wum0DYQ4=
-golang.org/x/sys v0.18.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/sys v0.19.0 h1:q5f1RH2jigJ1MoAWp2KTp3gm5zAGFUTarQZ5U386+4o=
+golang.org/x/sys v0.19.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
 golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

label.go

@@ -0,0 +1,89 @@
+package otelsql
+
+import (
+	"context"
+	"sync"
+
+	"go.opentelemetry.io/otel/attribute"
+)
+
+// Labeler is a helper to add attributes to a context.
+type Labeler struct {
+	mu         sync.Mutex
+	attributes []attribute.KeyValue
+}
+
+// Add attributes to a Labeler.
+func (l *Labeler) Add(ls ...attribute.KeyValue) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	l.attributes = append(l.attributes, ls...)
+}
+
+// Get returns a copy of the attributes added to the Labeler.
+func (l *Labeler) Get() []attribute.KeyValue {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	ret := make([]attribute.KeyValue, len(l.attributes))
+	copy(ret, l.attributes)
+
+	return ret
+}
+
+const (
+	labelerCtxMetrics = labelerContextKey("metrics")
+	labelerCtxTrace   = labelerContextKey("trace")
+)
+
+type labelerContextKey string
+
+// MetricsLabelsFromContext retrieves the labels from the provided context.
+func MetricsLabelsFromContext(ctx context.Context) []attribute.KeyValue {
+	l, _ := labelerFromContext(ctx, labelerCtxMetrics)
+
+	return l.Get()
+}
+
+// ContextWithMetricsLabels returns a new context with the labels added to the Labeler.
+func ContextWithMetricsLabels(ctx context.Context, labels ...attribute.KeyValue) context.Context {
+	return contextWithLabels(ctx, labelerCtxMetrics, labels...)
+}
+
+// TraceLabelsFromContext retrieves the labels from the provided context.
+func TraceLabelsFromContext(ctx context.Context) []attribute.KeyValue {
+	l, _ := labelerFromContext(ctx, labelerCtxTrace)
+
+	return l.Get()
+}
+
+// ContextWithTraceLabels returns a new context with the labels added to the Labeler.
+func ContextWithTraceLabels(ctx context.Context, labels ...attribute.KeyValue) context.Context {
+	return contextWithLabels(ctx, labelerCtxTrace, labels...)
+}
+
+// ContextWithTraceAndMetricsLabels returns a new context with the labels added to the Labeler.
+func ContextWithTraceAndMetricsLabels(ctx context.Context, labels ...attribute.KeyValue) context.Context {
+	ctx = ContextWithMetricsLabels(ctx, labels...)
+	ctx = ContextWithTraceLabels(ctx, labels...)
+
+	return ctx
+}
+
+func labelerFromContext(ctx context.Context, key labelerContextKey) (*Labeler, bool) { //nolint: unparam
+	l, ok := ctx.Value(key).(*Labeler)
+	if !ok {
+		l = &Labeler{}
+	}
+
+	return l, ok
+}
+
+func contextWithLabels(ctx context.Context, key labelerContextKey, labels ...attribute.KeyValue) context.Context {
+	l, _ := labelerFromContext(ctx, key)
+
+	l.Add(labels...)
+
+	return context.WithValue(ctx, key, l)
+}

label_test.go

@@ -0,0 +1,77 @@
+package otelsql_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"go.opentelemetry.io/otel/attribute"
+
+	"go.nhat.io/otelsql"
+)
+
+var (
+	label1 = attribute.String("key1", "value1")
+	label2 = attribute.Int("key2", 2)
+	label3 = attribute.Bool("key3", true)
+)
+
+func TestMetricsLabeler(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+
+	labels := otelsql.MetricsLabelsFromContext(ctx)
+	assert.Empty(t, labels)
+
+	ctx = otelsql.ContextWithMetricsLabels(ctx, label1, label2)
+
+	labels = otelsql.MetricsLabelsFromContext(ctx)
+	assert.Equal(t, []attribute.KeyValue{label1, label2}, labels)
+
+	ctx = otelsql.ContextWithMetricsLabels(ctx, label3)
+	assert.Equal(t, []attribute.KeyValue{label1, label2, label3}, otelsql.MetricsLabelsFromContext(ctx))
+
+	assert.Empty(t, otelsql.TraceLabelsFromContext(ctx))
+}
+
+func TestTraceLabeler(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+
+	labels := otelsql.TraceLabelsFromContext(ctx)
+	assert.Empty(t, labels)
+
+	ctx = otelsql.ContextWithTraceLabels(ctx, label1, label2)
+
+	labels = otelsql.TraceLabelsFromContext(ctx)
+	assert.Equal(t, []attribute.KeyValue{label1, label2}, labels)
+
+	ctx = otelsql.ContextWithTraceLabels(ctx, label3)
+	assert.Equal(t, []attribute.KeyValue{label1, label2, label3}, otelsql.TraceLabelsFromContext(ctx))
+
+	assert.Empty(t, otelsql.MetricsLabelsFromContext(ctx))
+}
+
+func TestMetricsAndTraceLabeler(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+
+	labels := otelsql.MetricsLabelsFromContext(ctx)
+	assert.Empty(t, labels)
+
+	labels = otelsql.TraceLabelsFromContext(ctx)
+	assert.Empty(t, labels)
+
+	ctx = otelsql.ContextWithTraceAndMetricsLabels(ctx, label1, label2)
+
+	assert.Equal(t, []attribute.KeyValue{label1, label2}, otelsql.MetricsLabelsFromContext(ctx))
+	assert.Equal(t, []attribute.KeyValue{label1, label2}, otelsql.TraceLabelsFromContext(ctx))
+
+	ctx = otelsql.ContextWithTraceAndMetricsLabels(ctx, label3)
+
+	assert.Equal(t, []attribute.KeyValue{label1, label2, label3}, otelsql.MetricsLabelsFromContext(ctx))
+	assert.Equal(t, []attribute.KeyValue{label1, label2, label3}, otelsql.TraceLabelsFromContext(ctx))
+}

ping_internal_test.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/metric/noop"
 	semconv "go.opentelemetry.io/otel/semconv/v1.20.0"
 
@@ -97,9 +98,10 @@ func TestPingStats(t *testing.T) {
 	t.Parallel()
 
 	testCases := []struct {
-		scenario string
-		ping     pingFunc
-		expected string
+		scenario  string
+		ping      pingFunc
+		ctxLabels []attribute.KeyValue
+		expected  string
 	}{
 		{
 			scenario: "error",
@@ -133,6 +135,24 @@ func TestPingStats(t *testing.T) {
 				}
 			]`,
 		},
+		{
+			scenario: "extra labels",
+			ping:     nopPing,
+			ctxLabels: []attribute.KeyValue{
+				attribute.String("extra", "label"),
+			},
+			expected: `[
+				{
+					"Name": "db.sql.client.calls{service.name=otelsql,instrumentation.name=ping_test,db.instance=test,db.operation=go.sql.ping,db.sql.status=OK,db.system=other_sql,extra=label}",
+					"Sum": 1
+				},
+				{
+					"Name": "db.sql.client.latency{service.name=otelsql,instrumentation.name=ping_test,db.instance=test,db.operation=go.sql.ping,db.sql.status=OK,db.system=other_sql,extra=label}",
+					"Sum": "<ignore-diff>",
+					"Count": 1
+				}
+			]`,
+		},
 	}
 
 	for _, tc := range testCases {
@@ -159,7 +179,9 @@ func TestPingStats(t *testing.T) {
 						pingStats(r),
 					}, tc.ping)
 
-					_ = ping(context.Background()) // nolint: errcheck
+					ctx := ContextWithMetricsLabels(context.Background(), tc.ctxLabels...)
+
+					_ = ping(ctx) // nolint: errcheck
 				})
 		})
 	}