forked from inkbe/opentracing-sql
/
conn.go
122 lines (110 loc) · 3.19 KB
/
conn.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
package sql
import (
"context"
"database/sql/driver"
)
// conn defines a tracing wrapper for driver.Conn.
type conn struct {
conn driver.Conn
tracer *tracer
}
// Prepare implements driver.Conn Prepare.
func (c *conn) Prepare(query string) (driver.Stmt, error) {
s, err := c.conn.Prepare(query)
if err != nil {
return nil, err
}
return &stmt{stmt: s, tracer: c.tracer}, nil
}
// Prepare implements driver.Conn Close.
func (c *conn) Close() error {
return c.conn.Close()
}
// Prepare implements driver.Conn Begin.
func (c *conn) Begin() (driver.Tx, error) {
t, err := c.conn.Begin()
if err != nil {
return nil, err
}
return &tx{tx: t, tracer: c.tracer}, nil
}
// BeginTx implements driver.ConnBeginTx BeginTx.
func (c *conn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
s := c.tracer.newSpan(ctx)
if connBeginTx, ok := c.conn.(driver.ConnBeginTx); ok {
t, err := connBeginTx.BeginTx(ctx, opts)
if err != nil {
return nil, err
}
return &tx{tx: t, tracer: c.tracer, span: s}, nil
}
return c.conn.Begin()
}
// PrepareContext implements driver.ConnPrepareContext PrepareContext.
func (c *conn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
if connPrepareContext, ok := c.conn.(driver.ConnPrepareContext); ok {
s, err := connPrepareContext.PrepareContext(ctx, query)
if err != nil {
return nil, err
}
return &stmt{stmt: s, tracer: c.tracer}, nil
}
return c.conn.Prepare(query)
}
// Exec implements driver.Execer Exec.
func (c *conn) Exec(query string, args []driver.Value) (driver.Result, error) {
if execer, ok := c.conn.(driver.Execer); ok {
return execer.Exec(query, args)
}
return nil, ErrUnsupported
}
// Exec implements driver.StmtExecContext ExecContext.
func (c *conn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
s := c.tracer.newSpan(ctx)
if c.tracer.saveQuery {
s.SetTag(spanTagQuery, query)
}
defer s.Finish()
if execerContext, ok := c.conn.(driver.ExecerContext); ok {
r, err := execerContext.ExecContext(ctx, query, args)
return r, err
}
values, err := namedValueToValue(args)
if err != nil {
return nil, err
}
return c.Exec(query, values)
}
// Ping implements driver.Pinger Ping.
func (c *conn) Ping(ctx context.Context) error {
if pinger, ok := c.conn.(driver.Pinger); ok {
s := c.tracer.newSpan(ctx)
defer s.Finish()
return pinger.Ping(ctx)
}
return ErrUnsupported
}
// Query implements driver.Queryer Query.
func (c *conn) Query(query string, args []driver.Value) (driver.Rows, error) {
if queryer, ok := c.conn.(driver.Queryer); ok {
return queryer.Query(query, args)
}
return nil, ErrUnsupported
}
// QueryContext implements Driver.QueryerContext QueryContext.
func (c *conn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (rows driver.Rows, err error) {
s := c.tracer.newSpan(ctx)
if c.tracer.saveQuery {
s.SetTag(spanTagQuery, query)
}
defer s.Finish()
if queryerContext, ok := c.conn.(driver.QueryerContext); ok {
rows, err := queryerContext.QueryContext(ctx, query, args)
return rows, err
}
values, err := namedValueToValue(args)
if err != nil {
return nil, err
}
return c.Query(query, values)
}