Support custom metrics and trace labels provided from the context

nhatthm · Mar 27, 2024 · 9897125 · 9897125
1 parent eaabba7
commit 9897125
Show file tree

Hide file tree

Showing 18 changed files with 425 additions and 110 deletions.
diff --git a/begin_internal_test.go b/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
 				})
 		})
 	}

diff --git a/exec_internal_test.go b/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
 				})
 		})
 	}

diff --git a/label.go b/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)
+}
diff --git a/label_test.go b/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))
+}
diff --git a/ping_internal_test.go b/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
 				})
 		})
 	}

diff --git a/prepare_internal_test.go b/prepare_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 TestPrepareStats(t *testing.T) {
 	testCases := []struct {
 		scenario string
 		prepare  prepareContextFunc
+		ctxLabel []attribute.KeyValue
 		expected string
 	}{
 		{
@@ -172,6 +174,24 @@ func TestPrepareStats(t *testing.T) {
 				}
 			]`,
 		},
+		{
+			scenario: "extra labels",
+			prepare:  nopPrepareContext,
+			ctxLabel: []attribute.KeyValue{
+				attribute.String("extra", "label"),
+			},
+			expected: `[
+				{
+					"Name": "db.sql.client.calls{service.name=otelsql,instrumentation.name=prepare_test,db.instance=test,db.operation=go.sql.prepare,db.sql.status=OK,db.system=other_sql,extra=label}",
+					"Sum": 1
+				},
+				{
+					"Name": "db.sql.client.latency{service.name=otelsql,instrumentation.name=prepare_test,db.instance=test,db.operation=go.sql.prepare,db.sql.status=OK,db.system=other_sql,extra=label}",
+					"Sum": "<ignore-diff>",
+					"Count": 1
+				}
+			]`,
+		},
 	}
 
 	for _, tc := range testCases {
@@ -198,7 +218,9 @@ func TestPrepareStats(t *testing.T) {
 						prepareStats(r),
 					}, tc.prepare)
 
-					_, _ = prepare(context.Background(), "") // nolint: errcheck
+					ctx := ContextWithMetricsLabels(context.Background(), tc.ctxLabel...)
+
+					_, _ = prepare(ctx, "") // nolint: errcheck
 				})
 		})
 	}