Skip to content

Commit 9897125

Browse files
Manish Danginhatthm
Manish Dangi
authored and
nhatthm
committed
Support custom metrics and trace labels provided from the context
1 parent eaabba7 commit 9897125

18 files changed

+425
-110
lines changed

Diff for: ‎begin_internal_test.go

+23-1
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88

99
"github.com/stretchr/testify/assert"
1010
"github.com/stretchr/testify/require"
11+
"go.opentelemetry.io/otel/attribute"
1112
"go.opentelemetry.io/otel/metric/noop"
1213
semconv "go.opentelemetry.io/otel/semconv/v1.20.0"
1314

@@ -138,6 +139,7 @@ func TestBeginStats(t *testing.T) {
138139
testCases := []struct {
139140
scenario string
140141
begin beginFunc
142+
ctxLabel []attribute.KeyValue
141143
expected string
142144
}{
143145
{
@@ -172,6 +174,24 @@ func TestBeginStats(t *testing.T) {
172174
}
173175
]`,
174176
},
177+
{
178+
scenario: "extra labels",
179+
begin: nopBegin,
180+
ctxLabel: []attribute.KeyValue{
181+
attribute.String("extra", "label"),
182+
},
183+
expected: `[
184+
{
185+
"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}",
186+
"Sum": 1
187+
},
188+
{
189+
"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}",
190+
"Sum": "<ignore-diff>",
191+
"Count": 1
192+
}
193+
]`,
194+
},
175195
}
176196

177197
for _, tc := range testCases {
@@ -198,7 +218,9 @@ func TestBeginStats(t *testing.T) {
198218
beginStats(r),
199219
}, tc.begin)
200220

201-
_, _ = begin(context.Background(), driver.TxOptions{}) // nolint: errcheck
221+
ctx := ContextWithMetricsLabels(context.Background(), tc.ctxLabel...)
222+
223+
_, _ = begin(ctx, driver.TxOptions{}) // nolint: errcheck
202224
})
203225
})
204226
}

Diff for: ‎exec_internal_test.go

+23-1
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88

99
"github.com/stretchr/testify/assert"
1010
"github.com/stretchr/testify/require"
11+
"go.opentelemetry.io/otel/attribute"
1112
"go.opentelemetry.io/otel/metric/noop"
1213
semconv "go.opentelemetry.io/otel/semconv/v1.20.0"
1314

@@ -112,6 +113,7 @@ func TestExecStats(t *testing.T) {
112113
testCases := []struct {
113114
scenario string
114115
execer execContextFunc
116+
ctxLabel []attribute.KeyValue
115117
expected string
116118
}{
117119
{
@@ -146,6 +148,24 @@ func TestExecStats(t *testing.T) {
146148
}
147149
]`,
148150
},
151+
{
152+
scenario: "extra labels",
153+
execer: nopExecContext,
154+
ctxLabel: []attribute.KeyValue{
155+
attribute.String("extra", "label"),
156+
},
157+
expected: `[
158+
{
159+
"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}",
160+
"Sum": 1
161+
},
162+
{
163+
"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}",
164+
"Sum": "<ignore-diff>",
165+
"Count": 1
166+
}
167+
]`,
168+
},
149169
}
150170

151171
for _, tc := range testCases {
@@ -172,7 +192,9 @@ func TestExecStats(t *testing.T) {
172192
execStats(r, metricMethodExec),
173193
}, tc.execer)
174194

175-
_, _ = exec(context.Background(), "", nil) // nolint: errcheck
195+
ctx := ContextWithMetricsLabels(context.Background(), tc.ctxLabel...)
196+
197+
_, _ = exec(ctx, "", nil) // nolint: errcheck
176198
})
177199
})
178200
}

Diff for: ‎label.go

