Skip to content

Commit 016fb0f

Browse files
basvanbeekbasvanbeek
committed
AllowRoot functionality to control if ocsql can create root spans or not
1 parent 2a482e4 commit 016fb0f

File tree

2 files changed

+123
-27
lines changed

2 files changed

+123
-27
lines changed

‎driver.go‎

Lines changed: 109 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -113,7 +113,7 @@ func (d ocDriver) Open(name string) (driver.Conn, error) {
113113
}
114114

115115
func (c ocConn) Ping(ctx context.Context) (err error) {
116-
if c.options.Ping {
116+
if c.options.Ping && (c.options.AllowRoot || trace.FromContext(ctx) != nil) {
117117
var span *trace.Span
118118
ctx, span = trace.StartSpan(ctx, "sql:ping")
119119
defer func() {
@@ -137,6 +137,10 @@ func (c ocConn) Ping(ctx context.Context) (err error) {
137137

138138
func (c ocConn) Exec(query string, args []driver.Value) (res driver.Result, err error) {
139139
if exec, ok := c.parent.(driver.Execer); ok {
140+
if !c.options.AllowRoot {
141+
return exec.Exec(query, args)
142+
}
143+
140144
ctx, span := trace.StartSpan(context.Background(), "sql:exec")
141145
attrs := []trace.Attribute{
142146
attrDeprecated,
@@ -169,8 +173,17 @@ func (c ocConn) Exec(query string, args []driver.Value) (res driver.Result, err
169173

170174
func (c ocConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (res driver.Result, err error) {
171175
if execCtx, ok := c.parent.(driver.ExecerContext); ok {
176+
parentSpan := trace.FromContext(ctx)
177+
if !c.options.AllowRoot && parentSpan == nil {
178+
return execCtx.ExecContext(ctx, query, args)
179+
}
180+
172181
var span *trace.Span
173-
ctx, span = trace.StartSpan(ctx, "sql:exec")
182+
if parentSpan == nil {
183+
ctx, span = trace.StartSpan(ctx, "sql:exec")
184+
} else {
185+
_, span = trace.StartSpan(ctx, "sql:exec")
186+
}
174187
if c.options.Query {
175188
attrs := []trace.Attribute{
176189
trace.StringAttribute("sql.query", query),
@@ -198,6 +211,10 @@ func (c ocConn) ExecContext(ctx context.Context, query string, args []driver.Nam
198211

199212
func (c ocConn) Query(query string, args []driver.Value) (rows driver.Rows, err error) {
200213
if queryer, ok := c.parent.(driver.Queryer); ok {
214+
if !c.options.AllowRoot {
215+
return queryer.Query(query, args)
216+
}
217+
201218
ctx, span := trace.StartSpan(context.Background(), "sql:query")
202219
attrs := []trace.Attribute{
203220
attrDeprecated,
@@ -231,8 +248,17 @@ func (c ocConn) Query(query string, args []driver.Value) (rows driver.Rows, err
231248

232249
func (c ocConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (rows driver.Rows, err error) {
233250
if queryerCtx, ok := c.parent.(driver.QueryerContext); ok {
251+
parentSpan := trace.FromContext(ctx)
252+
if !c.options.AllowRoot && parentSpan == nil {
253+
return queryerCtx.QueryContext(ctx, query, args)
254+
}
255+
234256
var span *trace.Span
235-
ctx, span = trace.StartSpan(ctx, "sql:query")
257+
if parentSpan == nil {
258+
ctx, span = trace.StartSpan(ctx, "sql:query")
259+
} else {
260+
_, span = trace.StartSpan(ctx, "sql:query")
261+
}
236262
if c.options.Query {
237263
attrs := []trace.Attribute{
238264
trace.StringAttribute("sql.query", query),
@@ -260,17 +286,19 @@ func (c ocConn) QueryContext(ctx context.Context, query string, args []driver.Na
260286
}
261287

262288
func (c ocConn) Prepare(query string) (stmt driver.Stmt, err error) {
263-
_, span := trace.StartSpan(context.Background(), "sql:prepare")
264-
attrs := []trace.Attribute{attrMissingContext}
265-
if c.options.Query {
266-
attrs = append(attrs, trace.StringAttribute("sql.query", query))
267-
}
268-
span.AddAttributes(attrs...)
289+
if c.options.AllowRoot {
290+
_, span := trace.StartSpan(context.Background(), "sql:prepare")
291+
attrs := []trace.Attribute{attrMissingContext}
292+
if c.options.Query {
293+
attrs = append(attrs, trace.StringAttribute("sql.query", query))
294+
}
295+
span.AddAttributes(attrs...)
269296

270-
defer func() {
271-
setSpanStatus(span, err)
272-
span.End()
273-
}()
297+
defer func() {
298+
setSpanStatus(span, err)
299+
span.End()
300+
}()
301+
}
274302

275303
stmt, err = c.parent.Prepare(query)
276304
if err != nil {
@@ -291,19 +319,23 @@ func (c *ocConn) Begin() (driver.Tx, error) {
291319

292320
func (c *ocConn) PrepareContext(ctx context.Context, query string) (stmt driver.Stmt, err error) {
293321
var span *trace.Span
294-
ctx, span = trace.StartSpan(ctx, "sql:prepare")
295-
if c.options.Query {
296-
span.AddAttributes(trace.StringAttribute("sql.query", query))
322+
if c.options.AllowRoot || trace.FromContext(ctx) != nil {
323+
ctx, span = trace.StartSpan(ctx, "sql:prepare")
324+
if c.options.Query {
325+
span.AddAttributes(trace.StringAttribute("sql.query", query))
326+
}
327+
defer func() {
328+
setSpanStatus(span, err)
329+
span.End()
330+
}()
297331
}
298-
defer func() {
299-
setSpanStatus(span, err)
300-
span.End()
301-
}()
302332

303333
if prepCtx, ok := c.parent.(driver.ConnPrepareContext); ok {
304334
stmt, err = prepCtx.PrepareContext(ctx, query)
305335
} else {
306-
span.AddAttributes(attrMissingContext)
336+
if span != nil {
337+
span.AddAttributes(attrMissingContext)
338+
}
307339
stmt, err = c.parent.Prepare(query)
308340
}
309341
if err != nil {
@@ -315,6 +347,13 @@ func (c *ocConn) PrepareContext(ctx context.Context, query string) (stmt driver.
315347
}
316348

317349
func (c *ocConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
350+
if !c.options.AllowRoot && trace.FromContext(ctx) == nil {
351+
if connBeginTx, ok := c.parent.(driver.ConnBeginTx); ok {
352+
return connBeginTx.BeginTx(ctx, opts)
353+
}
354+
return c.parent.Begin()
355+
}
356+
318357
if c.options.Transaction {
319358
if ctx == nil || ctx == context.TODO() {
320359
var appSpan *trace.Span
@@ -431,7 +470,11 @@ func (r ocResult) RowsAffected() (cnt int64, err error) {
431470
}
432471

433472
func (s ocStmt) Exec(args []driver.Value) (res driver.Result, err error) {
434-
_, span := trace.StartSpan(context.Background(), "sql:exec")
473+
if !s.options.AllowRoot {
474+
return s.parent.Exec(args)
475+
}
476+
477+
ctx, span := trace.StartSpan(context.Background(), "sql:exec")
435478
attrs := []trace.Attribute{
436479
attrDeprecated,
437480
trace.StringAttribute(
@@ -452,6 +495,11 @@ func (s ocStmt) Exec(args []driver.Value) (res driver.Result, err error) {
452495
}()
453496

454497
res, err = s.parent.Exec(args)
498+
if err != nil {
499+
return nil, err
500+
}
501+
502+
res, err = ocResult{parent: res, ctx: ctx, options: s.options}, nil
455503
return
456504
}
457505

@@ -464,7 +512,11 @@ func (s ocStmt) NumInput() int {
464512
}
465513

466514
func (s ocStmt) Query(args []driver.Value) (rows driver.Rows, err error) {
467-
_, span := trace.StartSpan(context.Background(), "sql:query")
515+
if !s.options.AllowRoot {
516+
return s.parent.Query(args)
517+
}
518+
519+
ctx, span := trace.StartSpan(context.Background(), "sql:query")
468520
attrs := []trace.Attribute{
469521
attrDeprecated,
470522
trace.StringAttribute(
@@ -485,12 +537,26 @@ func (s ocStmt) Query(args []driver.Value) (rows driver.Rows, err error) {
485537
}()
486538

487539
rows, err = s.parent.Query(args)
540+
if err != nil {
541+
return nil, err
542+
}
543+
rows, err = ocRows{parent: rows, ctx: ctx, options: s.options}, nil
488544
return
489545
}
490546

491547
func (s ocStmt) ExecContext(ctx context.Context, args []driver.NamedValue) (res driver.Result, err error) {
548+
parentSpan := trace.FromContext(ctx)
549+
if !s.options.AllowRoot && parentSpan == nil {
550+
// we already tested driver to implement StmtExecContext
551+
return s.parent.(driver.StmtExecContext).ExecContext(ctx, args)
552+
}
553+
492554
var span *trace.Span
493-
ctx, span = trace.StartSpan(ctx, "sql:exec")
555+
if parentSpan == nil {
556+
ctx, span = trace.StartSpan(ctx, "sql:exec")
557+
} else {
558+
_, span = trace.StartSpan(ctx, "sql:exec")
559+
}
494560
if s.options.Query {
495561
attrs := []trace.Attribute{trace.StringAttribute("sql.query", s.query)}
496562
if s.options.QueryParams {
@@ -507,12 +573,26 @@ func (s ocStmt) ExecContext(ctx context.Context, args []driver.NamedValue) (res
507573
// we already tested driver to implement StmtExecContext
508574
execContext := s.parent.(driver.StmtExecContext)
509575
res, err = execContext.ExecContext(ctx, args)
576+
if err != nil {
577+
return nil, err
578+
}
579+
res, err = ocResult{parent: res, ctx: ctx, options: s.options}, nil
510580
return
511581
}
512582

513583
func (s ocStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (rows driver.Rows, err error) {
584+
parentSpan := trace.FromContext(ctx)
585+
if !s.options.AllowRoot && parentSpan == nil {
586+
// we already tested driver to implement StmtQueryContext
587+
return s.parent.(driver.StmtQueryContext).QueryContext(ctx, args)
588+
}
589+
514590
var span *trace.Span
515-
ctx, span = trace.StartSpan(ctx, "sql:query")
591+
if parentSpan == nil {
592+
ctx, span = trace.StartSpan(ctx, "sql:query")
593+
} else {
594+
_, span = trace.StartSpan(ctx, "sql:query")
595+
}
516596
if s.options.Query {
517597
attrs := []trace.Attribute{trace.StringAttribute("sql.query", s.query)}
518598
if s.options.QueryParams {
@@ -529,6 +609,10 @@ func (s ocStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (row
529609
// we already tested driver to implement StmtQueryContext
530610
queryContext := s.parent.(driver.StmtQueryContext)
531611
rows, err = queryContext.QueryContext(ctx, args)
612+
if err != nil {
613+
return nil, err
614+
}
615+
rows, err = ocRows{parent: rows, ctx: ctx, options: s.options}, nil
532616
return
533617
}
534618

‎options.go‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ type TraceOption func(o *TraceOptions)
55

66
// TraceOptions holds configuration of our ocsql tracing middleware.
77
type TraceOptions struct {
8+
AllowRoot bool
89
Transaction bool
910
Ping bool
1011
RowsNext bool
@@ -17,6 +18,7 @@ type TraceOptions struct {
1718

1819
// TraceAll has all tracing options enabled.
1920
var TraceAll = TraceOptions{
21+
AllowRoot: true,
2022
Transaction: true,
2123
Ping: true,
2224
RowsNext: true,
@@ -27,14 +29,24 @@ var TraceAll = TraceOptions{
2729
QueryParams: true,
2830
}
2931

30-
// WithOptions sets our ocsql tracing middleware options.
32+
// WithOptions sets our ocsql tracing middleware options through a single
33+
// TraceOptions object.
3134
func WithOptions(options TraceOptions) TraceOption {
3235
return func(o *TraceOptions) {
3336
*o = options
3437
}
3538
}
3639

37-
// WithTransaction enables / disables tracing of transaction spans
40+
// WithAllowRoot when set to true will allow ocsql to create root spans. If
41+
// no context is provided to (the majority) of database/sql commands this will
42+
// result in many single span traces.
43+
func WithAllowRoot(b bool) TraceOption {
44+
return func(o *TraceOptions) {
45+
o.AllowRoot = b
46+
}
47+
}
48+
49+
// WithTransaction enables / disables creation of transaction spans.
3850
func WithTransaction(b bool) TraceOption {
3951
return func(o *TraceOptions) {
4052
o.Transaction = b

0 commit comments

Comments
 (0)