Skip to content

Commit 68ca6f8

Browse files
vearutopvearutop
committed
Merge remote-tracking branch 'upstream/master' into db-instance-name
2 parents c352499 + e766b83 commit 68ca6f8

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
@@ -133,7 +133,10 @@ func (c ocConn) Ping(ctx context.Context) (err error) {
133133

134134
if c.options.Ping && (c.options.AllowRoot || trace.FromContext(ctx) != nil) {
135135
var span *trace.Span
136-
ctx, span = trace.StartSpan(ctx, "sql:ping", trace.WithSpanKind(trace.SpanKindClient))
136+
ctx, span = trace.StartSpan(ctx, "sql:ping",
137+
trace.WithSpanKind(trace.SpanKindClient),
138+
trace.WithSampler(c.options.Sampler),
139+
)
137140
if len(c.options.DefaultAttributes) > 0 {
138141
span.AddAttributes(c.options.DefaultAttributes...)
139142
}
@@ -164,7 +167,10 @@ func (c ocConn) Exec(query string, args []driver.Value) (res driver.Result, err
164167
return exec.Exec(query, args)
165168
}
166169

167-
ctx, span := trace.StartSpan(context.Background(), "sql:exec", trace.WithSpanKind(trace.SpanKindClient))
170+
ctx, span := trace.StartSpan(context.Background(), "sql:exec",
171+
trace.WithSpanKind(trace.SpanKindClient),
172+
trace.WithSampler(c.options.Sampler),
173+
)
168174
attrs := make([]trace.Attribute, 0, len(c.options.DefaultAttributes)+2)
169175
attrs = append(attrs, c.options.DefaultAttributes...)
170176
attrs = append(
@@ -208,9 +214,15 @@ func (c ocConn) ExecContext(ctx context.Context, query string, args []driver.Nam
208214

209215
var span *trace.Span
210216
if parentSpan == nil {
211-
ctx, span = trace.StartSpan(ctx, "sql:exec", trace.WithSpanKind(trace.SpanKindClient))
217+
ctx, span = trace.StartSpan(ctx, "sql:exec",
218+
trace.WithSpanKind(trace.SpanKindClient),
219+
trace.WithSampler(c.options.Sampler),
220+
)
212221
} else {
213-
_, span = trace.StartSpan(ctx, "sql:exec", trace.WithSpanKind(trace.SpanKindClient))
222+
_, span = trace.StartSpan(ctx, "sql:exec",
223+
trace.WithSpanKind(trace.SpanKindClient),
224+
trace.WithSampler(c.options.Sampler),
225+
)
214226
}
215227
attrs := append([]trace.Attribute(nil), c.options.DefaultAttributes...)
216228
if c.options.Query {
@@ -244,7 +256,10 @@ func (c ocConn) Query(query string, args []driver.Value) (rows driver.Rows, err
244256
return queryer.Query(query, args)
245257
}
246258

247-
ctx, span := trace.StartSpan(context.Background(), "sql:query", trace.WithSpanKind(trace.SpanKindClient))
259+
ctx, span := trace.StartSpan(context.Background(), "sql:query",
260+
trace.WithSpanKind(trace.SpanKindClient),
261+
trace.WithSampler(c.options.Sampler),
262+
)
248263
attrs := make([]trace.Attribute, 0, len(c.options.DefaultAttributes)+2)
249264
attrs = append(attrs, c.options.DefaultAttributes...)
250265
attrs = append(
@@ -289,9 +304,15 @@ func (c ocConn) QueryContext(ctx context.Context, query string, args []driver.Na
289304

290305
var span *trace.Span
291306
if parentSpan == nil {
292-
ctx, span = trace.StartSpan(ctx, "sql:query", trace.WithSpanKind(trace.SpanKindClient))
307+
ctx, span = trace.StartSpan(ctx, "sql:query",
308+
trace.WithSpanKind(trace.SpanKindClient),
309+
trace.WithSampler(c.options.Sampler),
310+
)
293311
} else {
294-
_, span = trace.StartSpan(ctx, "sql:query", trace.WithSpanKind(trace.SpanKindClient))
312+
_, span = trace.StartSpan(ctx, "sql:query",
313+
trace.WithSpanKind(trace.SpanKindClient),
314+
trace.WithSampler(c.options.Sampler),
315+
)
295316
}
296317
attrs := append([]trace.Attribute(nil), c.options.DefaultAttributes...)
297318
if c.options.Query {
@@ -322,7 +343,10 @@ func (c ocConn) Prepare(query string) (stmt driver.Stmt, err error) {
322343
defer recordCallStats(context.Background(), "go.sql.prepare", c.options.InstanceName)(err)
323344

324345
if c.options.AllowRoot {
325-
_, span := trace.StartSpan(context.Background(), "sql:prepare", trace.WithSpanKind(trace.SpanKindClient))
346+
_, span := trace.StartSpan(context.Background(), "sql:prepare",
347+
trace.WithSpanKind(trace.SpanKindClient),
348+
trace.WithSampler(c.options.Sampler),
349+
)
326350
attrs := make([]trace.Attribute, 0, len(c.options.DefaultAttributes)+1)
327351
attrs = append(attrs, c.options.DefaultAttributes...)
328352
attrs = append(attrs, attrMissingContext)
@@ -360,7 +384,10 @@ func (c *ocConn) PrepareContext(ctx context.Context, query string) (stmt driver.
360384
var span *trace.Span
361385
attrs := append([]trace.Attribute(nil), c.options.DefaultAttributes...)
362386
if c.options.AllowRoot || trace.FromContext(ctx) != nil {
363-
ctx, span = trace.StartSpan(ctx, "sql:prepare", trace.WithSpanKind(trace.SpanKindClient))
387+
ctx, span = trace.StartSpan(ctx, "sql:prepare",
388+
trace.WithSpanKind(trace.SpanKindClient),
389+
trace.WithSampler(c.options.Sampler),
390+
)
364391
if c.options.Query {
365392
attrs = append(attrs, trace.StringAttribute("sql.query", query))
366393
}
@@ -402,10 +429,16 @@ func (c *ocConn) BeginTx(ctx context.Context, opts driver.TxOptions) (tx driver.
402429

403430
if ctx == nil || ctx == context.TODO() {
404431
ctx = context.Background()
405-
_, 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+
)
406436
attrs = append(attrs, attrMissingContext)
407437
} else {
408-
_, span = trace.StartSpan(ctx, "sql:begin_transaction", trace.WithSpanKind(trace.SpanKindClient))
438+
_, span = trace.StartSpan(ctx, "sql:begin_transaction",
439+
trace.WithSpanKind(trace.SpanKindClient),
440+
trace.WithSampler(c.options.Sampler),
441+
)
409442
}
410443
defer func() {
411444
if len(attrs) > 0 {
@@ -447,7 +480,10 @@ type ocResult struct {
447480

448481
func (r ocResult) LastInsertId() (id int64, err error) {
449482
if r.options.LastInsertID {
450-
_, span := trace.StartSpan(r.ctx, "sql:last_insert_id", trace.WithSpanKind(trace.SpanKindClient))
483+
_, span := trace.StartSpan(r.ctx, "sql:last_insert_id",
484+
trace.WithSpanKind(trace.SpanKindClient),
485+
trace.WithSampler(r.options.Sampler),
486+
)
451487
if len(r.options.DefaultAttributes) > 0 {
452488
span.AddAttributes(r.options.DefaultAttributes...)
453489
}
@@ -463,7 +499,10 @@ func (r ocResult) LastInsertId() (id int64, err error) {
463499

464500
func (r ocResult) RowsAffected() (cnt int64, err error) {
465501
if r.options.RowsAffected {
466-
_, span := trace.StartSpan(r.ctx, "sql:rows_affected", trace.WithSpanKind(trace.SpanKindClient))
502+
_, span := trace.StartSpan(r.ctx, "sql:rows_affected",
503+
trace.WithSpanKind(trace.SpanKindClient),
504+
trace.WithSampler(r.options.Sampler),
505+
)
467506
if len(r.options.DefaultAttributes) > 0 {
468507
span.AddAttributes(r.options.DefaultAttributes...)
469508
}
@@ -491,7 +530,10 @@ func (s ocStmt) Exec(args []driver.Value) (res driver.Result, err error) {
491530
return s.parent.Exec(args)
492531
}
493532

494-
ctx, span := trace.StartSpan(context.Background(), "sql:exec", trace.WithSpanKind(trace.SpanKindClient))
533+
ctx, span := trace.StartSpan(context.Background(), "sql:exec",
534+
trace.WithSpanKind(trace.SpanKindClient),
535+
trace.WithSampler(s.options.Sampler),
536+
)
495537
attrs := make([]trace.Attribute, 0, len(s.options.DefaultAttributes)+2)
496538
attrs = append(attrs, s.options.DefaultAttributes...)
497539
attrs = append(
@@ -538,7 +580,10 @@ func (s ocStmt) Query(args []driver.Value) (rows driver.Rows, err error) {
538580
return s.parent.Query(args)
539581
}
540582

541-
ctx, span := trace.StartSpan(context.Background(), "sql:query", trace.WithSpanKind(trace.SpanKindClient))
583+
ctx, span := trace.StartSpan(context.Background(), "sql:query",
584+
trace.WithSpanKind(trace.SpanKindClient),
585+
trace.WithSampler(s.options.Sampler),
586+
)
542587
attrs := make([]trace.Attribute, 0, len(s.options.DefaultAttributes)+2)
543588
attrs = append(attrs, s.options.DefaultAttributes...)
544589
attrs = append(
@@ -580,9 +625,15 @@ func (s ocStmt) ExecContext(ctx context.Context, args []driver.NamedValue) (res
580625

581626
var span *trace.Span
582627
if parentSpan == nil {
583-
ctx, span = trace.StartSpan(ctx, "sql:exec", trace.WithSpanKind(trace.SpanKindClient))
628+
ctx, span = trace.StartSpan(ctx, "sql:exec",
629+
trace.WithSpanKind(trace.SpanKindClient),
630+
trace.WithSampler(s.options.Sampler),
631+
)
584632
} else {
585-
_, span = trace.StartSpan(ctx, "sql:exec", trace.WithSpanKind(trace.SpanKindClient))
633+
_, span = trace.StartSpan(ctx, "sql:exec",
634+
trace.WithSpanKind(trace.SpanKindClient),
635+
trace.WithSampler(s.options.Sampler),
636+
)
586637
}
587638
attrs := append([]trace.Attribute(nil), s.options.DefaultAttributes...)
588639
if s.options.Query {
@@ -619,9 +670,15 @@ func (s ocStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (row
619670

620671
var span *trace.Span
621672
if parentSpan == nil {
622-
ctx, span = trace.StartSpan(ctx, "sql:query", trace.WithSpanKind(trace.SpanKindClient))
673+
ctx, span = trace.StartSpan(ctx, "sql:query",
674+
trace.WithSpanKind(trace.SpanKindClient),
675+
trace.WithSampler(s.options.Sampler),
676+
)
623677
} else {
624-
_, span = trace.StartSpan(ctx, "sql:query", trace.WithSpanKind(trace.SpanKindClient))
678+
_, span = trace.StartSpan(ctx, "sql:query",
679+
trace.WithSpanKind(trace.SpanKindClient),
680+
trace.WithSampler(s.options.Sampler),
681+
)
625682
}
626683
attrs := append([]trace.Attribute(nil), s.options.DefaultAttributes...)
627684
if s.options.Query {
@@ -736,7 +793,10 @@ func (r ocRows) Columns() []string {
736793

737794
func (r ocRows) Close() (err error) {
738795
if r.options.RowsClose {
739-
_, span := trace.StartSpan(r.ctx, "sql:rows_close", trace.WithSpanKind(trace.SpanKindClient))
796+
_, span := trace.StartSpan(r.ctx, "sql:rows_close",
797+
trace.WithSpanKind(trace.SpanKindClient),
798+
trace.WithSampler(r.options.Sampler),
799+
)
740800
if len(r.options.DefaultAttributes) > 0 {
741801
span.AddAttributes(r.options.DefaultAttributes...)
742802
}
@@ -752,7 +812,10 @@ func (r ocRows) Close() (err error) {
752812

753813
func (r ocRows) Next(dest []driver.Value) (err error) {
754814
if r.options.RowsNext {
755-
_, span := trace.StartSpan(r.ctx, "sql:rows_next", trace.WithSpanKind(trace.SpanKindClient))
815+
_, span := trace.StartSpan(r.ctx, "sql:rows_next",
816+
trace.WithSpanKind(trace.SpanKindClient),
817+
trace.WithSampler(r.options.Sampler),
818+
)
756819
if len(r.options.DefaultAttributes) > 0 {
757820
span.AddAttributes(r.options.DefaultAttributes...)
758821
}
@@ -798,7 +861,7 @@ func wrapRows(ctx context.Context, parent driver.Rows, options TraceOptions) dri
798861
return r
799862
}
800863

801-
// ocTx implemens driver.Tx
864+
// ocTx implements driver.Tx
802865
type ocTx struct {
803866
parent driver.Tx
804867
ctx context.Context
@@ -808,7 +871,10 @@ type ocTx struct {
808871
func (t ocTx) Commit() (err error) {
809872
defer recordCallStats(context.Background(), "go.sql.commit", t.options.InstanceName)(err)
810873

811-
_, span := trace.StartSpan(t.ctx, "sql:commit", trace.WithSpanKind(trace.SpanKindClient))
874+
_, span := trace.StartSpan(t.ctx, "sql:commit",
875+
trace.WithSpanKind(trace.SpanKindClient),
876+
trace.WithSampler(t.options.Sampler),
877+
)
812878
if len(t.options.DefaultAttributes) > 0 {
813879
span.AddAttributes(t.options.DefaultAttributes...)
814880
}
@@ -824,7 +890,10 @@ func (t ocTx) Commit() (err error) {
824890
func (t ocTx) Rollback() (err error) {
825891
defer recordCallStats(context.Background(), "go.sql.rollback", t.options.InstanceName)(err)
826892

827-
_, span := trace.StartSpan(t.ctx, "sql:rollback", trace.WithSpanKind(trace.SpanKindClient))
893+
_, span := trace.StartSpan(t.ctx, "sql:rollback",
894+
trace.WithSpanKind(trace.SpanKindClient),
895+
trace.WithSampler(t.options.Sampler),
896+
)
828897
if len(t.options.DefaultAttributes) > 0 {
829898
span.AddAttributes(t.options.DefaultAttributes...)
830899
}

‎options.go‎

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

5959
// DisableErrSkip, if set to true, will suppress driver.ErrSkip errors in spans.
6060
DisableErrSkip bool
61+
62+
// Sampler to use when creating spans.
63+
Sampler trace.Sampler
6164
}
6265

6366
// WithAllTraceOptions enables all available trace options.
@@ -172,6 +175,13 @@ func WithDisableErrSkip(b bool) TraceOption {
172175
}
173176
}
174177

178+
// WithSampler will be used on span creation.
179+
func WithSampler(sampler trace.Sampler) TraceOption {
180+
return func(o *TraceOptions) {
181+
o.Sampler = sampler
182+
}
183+
}
184+
175185
// WithInstanceName sets database instance name.
176186
func WithInstanceName(instanceName string) TraceOption {
177187
return func(o *TraceOptions) {

0 commit comments

Comments
 (0)