+89
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
package otelsql
2+
3+
import (
4+
"context"
5+
"sync"
6+
7+
"go.opentelemetry.io/otel/attribute"
8+
)
9+
10+
// Labeler is a helper to add attributes to a context.
11+
type Labeler struct {
12+
mu sync.Mutex
13+
attributes []attribute.KeyValue
14+
}
15+
16+
// Add attributes to a Labeler.
17+
func (l *Labeler) Add(ls ...attribute.KeyValue) {
18+
l.mu.Lock()
19+
defer l.mu.Unlock()
20+
21+
l.attributes = append(l.attributes, ls...)
22+
}
23+
24+
// Get returns a copy of the attributes added to the Labeler.
25+
func (l *Labeler) Get() []attribute.KeyValue {
26+
l.mu.Lock()
27+
defer l.mu.Unlock()
28+
29+
ret := make([]attribute.KeyValue, len(l.attributes))
30+
copy(ret, l.attributes)
31+
32+
return ret
33+
}
34+
35+
const (
36+
labelerCtxMetrics = labelerContextKey("metrics")
37+
labelerCtxTrace = labelerContextKey("trace")
38+
)
39+
40+
type labelerContextKey string
41+
42+
// MetricsLabelsFromContext retrieves the labels from the provided context.
43+
func MetricsLabelsFromContext(ctx context.Context) []attribute.KeyValue {
44+
l, _ := labelerFromContext(ctx, labelerCtxMetrics)
45+
46+
return l.Get()
47+
}
48+
49+
// ContextWithMetricsLabels returns a new context with the labels added to the Labeler.
50+
func ContextWithMetricsLabels(ctx context.Context, labels ...attribute.KeyValue) context.Context {
51+
return contextWithLabels(ctx, labelerCtxMetrics, labels...)
52+
}
53+
54+
// TraceLabelsFromContext retrieves the labels from the provided context.
55+
func TraceLabelsFromContext(ctx context.Context) []attribute.KeyValue {
56+
l, _ := labelerFromContext(ctx, labelerCtxTrace)
57+
58+
return l.Get()
59+
}
60+
61+
// ContextWithTraceLabels returns a new context with the labels added to the Labeler.
62+
func ContextWithTraceLabels(ctx context.Context, labels ...attribute.KeyValue) context.Context {
63+
return contextWithLabels(ctx, labelerCtxTrace, labels...)
64+
}
65+
66+
// ContextWithTraceAndMetricsLabels returns a new context with the labels added to the Labeler.
67+
func ContextWithTraceAndMetricsLabels(ctx context.Context, labels ...attribute.KeyValue) context.Context {
68+
ctx = ContextWithMetricsLabels(ctx, labels...)
69+
ctx = ContextWithTraceLabels(ctx, labels...)
70+
71+
return ctx
72+
}
73+
74+
func labelerFromContext(ctx context.Context, key labelerContextKey) (*Labeler, bool) { //nolint: unparam
75+
l, ok := ctx.Value(key).(*Labeler)
76+
if !ok {
77+
l = &Labeler{}
78+
}
79+
80+
return l, ok
81+
}
82+
83+
func contextWithLabels(ctx context.Context, key labelerContextKey, labels ...attribute.KeyValue) context.Context {
84+
l, _ := labelerFromContext(ctx, key)
85+
86+
l.Add(labels...)
87+
88+
return context.WithValue(ctx, key, l)
89+
}

Diff for: ‎label_test.go

+77
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
package otelsql_test
2+
3+
import (
4+
"context"
5+
"testing"
6+
7+
"github.com/stretchr/testify/assert"
8+
"go.opentelemetry.io/otel/attribute"
9+
10+
"go.nhat.io/otelsql"
11+
)
12+
13+
var (
14+
label1 = attribute.String("key1", "value1")
15+
label2 = attribute.Int("key2", 2)
16+
label3 = attribute.Bool("key3", true)
17+
)
18+
19+
func TestMetricsLabeler(t *testing.T) {
20+
t.Parallel()
21+
22+
ctx := context.Background()
23+
24+
labels := otelsql.MetricsLabelsFromContext(ctx)
25+
assert.Empty(t, labels)
26+
27+
ctx = otelsql.ContextWithMetricsLabels(ctx, label1, label2)
28+
29+
labels = otelsql.MetricsLabelsFromContext(ctx)
30+
assert.Equal(t, []attribute.KeyValue{label1, label2}, labels)
31+
32+
ctx = otelsql.ContextWithMetricsLabels(ctx, label3)
33+
assert.Equal(t, []attribute.KeyValue{label1, label2, label3}, otelsql.MetricsLabelsFromContext(ctx))
34+
35+
assert.Empty(t, otelsql.TraceLabelsFromContext(ctx))
36+
}
37+
38+
func TestTraceLabeler(t *testing.T) {
39+
t.Parallel()
40+
41+
ctx := context.Background()
42+
43+
labels := otelsql.TraceLabelsFromContext(ctx)
44+
assert.Empty(t, labels)
45+
46+
ctx = otelsql.ContextWithTraceLabels(ctx, label1, label2)
47+
48+
labels = otelsql.TraceLabelsFromContext(ctx)
49+
assert.Equal(t, []attribute.KeyValue{label1, label2}, labels)
50+
51+
ctx = otelsql.ContextWithTraceLabels(ctx, label3)
52+
assert.Equal(t, []attribute.KeyValue{label1, label2, label3}, otelsql.TraceLabelsFromContext(ctx))
53+
54+
assert.Empty(t, otelsql.MetricsLabelsFromContext(ctx))
55+
}
56+
57+
func TestMetricsAndTraceLabeler(t *testing.T) {
58+
t.Parallel()
59+
60+
ctx := context.Background()
61+
62+
labels := otelsql.MetricsLabelsFromContext(ctx)
63+
assert.Empty(t, labels)
64+
65+
labels = otelsql.TraceLabelsFromContext(ctx)
66+
assert.Empty(t, labels)
67+
68+
ctx = otelsql.ContextWithTraceAndMetricsLabels(ctx, label1, label2)
69+
70+
assert.Equal(t, []attribute.KeyValue{label1, label2}, otelsql.MetricsLabelsFromContext(ctx))
71+
assert.Equal(t, []attribute.KeyValue{label1, label2}, otelsql.TraceLabelsFromContext(ctx))
72+
73+
ctx = otelsql.ContextWithTraceAndMetricsLabels(ctx, label3)
74+
75+
assert.Equal(t, []attribute.KeyValue{label1, label2, label3}, otelsql.MetricsLabelsFromContext(ctx))
76+
assert.Equal(t, []attribute.KeyValue{label1, label2, label3}, otelsql.TraceLabelsFromContext(ctx))
77+
}

