forked from luna-duclos/instrumentedsql
-
Notifications
You must be signed in to change notification settings - Fork 0
/
connector.go
66 lines (53 loc) · 1.38 KB
/
connector.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
// +build go1.10
package instrumentedsql
import (
"context"
"database/sql/driver"
)
type wrappedConnector struct {
parent driver.Connector
driverRef *wrappedDriver
}
func (d wrappedDriver) OpenConnector(name string) (driver.Connector, error) {
driver, ok := d.parent.(driver.DriverContext)
if !ok {
return wrappedConnector{
parent: dsnConnector{dsn: name, driver: &d},
driverRef: &d,
}, nil
}
conn, err := driver.OpenConnector(name)
if err != nil {
return nil, err
}
return wrappedConnector{parent: conn, driverRef: &d}, nil
}
func (c wrappedConnector) Connect(ctx context.Context) (driver.Conn, error) {
conn, err := c.parent.Connect(ctx)
if err != nil {
return nil, err
}
return wrappedConn{opts: c.driverRef.opts, parent: conn}, nil
}
func (c wrappedConnector) Driver() driver.Driver {
return c.driverRef
}
func (c wrappedConn) ResetSession(ctx context.Context) error {
conn, ok := c.parent.(driver.SessionResetter)
if !ok {
return nil
}
return conn.ResetSession(ctx)
}
// dsnConnector is a fallback connector placed in position of wrappedConnector.parent
// when given Driver does not comply with DriverContext interface.
type dsnConnector struct {
dsn string
driver driver.Driver
}
func (t dsnConnector) Connect(_ context.Context) (driver.Conn, error) {
return t.driver.Open(t.dsn)
}
func (t dsnConnector) Driver() driver.Driver {
return t.driver
}