forked from DataDog/dd-trace-go
/
sql.go
60 lines (56 loc) · 2.04 KB
/
sql.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
// Package sql provides functions to trace the database/sql package (https://golang.org/pkg/database/sql).
// It will automatically augment operations such as connections, statements and transactions with tracing.
//
// We start by telling the package which driver we will be using. For example, if we are using "github.com/lib/pq",
// we would do as follows:
//
// sqltrace.Register("pq", pq.Driver{})
// db, err := sqltrace.Open("pq", "postgres://pqgotest:password@localhost...")
//
// The rest of our application would continue as usual, but with tracing enabled.
//
package sql
import (
"database/sql"
"database/sql/driver"
"errors"
)
// Register tells the sql integration package about the driver that we will be tracing. It must
// be called before Open, if that connection is to be traced. It uses the driverName suffixed
// with ".db" as the default service name.
func Register(driverName string, driver driver.Driver, opts ...RegisterOption) {
if driver == nil {
panic("sqltrace: Register driver is nil")
}
name := tracedDriverName(driverName)
if driverExists(name) {
// no problem, carry on
return
}
cfg := new(registerConfig)
defaults(cfg)
for _, fn := range opts {
fn(cfg)
}
if cfg.serviceName == "" {
cfg.serviceName = driverName + ".db"
}
sql.Register(name, &tracedDriver{
Driver: driver,
driverName: driverName,
config: cfg,
})
}
// errNotRegistered is returned when there is an attempt to open a database connection towards a driver
// that has not previously been registered using this package.
var errNotRegistered = errors.New("sqltrace: Register must be called before Open")
// Open returns connection to a DB using a the traced version of the given driver. In order for Open
// to work, the driver must first be registered using Register or RegisterWithServiceName. If this
// did not occur, Open will return an error.
func Open(driverName, dataSourceName string) (*sql.DB, error) {
name := tracedDriverName(driverName)
if !driverExists(name) {
return nil, errNotRegistered
}
return sql.Open(name, dataSourceName)
}