Diff for: ‎ping_internal_test.go

+26-4
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77

88
"github.com/stretchr/testify/assert"
99
"github.com/stretchr/testify/require"
10+
"go.opentelemetry.io/otel/attribute"
1011
"go.opentelemetry.io/otel/metric/noop"
1112
semconv "go.opentelemetry.io/otel/semconv/v1.20.0"
1213

@@ -97,9 +98,10 @@ func TestPingStats(t *testing.T) {
9798
t.Parallel()
9899

99100
testCases := []struct {
100-
scenario string
101-
ping pingFunc
102-
expected string
101+
scenario string
102+
ping pingFunc
103+
ctxLabels []attribute.KeyValue
104+
expected string
103105
}{
104106
{
105107
scenario: "error",
@@ -133,6 +135,24 @@ func TestPingStats(t *testing.T) {
133135
}
134136
]`,
135137
},
138+
{
139+
scenario: "extra labels",
140+
ping: nopPing,
141+
ctxLabels: []attribute.KeyValue{
142+
attribute.String("extra", "label"),
143+
},
144+
expected: `[
145+
{
146+
"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}",
147+
"Sum": 1
148+
},
149+
{
150+
"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}",
151+
"Sum": "<ignore-diff>",
152+
"Count": 1
153+
}
154+
]`,
155+
},
136156
}
137157

138158
for _, tc := range testCases {
@@ -159,7 +179,9 @@ func TestPingStats(t *testing.T) {
159179
pingStats(r),
160180
}, tc.ping)
161181

162-
_ = ping(context.Background()) // nolint: errcheck
182+
ctx := ContextWithMetricsLabels(context.Background(), tc.ctxLabels...)
183+
184+
_ = ping(ctx) // nolint: errcheck
163185
})
164186
})
165187
}

Diff for: ‎prepare_internal_test.go

+23-1
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88

99
"github.com/stretchr/testify/assert"
1010
"github.com/stretchr/testify/require"
11+
"go.opentelemetry.io/otel/attribute"
1112
"go.opentelemetry.io/otel/metric/noop"
1213
semconv "go.opentelemetry.io/otel/semconv/v1.20.0"
1314

@@ -138,6 +139,7 @@ func TestPrepareStats(t *testing.T) {
138139
testCases := []struct {
139140
scenario string
140141
prepare prepareContextFunc
142+
ctxLabel []attribute.KeyValue
141143
expected string
142144
}{
143145
{
@@ -172,6 +174,24 @@ func TestPrepareStats(t *testing.T) {
172174
}
173175
]`,
174176
},
177+
{
178+
scenario: "extra labels",
179+
prepare: nopPrepareContext,
180+
ctxLabel: []attribute.KeyValue{
181+
attribute.String("extra", "label"),
182+
},
183+
expected: `[
184+
{
185+
"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}",
186+
"Sum": 1
187+
},
188+
{
189+
"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}",
190+
"Sum": "<ignore-diff>",
191+
"Count": 1
192+
}
193+
]`,
194+
},
175195
}
176196

177197
for _, tc := range testCases {
@@ -198,7 +218,9 @@ func TestPrepareStats(t *testing.T) {
198218
prepareStats(r),
199219
}, tc.prepare)
200220

201-
_, _ = prepare(context.Background(), "") // nolint: errcheck
221+
ctx := ContextWithMetricsLabels(context.Background(), tc.ctxLabel...)
222+
223+
_, _ = prepare(ctx, "") // nolint: errcheck
202224
})
203225
})
204226
}

0 commit comments

Comments
 (0)