Skip to content

Commit e766b83

Browse files
odeke-emodeke-em
authored
Merge pull request #40 from alexsn/with_sampler
Add WithSampler option propagating to trace.StartSpan call
2 parents 41ee6eb + 420244e commit e766b83

File tree

2 files changed

+103
-24
lines changed

2 files changed

+103
-24
lines changed

‎driver.go‎

Lines changed: 93 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,10 @@ 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, "sql:ping",
131+
trace.WithSpanKind(trace.SpanKindClient),
132+
trace.WithSampler(c.options.Sampler),
133+
)
131134
if len(c.options.DefaultAttributes) > 0 {
132135
span.AddAttributes(c.options.DefaultAttributes...)
133136
}
@@ -158,7 +161,10 @@ func (c ocConn) Exec(query string, args []driver.Value) (res driver.Result, err
158161
return exec.Exec(query, args)
159162
}
160163

161-
ctx, span := trace.StartSpan(context.Background(), "sql:exec", trace.WithSpanKind(trace.SpanKindClient))
164+
ctx, span := trace.StartSpan(context.Background(), "sql:exec",
165+
trace.WithSpanKind(trace.SpanKindClient),
166+
trace.WithSampler(c.options.Sampler),
167+
)
162168
attrs := make([]trace.Attribute, 0, len(c.options.DefaultAttributes)+2)
163169
attrs = append(attrs, c.options.DefaultAttributes...)
164170
attrs = append(
@@ -202,9 +208,15 @@ func (c ocConn) ExecContext(ctx context.Context, query string, args []driver.Nam
202208

203209
var span *trace.Span
204210
if parentSpan == nil {
205-
ctx, span = trace.StartSpan(ctx, "sql:exec", trace.WithSpanKind(trace.SpanKindClient))
211+
ctx, span = trace.StartSpan(ctx, "sql:exec",
212+
trace.WithSpanKind(trace.SpanKindClient),
213+
trace.WithSampler(c.options.Sampler),
214+
)
206215
} else {
207-
_, span = trace.StartSpan(ctx, "sql:exec", trace.WithSpanKind(trace.SpanKindClient))
216+
_, span = trace.StartSpan(ctx, "sql:exec",
217+
trace.WithSpanKind(trace.SpanKindClient),
218+
trace.WithSampler(c.options.Sampler),
219+
)
208220
}
209221
attrs := append([]trace.Attribute(nil), c.options.DefaultAttributes...)
210222
if c.options.Query {
@@ -238,7 +250,10 @@ func (c ocConn) Query(query string, args []driver.Value) (rows driver.Rows, err
238250
return queryer.Query(query, args)
239251
}
240252

241-
ctx, span := trace.StartSpan(context.Background(), "sql:query", trace.WithSpanKind(trace.SpanKindClient))
253+
ctx, span := trace.StartSpan(context.Background(), "sql:query",
254+
trace.WithSpanKind(trace.SpanKindClient),
255+
trace.WithSampler(c.options.Sampler),
256+
)
242257
attrs := make([]trace.Attribute, 0, len(c.options.DefaultAttributes)+2)
243258
attrs = append(attrs, c.options.DefaultAttributes...)
244259
attrs = append(
@@ -283,9 +298,15 @@ func (c ocConn) QueryContext(ctx context.Context, query string, args []driver.Na
283298

284299
var span *trace.Span
285300
if parentSpan == nil {
286-
ctx, span = trace.StartSpan(ctx, "sql:query", trace.WithSpanKind(trace.SpanKindClient))
301+
ctx, span = trace.StartSpan(ctx, "sql:query",
302+
trace.WithSpanKind(trace.SpanKindClient),
303+
trace.WithSampler(c.options.Sampler),
304+
)
287305
} else {
288-
_, span = trace.StartSpan(ctx, "sql:query", trace.WithSpanKind(trace.SpanKindClient))
306+
_, span = trace.StartSpan(ctx, "sql:query",
307+
trace.WithSpanKind(trace.SpanKindClient),
308+
trace.WithSampler(c.options.Sampler),
309+
)
289310
}
290311
attrs := append([]trace.Attribute(nil), c.options.DefaultAttributes...)
291312
if c.options.Query {
@@ -316,7 +337,10 @@ func (c ocConn) Prepare(query string) (stmt driver.Stmt, err error) {
316337
defer recordCallStats(context.Background(), "go.sql.prepare")(err)
317338

318339
if c.options.AllowRoot {
319-
_, span := trace.StartSpan(context.Background(), "sql:prepare", trace.WithSpanKind(trace.SpanKindClient))
340+
_, span := trace.StartSpan(context.Background(), "sql:prepare",
341+
trace.WithSpanKind(trace.SpanKindClient),
342+
trace.WithSampler(c.options.Sampler),
343+
)
320344
attrs := make([]trace.Attribute, 0, len(c.options.DefaultAttributes)+1)
321345
attrs = append(attrs, c.options.DefaultAttributes...)
322346
attrs = append(attrs, attrMissingContext)
@@ -354,7 +378,10 @@ func (c *ocConn) PrepareContext(ctx context.Context, query string) (stmt driver.
354378
var span *trace.Span
355379
attrs := append([]trace.Attribute(nil), c.options.DefaultAttributes...)
356380
if c.options.AllowRoot || trace.FromContext(ctx) != nil {
357-
ctx, span = trace.StartSpan(ctx, "sql:prepare", trace.WithSpanKind(trace.SpanKindClient))
381+
ctx, span = trace.StartSpan(ctx, "sql:prepare",
382+
trace.WithSpanKind(trace.SpanKindClient),
383+
trace.WithSampler(c.options.Sampler),
384+
)
358385
if c.options.Query {
359386
attrs = append(attrs, trace.StringAttribute("sql.query", query))
360387
}
@@ -396,10 +423,16 @@ func (c *ocConn) BeginTx(ctx context.Context, opts driver.TxOptions) (tx driver.
396423

397424
if ctx == nil || ctx == context.TODO() {
398425
ctx = context.Background()
399-
_, span = trace.StartSpan(ctx, "sql:begin_transaction", trace.WithSpanKind(trace.SpanKindClient))
426+
_, span = trace.StartSpan(ctx, "sql:begin_transaction",
427+
trace.WithSpanKind(trace.SpanKindClient),
428+
trace.WithSampler(c.options.Sampler),
429+
)
400430
attrs = append(attrs, attrMissingContext)
401431
} else {
402-
_, span = trace.StartSpan(ctx, "sql:begin_transaction", trace.WithSpanKind(trace.SpanKindClient))
432+
_, span = trace.StartSpan(ctx, "sql:begin_transaction",
433+
trace.WithSpanKind(trace.SpanKindClient),
434+
trace.WithSampler(c.options.Sampler),
435+
)
403436
}
404437
defer func() {
405438
if len(attrs) > 0 {
@@ -441,7 +474,10 @@ type ocResult struct {
441474

442475
func (r ocResult) LastInsertId() (id int64, err error) {
443476
if r.options.LastInsertID {
444-
_, span := trace.StartSpan(r.ctx, "sql:last_insert_id", trace.WithSpanKind(trace.SpanKindClient))
477+
_, span := trace.StartSpan(r.ctx, "sql:last_insert_id",
478+
trace.WithSpanKind(trace.SpanKindClient),
479+
trace.WithSampler(r.options.Sampler),
480+
)
445481
if len(r.options.DefaultAttributes) > 0 {
446482
span.AddAttributes(r.options.DefaultAttributes...)
447483
}
@@ -457,7 +493,10 @@ func (r ocResult) LastInsertId() (id int64, err error) {
457493

458494
func (r ocResult) RowsAffected() (cnt int64, err error) {
459495
if r.options.RowsAffected {
460-
_, span := trace.StartSpan(r.ctx, "sql:rows_affected", trace.WithSpanKind(trace.SpanKindClient))
496+
_, span := trace.StartSpan(r.ctx, "sql:rows_affected",
497+
trace.WithSpanKind(trace.SpanKindClient),
498+
trace.WithSampler(r.options.Sampler),
499+
)
461500
if len(r.options.DefaultAttributes) > 0 {
462501
span.AddAttributes(r.options.DefaultAttributes...)
463502
}
@@ -485,7 +524,10 @@ func (s ocStmt) Exec(args []driver.Value) (res driver.Result, err error) {
485524
return s.parent.Exec(args)
486525
}
487526

488-
ctx, span := trace.StartSpan(context.Background(), "sql:exec", trace.WithSpanKind(trace.SpanKindClient))
527+
ctx, span := trace.StartSpan(context.Background(), "sql:exec",
528+
trace.WithSpanKind(trace.SpanKindClient),
529+
trace.WithSampler(s.options.Sampler),
530+
)
489531
attrs := make([]trace.Attribute, 0, len(s.options.DefaultAttributes)+2)
490532
attrs = append(attrs, s.options.DefaultAttributes...)
491533
attrs = append(
@@ -532,7 +574,10 @@ func (s ocStmt) Query(args []driver.Value) (rows driver.Rows, err error) {
532574
return s.parent.Query(args)
533575
}
534576

535-
ctx, span := trace.StartSpan(context.Background(), "sql:query", trace.WithSpanKind(trace.SpanKindClient))
577+
ctx, span := trace.StartSpan(context.Background(), "sql:query",
578+
trace.WithSpanKind(trace.SpanKindClient),
579+
trace.WithSampler(s.options.Sampler),
580+
)
536581
attrs := make([]trace.Attribute, 0, len(s.options.DefaultAttributes)+2)
537582
attrs = append(attrs, s.options.DefaultAttributes...)
538583
attrs = append(
@@ -574,9 +619,15 @@ func (s ocStmt) ExecContext(ctx context.Context, args []driver.NamedValue) (res
574619

575620
var span *trace.Span
576621
if parentSpan == nil {
577-
ctx, span = trace.StartSpan(ctx, "sql:exec", trace.WithSpanKind(trace.SpanKindClient))
622+
ctx, span = trace.StartSpan(ctx, "sql:exec",
623+
trace.WithSpanKind(trace.SpanKindClient),
624+
trace.WithSampler(s.options.Sampler),
625+
)
578626
} else {
579-
_, span = trace.StartSpan(ctx, "sql:exec", trace.WithSpanKind(trace.SpanKindClient))
627+
_, span = trace.StartSpan(ctx, "sql:exec",
628+
trace.WithSpanKind(trace.SpanKindClient),
629+
trace.WithSampler(s.options.Sampler),
630+
)
580631
}
581632
attrs := append([]trace.Attribute(nil), s.options.DefaultAttributes...)
582633
if s.options.Query {
@@ -613,9 +664,15 @@ func (s ocStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (row
613664

614665
var span *trace.Span
615666
if parentSpan == nil {
616-
ctx, span = trace.StartSpan(ctx, "sql:query", trace.WithSpanKind(trace.SpanKindClient))
667+
ctx, span = trace.StartSpan(ctx, "sql:query",
668+
trace.WithSpanKind(trace.SpanKindClient),
669+
trace.WithSampler(s.options.Sampler),
670+
)
617671
} else {
618-
_, span = trace.StartSpan(ctx, "sql:query", trace.WithSpanKind(trace.SpanKindClient))
672+
_, span = trace.StartSpan(ctx, "sql:query",
673+
trace.WithSpanKind(trace.SpanKindClient),
674+
trace.WithSampler(s.options.Sampler),
675+
)
619676
}
620677
attrs := append([]trace.Attribute(nil), s.options.DefaultAttributes...)
621678
if s.options.Query {
@@ -730,7 +787,10 @@ func (r ocRows) Columns() []string {
730787

731788
func (r ocRows) Close() (err error) {
732789
if r.options.RowsClose {
733-
_, span := trace.StartSpan(r.ctx, "sql:rows_close", trace.WithSpanKind(trace.SpanKindClient))
790+
_, span := trace.StartSpan(r.ctx, "sql:rows_close",
791+
trace.WithSpanKind(trace.SpanKindClient),
792+
trace.WithSampler(r.options.Sampler),
793+
)
734794
if len(r.options.DefaultAttributes) > 0 {
735795
span.AddAttributes(r.options.DefaultAttributes...)
736796
}
@@ -746,7 +806,10 @@ func (r ocRows) Close() (err error) {
746806

747807
func (r ocRows) Next(dest []driver.Value) (err error) {
748808
if r.options.RowsNext {
749-
_, span := trace.StartSpan(r.ctx, "sql:rows_next", trace.WithSpanKind(trace.SpanKindClient))
809+
_, span := trace.StartSpan(r.ctx, "sql:rows_next",
810+
trace.WithSpanKind(trace.SpanKindClient),
811+
trace.WithSampler(r.options.Sampler),
812+
)
750813
if len(r.options.DefaultAttributes) > 0 {
751814
span.AddAttributes(r.options.DefaultAttributes...)
752815
}
@@ -792,7 +855,7 @@ func wrapRows(ctx context.Context, parent driver.Rows, options TraceOptions) dri
792855
return r
793856
}
794857

795-
// ocTx implemens driver.Tx
858+
// ocTx implements driver.Tx
796859
type ocTx struct {
797860
parent driver.Tx
798861
ctx context.Context
@@ -802,7 +865,10 @@ type ocTx struct {
802865
func (t ocTx) Commit() (err error) {
803866
defer recordCallStats(context.Background(), "go.sql.commit")(err)
804867

805-
_, span := trace.StartSpan(t.ctx, "sql:commit", trace.WithSpanKind(trace.SpanKindClient))
868+
_, span := trace.StartSpan(t.ctx, "sql:commit",
869+
trace.WithSpanKind(trace.SpanKindClient),
870+
trace.WithSampler(t.options.Sampler),
871+
)
806872
if len(t.options.DefaultAttributes) > 0 {
807873
span.AddAttributes(t.options.DefaultAttributes...)
808874
}
@@ -818,7 +884,10 @@ func (t ocTx) Commit() (err error) {
818884
func (t ocTx) Rollback() (err error) {
819885
defer recordCallStats(context.Background(), "go.sql.rollback")(err)
820886

821-
_, span := trace.StartSpan(t.ctx, "sql:rollback", trace.WithSpanKind(trace.SpanKindClient))
887+
_, span := trace.StartSpan(t.ctx, "sql:rollback",
888+
trace.WithSpanKind(trace.SpanKindClient),
889+
trace.WithSampler(t.options.Sampler),
890+
)
822891
if len(t.options.DefaultAttributes) > 0 {
823892
span.AddAttributes(t.options.DefaultAttributes...)
824893
}

‎options.go‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,9 @@ type TraceOptions struct {
5353

5454
// DisableErrSkip, if set to true, will suppress driver.ErrSkip errors in spans.
5555
DisableErrSkip bool
56+
57+
// Sampler to use when creating spans.
58+
Sampler trace.Sampler
5659
}
5760

5861
// WithAllTraceOptions enables all available trace options.
@@ -166,3 +169,10 @@ func WithDisableErrSkip(b bool) TraceOption {
166169
o.DisableErrSkip = b
167170
}
168171
}
172+
173+
// WithSampler will be used on span creation.
174+
func WithSampler(sampler trace.Sampler) TraceOption {
175+
return func(o *TraceOptions) {
176+
o.Sampler = sampler
177+
}
178+
}

0 commit comments

Comments
 (0)