-
-
Notifications
You must be signed in to change notification settings - Fork 242
/
connection_instrumented.go
98 lines (82 loc) · 2.74 KB
/
connection_instrumented.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
package pop
import (
"database/sql"
"database/sql/driver"
"fmt"
"sync"
mysqld "github.com/go-sql-driver/mysql"
"github.com/gobuffalo/pop/v6/logging"
pgx "github.com/jackc/pgx/v4/stdlib"
"github.com/jmoiron/sqlx"
"github.com/luna-duclos/instrumentedsql"
)
const instrumentedDriverName = "instrumented-sql-driver"
var sqlDriverLock = sync.Mutex{}
func instrumentDriver(deets *ConnectionDetails, defaultDriverName string) (driverName, dialect string, err error) {
driverName = defaultDriverName
if deets.Driver != "" {
driverName = deets.Driver
}
dialect = driverName
if !deets.UseInstrumentedDriver {
if len(deets.InstrumentedDriverOptions) > 0 {
log(logging.Warn, "SQL driver instrumentation is disabled but `ConnectionDetails.InstrumentedDriverOptions` is not empty. Please double-check if this is a error.")
}
// If instrumentation is disabled, we just return the driver name we got (e.g. "pgx").
return driverName, dialect, nil
}
if len(deets.InstrumentedDriverOptions) == 0 {
log(logging.Warn, "SQL driver instrumentation was enabled but no options have been passed to `ConnectionDetails.InstrumentedDriverOptions`. Instrumentation will therefore not result in any output.")
}
var dr driver.Driver
var newDriverName string
switch CanonicalDialect(driverName) {
case nameCockroach:
fallthrough
case namePostgreSQL:
dr = new(pgx.Driver)
newDriverName = instrumentedDriverName + "-" + namePostgreSQL
case nameMariaDB:
fallthrough
case nameMySQL:
dr = mysqld.MySQLDriver{}
newDriverName = instrumentedDriverName + "-" + nameMySQL
case nameSQLite3:
var err error
dr, err = newSQLiteDriver()
if err != nil {
return "", "", err
}
newDriverName = instrumentedDriverName + "-" + nameSQLite3
}
sqlDriverLock.Lock()
defer sqlDriverLock.Unlock()
var found bool
for _, n := range sql.Drivers() {
if n == newDriverName {
found = true
break
}
}
if !found {
sql.Register(newDriverName, instrumentedsql.WrapDriver(dr, deets.InstrumentedDriverOptions...))
}
return newDriverName, dialect, nil
}
// openPotentiallyInstrumentedConnection first opens a raw SQL connection and then wraps it with `sqlx`.
//
// We do this because `sqlx` needs the database type in order to properly
// translate arguments (e.g. `?` to `$1`) in SQL queries. Because we use
// a custom driver name when using instrumentation, this detection would fail
// otherwise.
func openPotentiallyInstrumentedConnection(c dialect, dsn string) (*sqlx.DB, error) {
driverName, dialect, err := instrumentDriver(c.Details(), c.DefaultDriver())
if err != nil {
return nil, err
}
con, err := sql.Open(driverName, dsn)
if err != nil {
return nil, fmt.Errorf("could not open database connection: %w", err)
}
return sqlx.NewDb(con, dialect), nil
}