Skip to content

Commit a6884a0

Browse files
basvanbeekbasvanbeek
authored
Merge pull request #4 from bgpat/bgpat/attr
Added WithDefaultAttributes()
2 parents 2347323 + c6ee1db commit a6884a0

File tree

2 files changed

+88
-28
lines changed

2 files changed

+88
-28
lines changed

‎driver.go‎

Lines changed: 74 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,9 @@ 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+
if len(c.options.DefaultAttributes) > 0 {
100+
span.AddAttributes(c.options.DefaultAttributes...)
101+
}
99102
defer func() {
100103
if err != nil {
101104
span.SetStatus(trace.Status{
@@ -122,12 +125,15 @@ func (c ocConn) Exec(query string, args []driver.Value) (res driver.Result, err
122125
}
123126

124127
ctx, span := trace.StartSpan(context.Background(), "sql:exec")
125-
attrs := []trace.Attribute{
128+
attrs := make([]trace.Attribute, 0, len(c.options.DefaultAttributes)+2)
129+
attrs = append(attrs, c.options.DefaultAttributes...)
130+
attrs = append(
131+
attrs,
126132
attrDeprecated,
127133
trace.StringAttribute(
128134
"ocsql.deprecated", "driver does not support ExecerContext",
129135
),
130-
}
136+
)
131137
if c.options.Query {
132138
attrs = append(attrs, trace.StringAttribute("sql.query", query))
133139
if c.options.QueryParams {
@@ -164,15 +170,14 @@ func (c ocConn) ExecContext(ctx context.Context, query string, args []driver.Nam
164170
} else {
165171
_, span = trace.StartSpan(ctx, "sql:exec")
166172
}
173+
attrs := append([]trace.Attribute(nil), c.options.DefaultAttributes...)
167174
if c.options.Query {
168-
attrs := []trace.Attribute{
169-
trace.StringAttribute("sql.query", query),
170-
}
175+
attrs = append(attrs, trace.StringAttribute("sql.query", query))
171176
if c.options.QueryParams {
172177
attrs = append(attrs, namedParamsAttr(args)...)
173178
}
174-
span.AddAttributes(attrs...)
175179
}
180+
span.AddAttributes(attrs...)
176181

177182
defer func() {
178183
setSpanStatus(span, err)
@@ -196,12 +201,15 @@ func (c ocConn) Query(query string, args []driver.Value) (rows driver.Rows, err
196201
}
197202

198203
ctx, span := trace.StartSpan(context.Background(), "sql:query")
199-
attrs := []trace.Attribute{
204+
attrs := make([]trace.Attribute, 0, len(c.options.DefaultAttributes)+2)
205+
attrs = append(attrs, c.options.DefaultAttributes...)
206+
attrs = append(
207+
attrs,
200208
attrDeprecated,
201209
trace.StringAttribute(
202210
"ocsql.deprecated", "driver does not support QueryerContext",
203211
),
204-
}
212+
)
205213
if c.options.Query {
206214
attrs = append(attrs, trace.StringAttribute("sql.query", query))
207215
if c.options.QueryParams {
@@ -239,15 +247,14 @@ func (c ocConn) QueryContext(ctx context.Context, query string, args []driver.Na
239247
} else {
240248
_, span = trace.StartSpan(ctx, "sql:query")
241249
}
250+
attrs := append([]trace.Attribute(nil), c.options.DefaultAttributes...)
242251
if c.options.Query {
243-
attrs := []trace.Attribute{
244-
trace.StringAttribute("sql.query", query),
245-
}
252+
attrs = append(attrs, trace.StringAttribute("sql.query", query))
246253
if c.options.QueryParams {
247254
attrs = append(attrs, namedParamsAttr(args)...)
248255
}
249-
span.AddAttributes(attrs...)
250256
}
257+
span.AddAttributes(attrs...)
251258

252259
defer func() {
253260
setSpanStatus(span, err)
@@ -268,7 +275,9 @@ func (c ocConn) QueryContext(ctx context.Context, query string, args []driver.Na
268275
func (c ocConn) Prepare(query string) (stmt driver.Stmt, err error) {
269276
if c.options.AllowRoot {
270277
_, span := trace.StartSpan(context.Background(), "sql:prepare")
271-
attrs := []trace.Attribute{attrMissingContext}
278+
attrs := make([]trace.Attribute, 0, len(c.options.DefaultAttributes)+1)
279+
attrs = append(attrs, c.options.DefaultAttributes...)
280+
attrs = append(attrs, attrMissingContext)
272281
if c.options.Query {
273282
attrs = append(attrs, trace.StringAttribute("sql.query", query))
274283
}
@@ -299,10 +308,11 @@ func (c *ocConn) Begin() (driver.Tx, error) {
299308

300309
func (c *ocConn) PrepareContext(ctx context.Context, query string) (stmt driver.Stmt, err error) {
301310
var span *trace.Span
311+
attrs := append([]trace.Attribute(nil), c.options.DefaultAttributes...)
302312
if c.options.AllowRoot || trace.FromContext(ctx) != nil {
303313
ctx, span = trace.StartSpan(ctx, "sql:prepare")
304314
if c.options.Query {
305-
span.AddAttributes(trace.StringAttribute("sql.query", query))
315+
attrs = append(attrs, trace.StringAttribute("sql.query", query))
306316
}
307317
defer func() {
308318
setSpanStatus(span, err)
@@ -314,10 +324,11 @@ func (c *ocConn) PrepareContext(ctx context.Context, query string) (stmt driver.
314324
stmt, err = prepCtx.PrepareContext(ctx, query)
315325
} else {
316326
if span != nil {
317-
span.AddAttributes(attrMissingContext)
327+
attrs = append(attrs, attrMissingContext)
318328
}
319329
stmt, err = c.parent.Prepare(query)
320330
}
331+
span.AddAttributes(attrs...)
321332
if err != nil {
322333
return nil, err
323334
}
@@ -335,10 +346,16 @@ func (c *ocConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx,
335346
}
336347

337348
var span *trace.Span
349+
attrs := append([]trace.Attribute(nil), c.options.DefaultAttributes...)
350+
defer func() {
351+
if len(attrs) > 0 {
352+
span.AddAttributes(attrs...)
353+
}
354+
}()
338355
if ctx == nil || ctx == context.TODO() {
339356
ctx = context.Background()
340357
_, span = trace.StartSpan(ctx, "sql:begin_transaction")
341-
span.AddAttributes(attrMissingContext)
358+
attrs = append(attrs, attrMissingContext)
342359
} else {
343360
_, span = trace.StartSpan(ctx, "sql:begin_transaction")
344361
}
@@ -353,10 +370,13 @@ func (c *ocConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx,
353370
return ocTx{parent: tx, ctx: ctx}, nil
354371
}
355372

356-
span.AddAttributes(attrDeprecated)
357-
span.AddAttributes(trace.StringAttribute(
358-
"ocsql.deprecated", "driver does not support ConnBeginTx",
359-
))
373+
attrs = append(
374+
attrs,
375+
attrDeprecated,
376+
trace.StringAttribute(
377+
"ocsql.deprecated", "driver does not support ConnBeginTx",
378+
),
379+
)
360380
tx, err := c.parent.Begin()
361381
setSpanStatus(span, err)
362382
if err != nil {
@@ -375,6 +395,9 @@ type ocResult struct {
375395
func (r ocResult) LastInsertId() (id int64, err error) {
376396
if r.options.LastInsertID {
377397
_, span := trace.StartSpan(r.ctx, "sql:last_insert_id")
398+
if len(r.options.DefaultAttributes) > 0 {
399+
span.AddAttributes(r.options.DefaultAttributes...)
400+
}
378401
defer func() {
379402
setSpanStatus(span, err)
380403
span.End()
@@ -388,6 +411,9 @@ func (r ocResult) LastInsertId() (id int64, err error) {
388411
func (r ocResult) RowsAffected() (cnt int64, err error) {
389412
if r.options.RowsAffected {
390413
_, span := trace.StartSpan(r.ctx, "sql:rows_affected")
414+
if len(r.options.DefaultAttributes) > 0 {
415+
span.AddAttributes(r.options.DefaultAttributes...)
416+
}
391417
defer func() {
392418
setSpanStatus(span, err)
393419
span.End()
@@ -411,12 +437,15 @@ func (s ocStmt) Exec(args []driver.Value) (res driver.Result, err error) {
411437
}
412438

413439
ctx, span := trace.StartSpan(context.Background(), "sql:exec")
414-
attrs := []trace.Attribute{
440+
attrs := make([]trace.Attribute, 0, len(s.options.DefaultAttributes)+2)
441+
attrs = append(attrs, s.options.DefaultAttributes...)
442+
attrs = append(
443+
attrs,
415444
attrDeprecated,
416445
trace.StringAttribute(
417446
"ocsql.deprecated", "driver does not support StmtExecContext",
418447
),
419-
}
448+
)
420449
if s.options.Query {
421450
attrs = append(attrs, trace.StringAttribute("sql.query", s.query))
422451
if s.options.QueryParams {
@@ -453,12 +482,15 @@ func (s ocStmt) Query(args []driver.Value) (rows driver.Rows, err error) {
453482
}
454483

455484
ctx, span := trace.StartSpan(context.Background(), "sql:query")
456-
attrs := []trace.Attribute{
485+
attrs := make([]trace.Attribute, 0, len(s.options.DefaultAttributes)+2)
486+
attrs = append(attrs, s.options.DefaultAttributes...)
487+
attrs = append(
488+
attrs,
457489
attrDeprecated,
458490
trace.StringAttribute(
459491
"ocsql.deprecated", "driver does not support StmtQueryContext",
460492
),
461-
}
493+
)
462494
if s.options.Query {
463495
attrs = append(attrs, trace.StringAttribute("sql.query", s.query))
464496
if s.options.QueryParams {
@@ -493,13 +525,14 @@ func (s ocStmt) ExecContext(ctx context.Context, args []driver.NamedValue) (res
493525
} else {
494526
_, span = trace.StartSpan(ctx, "sql:exec")
495527
}
528+
attrs := append([]trace.Attribute(nil), s.options.DefaultAttributes...)
496529
if s.options.Query {
497-
attrs := []trace.Attribute{trace.StringAttribute("sql.query", s.query)}
530+
attrs = append(attrs, trace.StringAttribute("sql.query", s.query))
498531
if s.options.QueryParams {
499532
attrs = append(attrs, namedParamsAttr(args)...)
500533
}
501-
span.AddAttributes(attrs...)
502534
}
535+
span.AddAttributes(attrs...)
503536

504537
defer func() {
505538
setSpanStatus(span, err)
@@ -529,13 +562,14 @@ func (s ocStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (row
529562
} else {
530563
_, span = trace.StartSpan(ctx, "sql:query")
531564
}
565+
attrs := append([]trace.Attribute(nil), s.options.DefaultAttributes...)
532566
if s.options.Query {
533-
attrs := []trace.Attribute{trace.StringAttribute("sql.query", s.query)}
567+
attrs = append(attrs, trace.StringAttribute("sql.query", s.query))
534568
if s.options.QueryParams {
535569
attrs = append(attrs, namedParamsAttr(args)...)
536570
}
537-
span.AddAttributes(attrs...)
538571
}
572+
span.AddAttributes(attrs...)
539573

540574
defer func() {
541575
setSpanStatus(span, err)
@@ -566,6 +600,9 @@ func (r ocRows) Columns() []string {
566600
func (r ocRows) Close() (err error) {
567601
if r.options.RowsClose {
568602
_, span := trace.StartSpan(r.ctx, "sql:rows_close")
603+
if len(r.options.DefaultAttributes) > 0 {
604+
span.AddAttributes(r.options.DefaultAttributes...)
605+
}
569606
defer func() {
570607
setSpanStatus(span, err)
571608
span.End()
@@ -579,6 +616,9 @@ func (r ocRows) Close() (err error) {
579616
func (r ocRows) Next(dest []driver.Value) (err error) {
580617
if r.options.RowsNext {
581618
_, span := trace.StartSpan(r.ctx, "sql:rows_next")
619+
if len(r.options.DefaultAttributes) > 0 {
620+
span.AddAttributes(r.options.DefaultAttributes...)
621+
}
582622
defer func() {
583623
if err == io.EOF {
584624
// not an error; expected to happen during iteration
@@ -603,6 +643,9 @@ type ocTx struct {
603643

604644
func (t ocTx) Commit() (err error) {
605645
_, span := trace.StartSpan(t.ctx, "sql:commit")
646+
if len(t.options.DefaultAttributes) > 0 {
647+
span.AddAttributes(t.options.DefaultAttributes...)
648+
}
606649
defer func() {
607650
setSpanStatus(span, err)
608651
span.End()
@@ -614,6 +657,9 @@ func (t ocTx) Commit() (err error) {
614657

615658
func (t ocTx) Rollback() (err error) {
616659
_, span := trace.StartSpan(t.ctx, "sql:rollback")
660+
if len(t.options.DefaultAttributes) > 0 {
661+
span.AddAttributes(t.options.DefaultAttributes...)
662+
}
617663
defer func() {
618664
setSpanStatus(span, err)
619665
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 = attrs
154+
}
155+
}

0 commit comments

Comments
 (0)