forked from Sonek-HoangBui/apm-agent-go
-
Notifications
You must be signed in to change notification settings - Fork 0
/
driver.go
138 lines (124 loc) · 3.72 KB
/
driver.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
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
package apmsql
import (
"database/sql"
"database/sql/driver"
"fmt"
"reflect"
"strings"
)
// DriverPrefix should be used as a driver name prefix when
// registering via sql.Register.
const DriverPrefix = "elasticapm/"
// Register registers a traced version of the given driver.
//
// The name and driver values should be the same as given to
// sql.Register: the name of the driver (e.g. "postgres"), and
// the driver (e.g. &github.com/lib/pq.Driver{}).
func Register(name string, driver driver.Driver, opts ...WrapOption) {
wrapped := Wrap(driver, opts...)
sql.Register(DriverPrefix+name, wrapped)
}
// Open opens a database with the given driver and data source names,
// as in sql.Open. The driver name should be one registered via the
// Register function in this package.
func Open(driverName, dataSourceName string) (*sql.DB, error) {
return sql.Open(DriverPrefix+driverName, dataSourceName)
}
// Wrap wraps a database/sql/driver.Driver such that
// the driver's database methods are traced. The tracer
// will be obtained from the context supplied to methods
// that accept it.
func Wrap(driver driver.Driver, opts ...WrapOption) driver.Driver {
d := &tracingDriver{
Driver: driver,
}
for _, opt := range opts {
opt(d)
}
if d.driverName == "" {
d.driverName = driverName(driver)
}
if d.dsnParser == nil {
d.dsnParser = genericDSNParser
}
// store span types to avoid repeat allocations
d.connectSpanType = d.formatSpanType("connect")
d.pingSpanType = d.formatSpanType("ping")
d.prepareSpanType = d.formatSpanType("prepare")
d.querySpanType = d.formatSpanType("query")
d.execSpanType = d.formatSpanType("exec")
return d
}
// WrapOption is an option that can be supplied to Wrap.
type WrapOption func(*tracingDriver)
// WithDriverName returns a WrapOption which sets the underlying
// driver name to the specified value. If WithDriverName is not
// supplied to Wrap, the driver name will be inferred from the
// driver supplied to Wrap.
func WithDriverName(name string) WrapOption {
return func(d *tracingDriver) {
d.driverName = name
}
}
// WithDSNParser returns a WrapOption which sets the function to
// use for parsing the data source name. If WithDSNParser is not
// supplied to Wrap, the function to use will be inferred from
// the driver name.
func WithDSNParser(f DSNParserFunc) WrapOption {
return func(d *tracingDriver) {
d.dsnParser = f
}
}
type tracingDriver struct {
driver.Driver
driverName string
dsnParser DSNParserFunc
connectSpanType string
execSpanType string
pingSpanType string
prepareSpanType string
querySpanType string
}
func (d *tracingDriver) formatSpanType(suffix string) string {
return fmt.Sprintf("db.%s.%s", d.driverName, suffix)
}
// querySignature returns the value to use in Span.Name for
// a database query.
func (d *tracingDriver) querySignature(query string) string {
// TODO(axw) parse statement. Create a WrapOption for
// consumers to provide a driver-specific parser if
// necessary.
fields := strings.Fields(query)
if len(fields) == 0 {
return ""
}
return strings.ToUpper(fields[0])
}
func (d *tracingDriver) Open(name string) (driver.Conn, error) {
conn, err := d.Driver.Open(name)
if err != nil {
return nil, err
}
return newConn(conn, d, d.dsnParser(name)), nil
}
func driverName(d driver.Driver) string {
t := reflect.TypeOf(d)
if t.Kind() == reflect.Ptr {
t = t.Elem()
}
switch t.Name() {
case "SQLiteDriver":
return "sqlite3"
case "MySQLDriver":
return "mysql"
case "Driver":
// Check suffix in case of vendoring.
if strings.HasSuffix(t.PkgPath(), "github.com/lib/pq") {
return "postgresql"
}
}
// TODO include the package path of the driver in context
// so we can easily determine how the rules above should
// be updated.
return "generic"
}