Skip to content

Commit de664d7

Browse files
vearutopvearutop
committed
add option to format span name with context
1 parent 28a9935 commit de664d7

File tree

3 files changed

+100
-24
lines changed

3 files changed

+100
-24
lines changed

‎README.md‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -196,3 +196,37 @@ func (s *svc) GetDevice(ctx context.Context, id int) (*Device, error) {
196196
return device
197197
}
198198
```
199+
200+
With context you can also expose more relevant span names:
201+
202+
```go
203+
// Define a function to build meaningful span name from context
204+
formatSpanName := func(ctx context.Context, baseName string) string {
205+
ctxName := ctx.Value("txName")
206+
if txName, ok := ctxName.(string); ok {
207+
return fmt.Sprintf("%s:%s", txName, baseName)
208+
}
209+
return baseName
210+
}
211+
212+
// Register our ocsql wrapper for the provided Postgres driver.
213+
driverName, err := ocsql.Register("postgres", ocsql.WithAllTraceOptions(),
214+
ocsql.WithFormatSpanNameFunc(formatSpanName),
215+
)
216+
if err != nil {
217+
...
218+
}
219+
220+
// Connect to a Postgres database using the ocsql driver wrapper.
221+
db, err := sql.Open(driverName, "postgres://localhost:5432/my_database")
222+
if err != nil {
223+
...
224+
}
225+
226+
227+
// Add relevant information to context
228+
ctx := context.WithValue(context.Background(), "txName", "orderCreation")
229+
230+
// Use context with database
231+
db.ExecContext(ctx, "INSERT INTO order ...")
232+
```

‎driver.go‎

Lines changed: 42 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,8 @@ func (c ocConn) Ping(ctx context.Context) (err error) {
127127

128128
if c.options.Ping && (c.options.AllowRoot || trace.FromContext(ctx) != nil) {
129129
var span *trace.Span
130-
ctx, span = trace.StartSpan(ctx, "sql:ping", trace.WithSpanKind(trace.SpanKindClient))
130+
ctx, span = trace.StartSpan(ctx, c.options.formatSpanName(ctx, "sql:ping"),
131+
trace.WithSpanKind(trace.SpanKindClient))
131132
if len(c.options.DefaultAttributes) > 0 {
132133
span.AddAttributes(c.options.DefaultAttributes...)
133134
}
@@ -158,7 +159,8 @@ func (c ocConn) Exec(query string, args []driver.Value) (res driver.Result, err
158159
return exec.Exec(query, args)
159160
}
160161

161-
ctx, span := trace.StartSpan(context.Background(), "sql:exec", trace.WithSpanKind(trace.SpanKindClient))
162+
ctx, span := trace.StartSpan(context.Background(), c.options.formatSpanName(context.Background(), "sql:exec"),
163+
trace.WithSpanKind(trace.SpanKindClient))
162164
attrs := make([]trace.Attribute, 0, len(c.options.DefaultAttributes)+2)
163165
attrs = append(attrs, c.options.DefaultAttributes...)
164166
attrs = append(
@@ -200,11 +202,12 @@ func (c ocConn) ExecContext(ctx context.Context, query string, args []driver.Nam
200202
return execCtx.ExecContext(ctx, query, args)
201203
}
202204

205+
spanName := c.options.formatSpanName(ctx, "sql:exec")
203206
var span *trace.Span
204207
if parentSpan == nil {
205-
ctx, span = trace.StartSpan(ctx, "sql:exec", trace.WithSpanKind(trace.SpanKindClient))
208+
ctx, span = trace.StartSpan(ctx, spanName, trace.WithSpanKind(trace.SpanKindClient))
206209
} else {
207-
_, span = trace.StartSpan(ctx, "sql:exec", trace.WithSpanKind(trace.SpanKindClient))
210+
_, span = trace.StartSpan(ctx, spanName, trace.WithSpanKind(trace.SpanKindClient))
208211
}
209212
attrs := append([]trace.Attribute(nil), c.options.DefaultAttributes...)
210213
if c.options.Query {
@@ -238,7 +241,8 @@ func (c ocConn) Query(query string, args []driver.Value) (rows driver.Rows, err
238241
return queryer.Query(query, args)
239242
}
240243

241-
ctx, span := trace.StartSpan(context.Background(), "sql:query", trace.WithSpanKind(trace.SpanKindClient))
244+
ctx, span := trace.StartSpan(context.Background(), c.options.formatSpanName(context.Background(), "sql:query"),
245+
trace.WithSpanKind(trace.SpanKindClient))
242246
attrs := make([]trace.Attribute, 0, len(c.options.DefaultAttributes)+2)
243247
attrs = append(attrs, c.options.DefaultAttributes...)
244248
attrs = append(
@@ -283,9 +287,11 @@ func (c ocConn) QueryContext(ctx context.Context, query string, args []driver.Na
283287

284288
var span *trace.Span
285289
if parentSpan == nil {
286-
ctx, span = trace.StartSpan(ctx, "sql:query", trace.WithSpanKind(trace.SpanKindClient))
290+
ctx, span = trace.StartSpan(ctx, c.options.formatSpanName(ctx, "sql:query"),
291+
trace.WithSpanKind(trace.SpanKindClient))
287292
} else {
288-
_, span = trace.StartSpan(ctx, "sql:query", trace.WithSpanKind(trace.SpanKindClient))
293+
_, span = trace.StartSpan(ctx, c.options.formatSpanName(ctx, "sql:query"),
294+
trace.WithSpanKind(trace.SpanKindClient))
289295
}
290296
attrs := append([]trace.Attribute(nil), c.options.DefaultAttributes...)
291297
if c.options.Query {
@@ -316,7 +322,8 @@ func (c ocConn) Prepare(query string) (stmt driver.Stmt, err error) {
316322
defer recordCallStats(context.Background(), "go.sql.prepare")(err)
317323

318324
if c.options.AllowRoot {
319-
_, span := trace.StartSpan(context.Background(), "sql:prepare", trace.WithSpanKind(trace.SpanKindClient))
325+
_, span := trace.StartSpan(context.Background(), c.options.formatSpanName(context.Background(), "sql:prepare"),
326+
trace.WithSpanKind(trace.SpanKindClient))
320327
attrs := make([]trace.Attribute, 0, len(c.options.DefaultAttributes)+1)
321328
attrs = append(attrs, c.options.DefaultAttributes...)
322329
attrs = append(attrs, attrMissingContext)
@@ -354,7 +361,8 @@ func (c *ocConn) PrepareContext(ctx context.Context, query string) (stmt driver.
354361
var span *trace.Span
355362
attrs := append([]trace.Attribute(nil), c.options.DefaultAttributes...)
356363
if c.options.AllowRoot || trace.FromContext(ctx) != nil {
357-
ctx, span = trace.StartSpan(ctx, "sql:prepare", trace.WithSpanKind(trace.SpanKindClient))
364+
ctx, span = trace.StartSpan(ctx, c.options.formatSpanName(ctx, "sql:prepare"),
365+
trace.WithSpanKind(trace.SpanKindClient))
358366
if c.options.Query {
359367
attrs = append(attrs, trace.StringAttribute("sql.query", query))
360368
}
@@ -400,11 +408,10 @@ func (c *ocConn) BeginTx(ctx context.Context, opts driver.TxOptions) (tx driver.
400408
}()
401409
if ctx == nil || ctx == context.TODO() {
402410
ctx = context.Background()
403-
_, span = trace.StartSpan(ctx, "sql:begin_transaction", trace.WithSpanKind(trace.SpanKindClient))
404411
attrs = append(attrs, attrMissingContext)
405-
} else {
406-
_, span = trace.StartSpan(ctx, "sql:begin_transaction", trace.WithSpanKind(trace.SpanKindClient))
407412
}
413+
_, span = trace.StartSpan(ctx, c.options.formatSpanName(ctx, "sql:begin_transaction"),
414+
trace.WithSpanKind(trace.SpanKindClient))
408415
defer span.End()
409416

410417
if connBeginTx, ok := c.parent.(driver.ConnBeginTx); ok {
@@ -440,7 +447,8 @@ type ocResult struct {
440447

441448
func (r ocResult) LastInsertId() (id int64, err error) {
442449
if r.options.LastInsertID {
443-
_, span := trace.StartSpan(r.ctx, "sql:last_insert_id", trace.WithSpanKind(trace.SpanKindClient))
450+
_, span := trace.StartSpan(r.ctx, r.options.formatSpanName(r.ctx, "sql:last_insert_id"),
451+
trace.WithSpanKind(trace.SpanKindClient))
444452
if len(r.options.DefaultAttributes) > 0 {
445453
span.AddAttributes(r.options.DefaultAttributes...)
446454
}
@@ -456,7 +464,8 @@ func (r ocResult) LastInsertId() (id int64, err error) {
456464

457465
func (r ocResult) RowsAffected() (cnt int64, err error) {
458466
if r.options.RowsAffected {
459-
_, span := trace.StartSpan(r.ctx, "sql:rows_affected", trace.WithSpanKind(trace.SpanKindClient))
467+
_, span := trace.StartSpan(r.ctx, r.options.formatSpanName(r.ctx, "sql:rows_affected"),
468+
trace.WithSpanKind(trace.SpanKindClient))
460469
if len(r.options.DefaultAttributes) > 0 {
461470
span.AddAttributes(r.options.DefaultAttributes...)
462471
}
@@ -484,7 +493,8 @@ func (s ocStmt) Exec(args []driver.Value) (res driver.Result, err error) {
484493
return s.parent.Exec(args)
485494
}
486495

487-
ctx, span := trace.StartSpan(context.Background(), "sql:exec", trace.WithSpanKind(trace.SpanKindClient))
496+
ctx, span := trace.StartSpan(context.Background(), s.options.formatSpanName(context.Background(), "sql:exec"),
497+
trace.WithSpanKind(trace.SpanKindClient))
488498
attrs := make([]trace.Attribute, 0, len(s.options.DefaultAttributes)+2)
489499
attrs = append(attrs, s.options.DefaultAttributes...)
490500
attrs = append(
@@ -531,7 +541,8 @@ func (s ocStmt) Query(args []driver.Value) (rows driver.Rows, err error) {
531541
return s.parent.Query(args)
532542
}
533543

534-
ctx, span := trace.StartSpan(context.Background(), "sql:query", trace.WithSpanKind(trace.SpanKindClient))
544+
ctx, span := trace.StartSpan(context.Background(), s.options.formatSpanName(context.Background(), "sql:query"),
545+
trace.WithSpanKind(trace.SpanKindClient))
535546
attrs := make([]trace.Attribute, 0, len(s.options.DefaultAttributes)+2)
536547
attrs = append(attrs, s.options.DefaultAttributes...)
537548
attrs = append(
@@ -571,11 +582,12 @@ func (s ocStmt) ExecContext(ctx context.Context, args []driver.NamedValue) (res
571582
return s.parent.(driver.StmtExecContext).ExecContext(ctx, args)
572583
}
573584

585+
spanName := s.options.formatSpanName(ctx, "sql:exec")
574586
var span *trace.Span
575587
if parentSpan == nil {
576-
ctx, span = trace.StartSpan(ctx, "sql:exec", trace.WithSpanKind(trace.SpanKindClient))
588+
ctx, span = trace.StartSpan(ctx, spanName, trace.WithSpanKind(trace.SpanKindClient))
577589
} else {
578-
_, span = trace.StartSpan(ctx, "sql:exec", trace.WithSpanKind(trace.SpanKindClient))
590+
_, span = trace.StartSpan(ctx, spanName, trace.WithSpanKind(trace.SpanKindClient))
579591
}
580592
attrs := append([]trace.Attribute(nil), s.options.DefaultAttributes...)
581593
if s.options.Query {
@@ -612,9 +624,11 @@ func (s ocStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (row
612624

613625
var span *trace.Span
614626
if parentSpan == nil {
615-
ctx, span = trace.StartSpan(ctx, "sql:query", trace.WithSpanKind(trace.SpanKindClient))
627+
ctx, span = trace.StartSpan(ctx, s.options.formatSpanName(ctx, "sql:query"),
628+
trace.WithSpanKind(trace.SpanKindClient))
616629
} else {
617-
_, span = trace.StartSpan(ctx, "sql:query", trace.WithSpanKind(trace.SpanKindClient))
630+
_, span = trace.StartSpan(ctx, s.options.formatSpanName(ctx, "sql:query"),
631+
trace.WithSpanKind(trace.SpanKindClient))
618632
}
619633
attrs := append([]trace.Attribute(nil), s.options.DefaultAttributes...)
620634
if s.options.Query {
@@ -729,7 +743,8 @@ func (r ocRows) Columns() []string {
729743

730744
func (r ocRows) Close() (err error) {
731745
if r.options.RowsClose {
732-
_, span := trace.StartSpan(r.ctx, "sql:rows_close", trace.WithSpanKind(trace.SpanKindClient))
746+
_, span := trace.StartSpan(r.ctx, r.options.formatSpanName(r.ctx, "sql:rows_close"),
747+
trace.WithSpanKind(trace.SpanKindClient))
733748
if len(r.options.DefaultAttributes) > 0 {
734749
span.AddAttributes(r.options.DefaultAttributes...)
735750
}
@@ -745,7 +760,8 @@ func (r ocRows) Close() (err error) {
745760

746761
func (r ocRows) Next(dest []driver.Value) (err error) {
747762
if r.options.RowsNext {
748-
_, span := trace.StartSpan(r.ctx, "sql:rows_next", trace.WithSpanKind(trace.SpanKindClient))
763+
_, span := trace.StartSpan(r.ctx, r.options.formatSpanName(r.ctx, "sql:rows_next"),
764+
trace.WithSpanKind(trace.SpanKindClient))
749765
if len(r.options.DefaultAttributes) > 0 {
750766
span.AddAttributes(r.options.DefaultAttributes...)
751767
}
@@ -801,7 +817,8 @@ type ocTx struct {
801817
func (t ocTx) Commit() (err error) {
802818
defer recordCallStats(context.Background(), "go.sql.commit")(err)
803819

804-
_, span := trace.StartSpan(t.ctx, "sql:commit", trace.WithSpanKind(trace.SpanKindClient))
820+
_, span := trace.StartSpan(t.ctx, t.options.formatSpanName(t.ctx, "sql:commit"),
821+
trace.WithSpanKind(trace.SpanKindClient))
805822
if len(t.options.DefaultAttributes) > 0 {
806823
span.AddAttributes(t.options.DefaultAttributes...)
807824
}
@@ -817,7 +834,8 @@ func (t ocTx) Commit() (err error) {
817834
func (t ocTx) Rollback() (err error) {
818835
defer recordCallStats(context.Background(), "go.sql.rollback")(err)
819836

820-
_, span := trace.StartSpan(t.ctx, "sql:rollback", trace.WithSpanKind(trace.SpanKindClient))
837+
_, span := trace.StartSpan(t.ctx, t.options.formatSpanName(t.ctx, "sql:rollback"),
838+
trace.WithSpanKind(trace.SpanKindClient))
821839
if len(t.options.DefaultAttributes) > 0 {
822840
span.AddAttributes(t.options.DefaultAttributes...)
823841
}

‎options.go‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
package ocsql
22

33
import (
4+
"context"
5+
46
"go.opencensus.io/trace"
57
)
68

@@ -50,12 +52,26 @@ type TraceOptions struct {
5052

5153
// DefaultAttributes will be set to each span as default.
5254
DefaultAttributes []trace.Attribute
55+
56+
// FormatSpanName holds the function to use for generating the span name
57+
// from the information found in the context.
58+
FormatSpanName func(ctx context.Context, baseName string) string
59+
}
60+
61+
// formatSpanName decorates span name
62+
func (o *TraceOptions) formatSpanName(ctx context.Context, baseName string) string {
63+
if o.FormatSpanName != nil {
64+
return o.FormatSpanName(ctx, baseName)
65+
}
66+
return baseName
5367
}
5468

5569
// WithAllTraceOptions enables all available trace options.
5670
func WithAllTraceOptions() TraceOption {
5771
return func(o *TraceOptions) {
72+
f := o.FormatSpanName
5873
*o = AllTraceOptions
74+
o.FormatSpanName = f
5975
}
6076
}
6177

@@ -156,3 +172,11 @@ func WithDefaultAttributes(attrs ...trace.Attribute) TraceOption {
156172
o.DefaultAttributes = attrs
157173
}
158174
}
175+
176+
// WithFormatSpanNameFunc enables custom span names, built with information from the context.
177+
// When context is not available context.Background() is used.
178+
func WithFormatSpanNameFunc(f func(ctx context.Context, baseName string) string) TraceOption {
179+
return func(o *TraceOptions) {
180+
o.FormatSpanName = f
181+
}
182+
}

0 commit comments

Comments
 (0)