Skip to content

Commit 42810e8

Browse files
vearutopvearutop
committed
Add option to format span name with context
1 parent 41ee6eb commit 42810e8

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
@@ -226,3 +226,37 @@ func (s *svc) GetDevice(ctx context.Context, id int) (*Device, error) {
226226
return device
227227
}
228228
```
229+
230+
With context you can also expose more relevant span names:
231+
232+
```go
233+
// Define a function to build meaningful span name from context
234+
formatSpanName := func(ctx context.Context, baseName string) string {
235+
ctxName := ctx.Value("txName")
236+
if txName, ok := ctxName.(string); ok {
237+
return fmt.Sprintf("%s:%s", txName, baseName)
238+
}
239+
return baseName
240+
}
241+
242+
// Register our ocsql wrapper for the provided Postgres driver.
243+
driverName, err := ocsql.Register("postgres", ocsql.WithAllTraceOptions(),
244+
ocsql.WithFormatSpanNameFunc(formatSpanName),
245+
)
246+
if err != nil {
247+
...
248+
}
249+
250+
// Connect to a Postgres database using the ocsql driver wrapper.
251+
db, err := sql.Open(driverName, "postgres://localhost:5432/my_database")
252+
if err != nil {
253+
...
254+
}
255+
256+
257+
// Add relevant information to context
258+
ctx := context.WithValue(context.Background(), "txName", "orderCreation")
259+
260+
// Use context with database
261+
db.ExecContext(ctx, "INSERT INTO order ...")
262+
```

‎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
}
@@ -396,11 +404,10 @@ func (c *ocConn) BeginTx(ctx context.Context, opts driver.TxOptions) (tx driver.
396404

397405
if ctx == nil || ctx == context.TODO() {
398406
ctx = context.Background()
399-
_, span = trace.StartSpan(ctx, "sql:begin_transaction", trace.WithSpanKind(trace.SpanKindClient))
400407
attrs = append(attrs, attrMissingContext)
401-
} else {
402-
_, span = trace.StartSpan(ctx, "sql:begin_transaction", trace.WithSpanKind(trace.SpanKindClient))
403408
}
409+
_, span = trace.StartSpan(ctx, c.options.formatSpanName(ctx, "sql:begin_transaction"),
410+
trace.WithSpanKind(trace.SpanKindClient))
404411
defer func() {
405412
if len(attrs) > 0 {
406413
span.AddAttributes(attrs...)
@@ -441,7 +448,8 @@ type ocResult struct {
441448

442449
func (r ocResult) LastInsertId() (id int64, err error) {
443450
if r.options.LastInsertID {
444-
_, span := trace.StartSpan(r.ctx, "sql:last_insert_id", trace.WithSpanKind(trace.SpanKindClient))
451+
_, span := trace.StartSpan(r.ctx, r.options.formatSpanName(r.ctx, "sql:last_insert_id"),
452+
trace.WithSpanKind(trace.SpanKindClient))
445453
if len(r.options.DefaultAttributes) > 0 {
446454
span.AddAttributes(r.options.DefaultAttributes...)
447455
}
@@ -457,7 +465,8 @@ func (r ocResult) LastInsertId() (id int64, err error) {
457465

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

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

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

586+
spanName := s.options.formatSpanName(ctx, "sql:exec")
575587
var span *trace.Span
576588
if parentSpan == nil {
577-
ctx, span = trace.StartSpan(ctx, "sql:exec", trace.WithSpanKind(trace.SpanKindClient))
589+
ctx, span = trace.StartSpan(ctx, spanName, trace.WithSpanKind(trace.SpanKindClient))
578590
} else {
579-
_, span = trace.StartSpan(ctx, "sql:exec", trace.WithSpanKind(trace.SpanKindClient))
591+
_, span = trace.StartSpan(ctx, spanName, trace.WithSpanKind(trace.SpanKindClient))
580592
}
581593
attrs := append([]trace.Attribute(nil), s.options.DefaultAttributes...)
582594
if s.options.Query {
@@ -613,9 +625,11 @@ func (s ocStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (row
613625

614626
var span *trace.Span
615627
if parentSpan == nil {
616-
ctx, span = trace.StartSpan(ctx, "sql:query", trace.WithSpanKind(trace.SpanKindClient))
628+
ctx, span = trace.StartSpan(ctx, s.options.formatSpanName(ctx, "sql:query"),
629+
trace.WithSpanKind(trace.SpanKindClient))
617630
} else {
618-
_, span = trace.StartSpan(ctx, "sql:query", trace.WithSpanKind(trace.SpanKindClient))
631+
_, span = trace.StartSpan(ctx, s.options.formatSpanName(ctx, "sql:query"),
632+
trace.WithSpanKind(trace.SpanKindClient))
619633
}
620634
attrs := append([]trace.Attribute(nil), s.options.DefaultAttributes...)
621635
if s.options.Query {
@@ -730,7 +744,8 @@ func (r ocRows) Columns() []string {
730744

731745
func (r ocRows) Close() (err error) {
732746
if r.options.RowsClose {
733-
_, span := trace.StartSpan(r.ctx, "sql:rows_close", trace.WithSpanKind(trace.SpanKindClient))
747+
_, span := trace.StartSpan(r.ctx, r.options.formatSpanName(r.ctx, "sql:rows_close"),
748+
trace.WithSpanKind(trace.SpanKindClient))
734749
if len(r.options.DefaultAttributes) > 0 {
735750
span.AddAttributes(r.options.DefaultAttributes...)
736751
}
@@ -746,7 +761,8 @@ func (r ocRows) Close() (err error) {
746761

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

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

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

‎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

@@ -53,12 +55,26 @@ type TraceOptions struct {
5355

5456
// DisableErrSkip, if set to true, will suppress driver.ErrSkip errors in spans.
5557
DisableErrSkip bool
58+
59+
// FormatSpanName holds the function to use for generating the span name
60+
// from the information found in the context.
61+
FormatSpanName func(ctx context.Context, baseName string) string
62+
}
63+
64+
// formatSpanName decorates span name
65+
func (o *TraceOptions) formatSpanName(ctx context.Context, baseName string) string {
66+
if o.FormatSpanName != nil {
67+
return o.FormatSpanName(ctx, baseName)
68+
}
69+
return baseName
5670
}
5771

5872
// WithAllTraceOptions enables all available trace options.
5973
func WithAllTraceOptions() TraceOption {
6074
return func(o *TraceOptions) {
75+
f := o.FormatSpanName
6176
*o = AllTraceOptions
77+
o.FormatSpanName = f
6278
}
6379
}
6480

@@ -166,3 +182,11 @@ func WithDisableErrSkip(b bool) TraceOption {
166182
o.DisableErrSkip = b
167183
}
168184
}
185+
186+
// WithFormatSpanNameFunc enables custom span names, built with information from the context.
187+
// When context is not available context.Background() is used.
188+
func WithFormatSpanNameFunc(f func(ctx context.Context, baseName string) string) TraceOption {
189+
return func(o *TraceOptions) {
190+
o.FormatSpanName = f
191+
}
192+
}

0 commit comments

Comments
 (0)