Skip to content

Commit 1b1c082

Browse files
bgpatbgpat
committed
Add option for default attributes
1 parent d946a62 commit 1b1c082

File tree

2 files changed

+60
-28
lines changed

2 files changed

+60
-28
lines changed

‎driver.go‎

Lines changed: 46 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,7 @@ func (c ocConn) Ping(ctx context.Context) (err error) {
9696
if c.options.Ping && (c.options.AllowRoot || trace.FromContext(ctx) != nil) {
9797
var span *trace.Span
9898
ctx, span = trace.StartSpan(ctx, "sql:ping")
99+
span.AddAttributes(c.options.DefaultAttributes...)
99100
defer func() {
100101
if err != nil {
101102
span.SetStatus(trace.Status{
@@ -122,12 +123,13 @@ func (c ocConn) Exec(query string, args []driver.Value) (res driver.Result, err
122123
}
123124

124125
ctx, span := trace.StartSpan(context.Background(), "sql:exec")
125-
attrs := []trace.Attribute{
126+
attrs := append(
127+
c.options.DefaultAttributes,
126128
attrDeprecated,
127129
trace.StringAttribute(
128130
"ocsql.deprecated", "driver does not support ExecerContext",
129131
),
130-
}
132+
)
131133
if c.options.Query {
132134
attrs = append(attrs, trace.StringAttribute("sql.query", query))
133135
if c.options.QueryParams {
@@ -164,15 +166,14 @@ func (c ocConn) ExecContext(ctx context.Context, query string, args []driver.Nam
164166
} else {
165167
_, span = trace.StartSpan(ctx, "sql:exec")
166168
}
169+
attrs := c.options.DefaultAttributes
167170
if c.options.Query {
168-
attrs := []trace.Attribute{
169-
trace.StringAttribute("sql.query", query),
170-
}
171+
attrs = append(attrs, trace.StringAttribute("sql.query", query))
171172
if c.options.QueryParams {
172173
attrs = append(attrs, namedParamsAttr(args)...)
173174
}
174-
span.AddAttributes(attrs...)
175175
}
176+
span.AddAttributes(attrs...)
176177

177178
defer func() {
178179
setSpanStatus(span, err)
@@ -196,12 +197,13 @@ func (c ocConn) Query(query string, args []driver.Value) (rows driver.Rows, err
196197
}
197198

198199
ctx, span := trace.StartSpan(context.Background(), "sql:query")
199-
attrs := []trace.Attribute{
200+
attrs := append(
201+
c.options.DefaultAttributes,
200202
attrDeprecated,
201203
trace.StringAttribute(
202204
"ocsql.deprecated", "driver does not support QueryerContext",
203205
),
204-
}
206+
)
205207
if c.options.Query {
206208
attrs = append(attrs, trace.StringAttribute("sql.query", query))
207209
if c.options.QueryParams {
@@ -239,15 +241,14 @@ func (c ocConn) QueryContext(ctx context.Context, query string, args []driver.Na
239241
} else {
240242
_, span = trace.StartSpan(ctx, "sql:query")
241243
}
244+
attrs := c.options.DefaultAttributes
242245
if c.options.Query {
243-
attrs := []trace.Attribute{
244-
trace.StringAttribute("sql.query", query),
245-
}
246+
attrs = append(attrs, trace.StringAttribute("sql.query", query))
246247
if c.options.QueryParams {
247248
attrs = append(attrs, namedParamsAttr(args)...)
248249
}
249-
span.AddAttributes(attrs...)
250250
}
251+
span.AddAttributes(attrs...)
251252

252253
defer func() {
253254
setSpanStatus(span, err)
@@ -268,7 +269,7 @@ func (c ocConn) QueryContext(ctx context.Context, query string, args []driver.Na
268269
func (c ocConn) Prepare(query string) (stmt driver.Stmt, err error) {
269270
if c.options.AllowRoot {
270271
_, span := trace.StartSpan(context.Background(), "sql:prepare")
271-
attrs := []trace.Attribute{attrMissingContext}
272+
attrs := append(c.options.DefaultAttributes, attrMissingContext)
272273
if c.options.Query {
273274
attrs = append(attrs, trace.StringAttribute("sql.query", query))
274275
}
@@ -299,10 +300,11 @@ func (c *ocConn) Begin() (driver.Tx, error) {
299300

300301
func (c *ocConn) PrepareContext(ctx context.Context, query string) (stmt driver.Stmt, err error) {
301302
var span *trace.Span
303+
attrs := c.options.DefaultAttributes
302304
if c.options.AllowRoot || trace.FromContext(ctx) != nil {
303305
ctx, span = trace.StartSpan(ctx, "sql:prepare")
304306
if c.options.Query {
305-
span.AddAttributes(trace.StringAttribute("sql.query", query))
307+
attrs = append(attrs, trace.StringAttribute("sql.query", query))
306308
}
307309
defer func() {
308310
setSpanStatus(span, err)
@@ -314,10 +316,11 @@ func (c *ocConn) PrepareContext(ctx context.Context, query string) (stmt driver.
314316
stmt, err = prepCtx.PrepareContext(ctx, query)
315317
} else {
316318
if span != nil {
317-
span.AddAttributes(attrMissingContext)
319+
attrs = append(attrs, attrMissingContext)
318320
}
319321
stmt, err = c.parent.Prepare(query)
320322
}
323+
span.AddAttributes(attrs...)
321324
if err != nil {
322325
return nil, err
323326
}
@@ -335,10 +338,12 @@ func (c *ocConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx,
335338
}
336339

337340
var span *trace.Span
341+
attrs := c.options.DefaultAttributes
342+
defer span.AddAttributes(attrs...)
338343
if ctx == nil || ctx == context.TODO() {
339344
ctx = context.Background()
340345
_, span = trace.StartSpan(ctx, "sql:begin_transaction")
341-
span.AddAttributes(attrMissingContext)
346+
attrs = append(attrs, attrMissingContext)
342347
} else {
343348
_, span = trace.StartSpan(ctx, "sql:begin_transaction")
344349
}
@@ -353,10 +358,13 @@ func (c *ocConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx,
353358
return ocTx{parent: tx, ctx: ctx}, nil
354359
}
355360

356-
span.AddAttributes(attrDeprecated)
357-
span.AddAttributes(trace.StringAttribute(
358-
"ocsql.deprecated", "driver does not support ConnBeginTx",
359-
))
361+
attrs = append(
362+
attrs,
363+
attrDeprecated,
364+
trace.StringAttribute(
365+
"ocsql.deprecated", "driver does not support ConnBeginTx",
366+
),
367+
)
360368
tx, err := c.parent.Begin()
361369
setSpanStatus(span, err)
362370
if err != nil {
@@ -375,6 +383,7 @@ type ocResult struct {
375383
func (r ocResult) LastInsertId() (id int64, err error) {
376384
if r.options.LastInsertID {
377385
_, span := trace.StartSpan(r.ctx, "sql:last_insert_id")
386+
span.AddAttributes(r.options.DefaultAttributes...)
378387
defer func() {
379388
setSpanStatus(span, err)
380389
span.End()
@@ -388,6 +397,7 @@ func (r ocResult) LastInsertId() (id int64, err error) {
388397
func (r ocResult) RowsAffected() (cnt int64, err error) {
389398
if r.options.RowsAffected {
390399
_, span := trace.StartSpan(r.ctx, "sql:rows_affected")
400+
span.AddAttributes(r.options.DefaultAttributes...)
391401
defer func() {
392402
setSpanStatus(span, err)
393403
span.End()
@@ -411,12 +421,13 @@ func (s ocStmt) Exec(args []driver.Value) (res driver.Result, err error) {
411421
}
412422

413423
ctx, span := trace.StartSpan(context.Background(), "sql:exec")
414-
attrs := []trace.Attribute{
424+
attrs := append(
425+
s.options.DefaultAttributes,
415426
attrDeprecated,
416427
trace.StringAttribute(
417428
"ocsql.deprecated", "driver does not support StmtExecContext",
418429
),
419-
}
430+
)
420431
if s.options.Query {
421432
attrs = append(attrs, trace.StringAttribute("sql.query", s.query))
422433
if s.options.QueryParams {
@@ -453,12 +464,13 @@ func (s ocStmt) Query(args []driver.Value) (rows driver.Rows, err error) {
453464
}
454465

455466
ctx, span := trace.StartSpan(context.Background(), "sql:query")
456-
attrs := []trace.Attribute{
467+
attrs := append(
468+
s.options.DefaultAttributes,
457469
attrDeprecated,
458470
trace.StringAttribute(
459471
"ocsql.deprecated", "driver does not support StmtQueryContext",
460472
),
461-
}
473+
)
462474
if s.options.Query {
463475
attrs = append(attrs, trace.StringAttribute("sql.query", s.query))
464476
if s.options.QueryParams {
@@ -493,13 +505,14 @@ func (s ocStmt) ExecContext(ctx context.Context, args []driver.NamedValue) (res
493505
} else {
494506
_, span = trace.StartSpan(ctx, "sql:exec")
495507
}
508+
attrs := s.options.DefaultAttributes
496509
if s.options.Query {
497-
attrs := []trace.Attribute{trace.StringAttribute("sql.query", s.query)}
510+
attrs = append(attrs, trace.StringAttribute("sql.query", s.query))
498511
if s.options.QueryParams {
499512
attrs = append(attrs, namedParamsAttr(args)...)
500513
}
501-
span.AddAttributes(attrs...)
502514
}
515+
span.AddAttributes(attrs...)
503516

504517
defer func() {
505518
setSpanStatus(span, err)
@@ -529,13 +542,14 @@ func (s ocStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (row
529542
} else {
530543
_, span = trace.StartSpan(ctx, "sql:query")
531544
}
545+
attrs := s.options.DefaultAttributes
532546
if s.options.Query {
533-
attrs := []trace.Attribute{trace.StringAttribute("sql.query", s.query)}
547+
attrs = append(attrs, trace.StringAttribute("sql.query", s.query))
534548
if s.options.QueryParams {
535549
attrs = append(attrs, namedParamsAttr(args)...)
536550
}
537-
span.AddAttributes(attrs...)
538551
}
552+
span.AddAttributes(attrs...)
539553

540554
defer func() {
541555
setSpanStatus(span, err)
@@ -566,6 +580,7 @@ func (r ocRows) Columns() []string {
566580
func (r ocRows) Close() (err error) {
567581
if r.options.RowsClose {
568582
_, span := trace.StartSpan(r.ctx, "sql:rows_close")
583+
span.AddAttributes(r.options.DefaultAttributes...)
569584
defer func() {
570585
setSpanStatus(span, err)
571586
span.End()
@@ -579,6 +594,7 @@ func (r ocRows) Close() (err error) {
579594
func (r ocRows) Next(dest []driver.Value) (err error) {
580595
if r.options.RowsNext {
581596
_, span := trace.StartSpan(r.ctx, "sql:rows_next")
597+
span.AddAttributes(r.options.DefaultAttributes...)
582598
defer func() {
583599
if err == io.EOF {
584600
// not an error; expected to happen during iteration
@@ -603,6 +619,7 @@ type ocTx struct {
603619

604620
func (t ocTx) Commit() (err error) {
605621
_, span := trace.StartSpan(t.ctx, "sql:commit")
622+
span.AddAttributes(t.options.DefaultAttributes...)
606623
defer func() {
607624
setSpanStatus(span, err)
608625
span.End()
@@ -614,6 +631,7 @@ func (t ocTx) Commit() (err error) {
614631

615632
func (t ocTx) Rollback() (err error) {
616633
_, span := trace.StartSpan(t.ctx, "sql:rollback")
634+
span.AddAttributes(t.options.DefaultAttributes...)
617635
defer func() {
618636
setSpanStatus(span, err)
619637
span.End()

‎options.go‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
11
package ocsql
22

3+
import (
4+
"go.opencensus.io/trace"
5+
)
6+
37
// TraceOption allows for managing ocsql configuration using functional options.
48
type TraceOption func(o *TraceOptions)
59

@@ -43,6 +47,9 @@ type TraceOptions struct {
4347
// parameters recorded with respect to security.
4448
// This setting is a noop if the Query option is set to false.
4549
QueryParams bool
50+
51+
// DefaultAttributes will be set to each span as default.
52+
DefaultAttributes []trace.Attribute
4653
}
4754

4855
// WithAllTraceOptions enables all available trace options.
@@ -139,3 +146,10 @@ func WithQueryParams(b bool) TraceOption {
139146
o.QueryParams = b
140147
}
141148
}
149+
150+
// WithDefaultAttributes will be set to each span as default.
151+
func WithDefaultAttributes(attrs ...trace.Attribute) TraceOption {
152+
return func(o *TraceOptions) {
153+
o.DefaultAttributes = append(o.DefaultAttributes, attrs...)
154+
}
155+
}

0 commit comments

Comments
 (0)