Skip to content

Commit 0f7a0bb

Browse files
jcchavezsjcchavezs
committed
fix: fixes span error for ErrSkip errors from the driver.
1 parent 7900e02 commit 0f7a0bb

2 files changed

Lines changed: 136 additions & 48 deletions

File tree

‎driver.go‎

Lines changed: 76 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"database/sql/driver"
77
"fmt"
88
"sync"
9+
"time"
910

1011
zipkin "github.com/openzipkin/zipkin-go"
1112
zipkinmodel "github.com/openzipkin/zipkin-go/model"
@@ -110,21 +111,32 @@ func (c zConn) Exec(query string, args []driver.Value) (res driver.Result, err e
110111

111112
func (c zConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (res driver.Result, err error) {
112113
if execCtx, ok := c.parent.(driver.ExecerContext); ok {
113-
if zipkin.SpanFromContext(ctx) == nil && !c.options.AllowRootSpan {
114+
parentSpan := zipkin.SpanFromContext(ctx)
115+
if parentSpan == nil && !c.options.AllowRootSpan {
114116
return execCtx.ExecContext(ctx, query, args)
115117
}
116118

117-
span, _ := c.tracer.StartSpanFromContext(ctx, "sql/exec", zipkin.Kind(zipkinmodel.Client))
118-
defer span.Finish()
119-
120-
if c.options.TagQuery {
121-
span.Tag("sql.query", query)
122-
}
119+
var (
120+
startTime time.Time
121+
)
123122

124-
setSpanDefaultTags(span, c.options.DefaultTags)
123+
defer func() {
124+
if err == nil || err != driver.ErrSkip {
125+
var span zipkin.Span
126+
span, _ = c.tracer.StartSpanFromContext(ctx, "sql/exec", zipkin.Kind(zipkinmodel.Client), zipkin.StartTime(startTime))
127+
128+
if c.options.TagQuery {
129+
span.Tag("sql.query", query)
130+
}
131+
setSpanDefaultTags(span, c.options.DefaultTags)
132+
133+
setSpanError(span, err)
134+
span.Finish()
135+
}
136+
}()
125137

138+
startTime = time.Now()
126139
if res, err = execCtx.ExecContext(ctx, query, args); err != nil {
127-
zipkin.TagError.Set(span, err.Error())
128140
return nil, err
129141
}
130142

@@ -149,22 +161,27 @@ func (c zConn) QueryContext(ctx context.Context, query string, args []driver.Nam
149161
return queryerCtx.QueryContext(ctx, query, args)
150162
}
151163

152-
var span zipkin.Span
153-
if parentSpan == nil {
154-
span, ctx = c.tracer.StartSpanFromContext(ctx, "sql/query", zipkin.Kind(zipkinmodel.Client))
155-
} else {
156-
span, _ = c.tracer.StartSpanFromContext(ctx, "sql/query", zipkin.Kind(zipkinmodel.Client))
157-
}
158-
defer span.Finish()
164+
var (
165+
startTime time.Time
166+
)
159167

160-
if c.options.TagQuery {
161-
span.Tag("sql.query", query)
162-
}
163-
164-
setSpanDefaultTags(span, c.options.DefaultTags)
168+
defer func() {
169+
if err == nil || err != driver.ErrSkip {
170+
var span zipkin.Span
171+
span, _ = c.tracer.StartSpanFromContext(ctx, "sql/query", zipkin.Kind(zipkinmodel.Client), zipkin.StartTime(startTime))
172+
173+
if c.options.TagQuery {
174+
span.Tag("sql.query", query)
175+
}
176+
setSpanDefaultTags(span, c.options.DefaultTags)
177+
178+
setSpanError(span, err)
179+
span.Finish()
180+
}
181+
}()
165182

183+
startTime = time.Now()
166184
if rows, err = queryerCtx.QueryContext(ctx, query, args); err != nil {
167-
zipkin.TagError.Set(span, err.Error())
168185
return nil, err
169186
}
170187

@@ -175,6 +192,19 @@ func (c zConn) QueryContext(ctx context.Context, query string, args []driver.Nam
175192
}
176193

177194
func (c zConn) Prepare(query string) (stmt driver.Stmt, err error) {
195+
if c.options.AllowRootSpan {
196+
span := c.tracer.StartSpan("sql/prepare", zipkin.Kind(zipkinmodel.Client))
197+
198+
if c.options.TagQuery {
199+
span.Tag("sql.query", query)
200+
}
201+
setSpanDefaultTags(span, c.options.DefaultTags)
202+
defer func() {
203+
setSpanError(span, err)
204+
span.Finish()
205+
}()
206+
}
207+
178208
stmt, err = c.parent.Prepare(query)
179209
if err != nil {
180210
return nil, err
@@ -192,12 +222,33 @@ func (c *zConn) Begin() (driver.Tx, error) {
192222
return c.Begin()
193223
}
194224

195-
func (c *zConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
225+
func (c *zConn) PrepareContext(ctx context.Context, query string) (stmt driver.Stmt, err error) {
226+
var span zipkin.Span
227+
setSpanDefaultTags(span, c.options.DefaultTags)
228+
if c.options.AllowRootSpan || zipkin.SpanFromContext(ctx) != nil {
229+
span, ctx = c.tracer.StartSpanFromContext(ctx, "sql/prepare", zipkin.Kind(zipkinmodel.Client))
230+
if c.options.TagQuery {
231+
span.Tag("sql.query", query)
232+
}
233+
234+
defer func() {
235+
setSpanError(span, err)
236+
span.Finish()
237+
}()
238+
}
239+
196240
if prepCtx, ok := c.parent.(driver.ConnPrepareContext); ok {
197-
return prepCtx.PrepareContext(ctx, query)
241+
stmt, err = prepCtx.PrepareContext(ctx, query)
242+
} else {
243+
stmt, err = c.parent.Prepare(query)
198244
}
199245

200-
return c.parent.Prepare(query)
246+
if err != nil {
247+
return nil, err
248+
}
249+
250+
stmt = wrapStmt(stmt, query, c.tracer, c.options)
251+
return
201252
}
202253

203254
func (c *zConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {

‎driver_test.go‎

Lines changed: 60 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ package zipkinsql
33
import (
44
"context"
55
"database/sql"
6-
"fmt"
76
"testing"
87

98
_ "github.com/mattn/go-sqlite3"
@@ -41,7 +40,7 @@ func TestQuerySuccess(t *testing.T) {
4140
for _, c := range testCases {
4241
db, _, recorder := createDB(t, c.opts...)
4342

44-
rows, err := db.Query("SELECT 1")
43+
rows, err := db.Query("SELECT 1 WHERE 1 = ?", 1)
4544
if err != nil {
4645
t.Fatalf("unexpected error: %s", err.Error())
4746
}
@@ -66,6 +65,10 @@ func TestQuerySuccess(t *testing.T) {
6665
if want, have := "sql/query", spans[0].Name; want != have {
6766
t.Fatalf("unexpected span name, want: %s, have: %s", want, have)
6867
}
68+
69+
if errMsg, ok := spans[0].Tags["error"]; ok {
70+
t.Fatalf("unexpected error: %s", errMsg)
71+
}
6972
}
7073

7174
db.Close()
@@ -81,7 +84,7 @@ func TestQueryContextSuccess(t *testing.T) {
8184
for _, c := range testCases {
8285
db, _, recorder := createDB(t, c.opts...)
8386

84-
rows, err := db.QueryContext(ctx, "SELECT 1")
87+
rows, err := db.QueryContext(ctx, "SELECT 1 WHERE 1 = ?", 1)
8588
if err != nil {
8689
t.Fatalf("unexpected error: %s", err.Error())
8790
}
@@ -106,6 +109,10 @@ func TestQueryContextSuccess(t *testing.T) {
106109
if want, have := "sql/query", spans[0].Name; want != have {
107110
t.Fatalf("unexpected span name, want: %s, have: %s", want, have)
108111
}
112+
113+
if errMsg, ok := spans[0].Tags["error"]; ok {
114+
t.Fatalf("unexpected error: %s", errMsg)
115+
}
109116
}
110117

111118
db.Close()
@@ -119,7 +126,7 @@ func TestQueryContextPropagationSuccess(t *testing.T) {
119126

120127
span, ctx := tracer.StartSpanFromContext(ctx, "root")
121128

122-
rows, err := db.QueryContext(ctx, "SELECT 1")
129+
rows, err := db.QueryContext(ctx, "SELECT 1 WHERE 1 = ?", 1)
123130
if err != nil {
124131
t.Fatalf("unexpected error: %s", err.Error())
125132
}
@@ -146,6 +153,10 @@ func TestQueryContextPropagationSuccess(t *testing.T) {
146153
t.Fatalf("unexpected span name, want: %s, have: %s", want, have)
147154
}
148155

156+
if errMsg, ok := spans[0].Tags["error"]; ok {
157+
t.Fatalf("unexpected error: %s", errMsg)
158+
}
159+
149160
db.Close()
150161
recorder.Close()
151162
}
@@ -164,6 +175,7 @@ func TestExecContextSuccess(t *testing.T) {
164175
db, _, recorder := createDB(t, c.opts...)
165176

166177
sqlStmt := `
178+
drop table if exists foo;
167179
create table foo (id integer not null primary key, name text);
168180
delete from foo;
169181
`
@@ -192,6 +204,10 @@ func TestExecContextSuccess(t *testing.T) {
192204
if want, have := "sql/exec", spans[0].Name; want != have {
193205
t.Fatalf("unexpected span name, want: %s, have: %s", want, have)
194206
}
207+
208+
if errMsg, ok := spans[0].Tags["error"]; ok {
209+
t.Fatalf("unexpected error: %s", errMsg)
210+
}
195211
}
196212

197213
db.Close()
@@ -204,13 +220,14 @@ func TestTxWithCommitSuccess(t *testing.T) {
204220

205221
testCases := []testCase{
206222
{[]TraceOption{WithAllowRootSpan(false)}, 0},
207-
{[]TraceOption{WithAllowRootSpan(true)}, 3},
223+
{[]TraceOption{WithAllowRootSpan(true)}, 5},
208224
}
209225

210226
for _, c := range testCases {
211227
db, _, recorder := createDB(t, c.opts...)
212228

213229
sqlStmt := `
230+
drop table if exists foo;
214231
create table foo (id integer not null primary key, name text);
215232
delete from foo;
216233
`
@@ -229,11 +246,9 @@ func TestTxWithCommitSuccess(t *testing.T) {
229246
t.Fatalf("unexpected error: %s", err.Error())
230247
}
231248
defer stmt.Close()
232-
for i := 0; i < 100; i++ {
233-
_, err = stmt.Exec(i, fmt.Sprintf("こんにちわ世界%03d", i))
234-
if err != nil {
235-
t.Fatalf("unexpected error: %s", err.Error())
236-
}
249+
_, err = stmt.Exec("1", "こんにちわ世界")
250+
if err != nil {
251+
t.Fatalf("unexpected error: %s", err.Error())
237252
}
238253
tx.Commit()
239254

@@ -247,12 +262,25 @@ func TestTxWithCommitSuccess(t *testing.T) {
247262
t.Fatalf("unexpected first span name, want: %s, have: %s", want, have)
248263
}
249264
if want, have := "sql/begin_transaction", spans[1].Name; want != have {
250-
t.Fatalf("unexpected first span name, want: %s, have: %s", want, have)
251-
}
252-
if want, have := "sql/commit", spans[2].Name; want != have {
253265
t.Fatalf("unexpected second span name, want: %s, have: %s", want, have)
254266
}
267+
if want, have := "sql/prepare", spans[2].Name; want != have {
268+
t.Fatalf("unexpected third span name, want: %s, have: %s", want, have)
269+
}
270+
if want, have := "sql/exec", spans[3].Name; want != have {
271+
t.Fatalf("unexpected fourth span name, want: %s, have: %s", want, have)
272+
}
273+
if want, have := "sql/commit", spans[4].Name; want != have {
274+
t.Fatalf("unexpected fifth span name, want: %s, have: %s", want, have)
275+
}
276+
277+
for i := 0; i < 5; i++ {
278+
if errMsg, ok := spans[i].Tags["error"]; ok {
279+
t.Fatalf("unexpected error: %s", errMsg)
280+
}
281+
}
255282
}
283+
256284
db.Close()
257285
recorder.Close()
258286
}
@@ -263,13 +291,14 @@ func TestTxWithRollbackSuccess(t *testing.T) {
263291

264292
testCases := []testCase{
265293
{[]TraceOption{WithAllowRootSpan(false)}, 0},
266-
{[]TraceOption{WithAllowRootSpan(true)}, 3},
294+
{[]TraceOption{WithAllowRootSpan(true)}, 5},
267295
}
268296

269297
for _, c := range testCases {
270298
db, _, recorder := createDB(t, c.opts...)
271299

272300
sqlStmt := `
301+
drop table if exists foo;
273302
create table foo (id integer not null primary key, name text);
274303
delete from foo;
275304
`
@@ -288,12 +317,7 @@ func TestTxWithRollbackSuccess(t *testing.T) {
288317
t.Fatalf("unexpected error: %s", err.Error())
289318
}
290319
defer stmt.Close()
291-
for i := 0; i < 100; i++ {
292-
_, err = stmt.Exec(i, fmt.Sprintf("こんにちわ世界%03d", i))
293-
if err != nil {
294-
t.Fatalf("unexpected error: %s", err.Error())
295-
}
296-
}
320+
_, err = stmt.Exec("1", "こんにちわ世界")
297321
tx.Rollback()
298322

299323
spans := recorder.Flush()
@@ -306,12 +330,25 @@ func TestTxWithRollbackSuccess(t *testing.T) {
306330
t.Fatalf("unexpected first span name, want: %s, have: %s", want, have)
307331
}
308332
if want, have := "sql/begin_transaction", spans[1].Name; want != have {
309-
t.Fatalf("unexpected first span name, want: %s, have: %s", want, have)
310-
}
311-
if want, have := "sql/rollback", spans[2].Name; want != have {
312333
t.Fatalf("unexpected second span name, want: %s, have: %s", want, have)
313334
}
335+
if want, have := "sql/prepare", spans[2].Name; want != have {
336+
t.Fatalf("unexpected third span name, want: %s, have: %s", want, have)
337+
}
338+
if want, have := "sql/exec", spans[3].Name; want != have {
339+
t.Fatalf("unexpected fourth span name, want: %s, have: %s", want, have)
340+
}
341+
if want, have := "sql/rollback", spans[4].Name; want != have {
342+
t.Fatalf("unexpected fifth span name, want: %s, have: %s", want, have)
343+
}
344+
345+
for i := 0; i < c.expectedSpans; i++ {
346+
if errMsg, ok := spans[i].Tags["error"]; ok {
347+
t.Fatalf("unexpected error: %s", errMsg)
348+
}
349+
}
314350
}
351+
315352
db.Close()
316353
recorder.Close()
317354
}

0 commit comments

Comments
 (0)