Merge pull request #112 from scribd/laynax/SERF-3423/trace-driver

[SERF-3423] Adding MySQL driver instrumentation

Author: laynax

pkg/database/gorm.go

@@ -1,11 +1,14 @@
 package database
 
 import (
+	"database/sql"
 	"fmt"
 	"strconv"
 	"time"
 
 	"github.com/DATA-DOG/go-txdb"
+	mysqldriver "github.com/go-sql-driver/mysql"
+	sqltrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql"
 	gormtrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/gorm.io/gorm.v1"
 	"gorm.io/driver/mysql"
 	"gorm.io/gorm"
@@ -15,24 +18,7 @@ const testEnv = "test"
 
 // NewConnection returns a new instrumented Gorm database connection.
 func NewConnection(config *Config, environment, appName string) (*gorm.DB, error) {
-	serviceName := fmt.Sprintf("%s-mysql", appName)
-	dialector := getDialectorFromConfig(config, environment)
-
-	db, err := gormtrace.Open(dialector, nil, gormtrace.WithServiceName(serviceName))
-	if err != nil {
-		return nil, err
-	}
-
-	if err := databasePoolSettings(db, config); err != nil {
-		return nil, err
-	}
-
-	return db, nil
-}
-
-func getDialectorFromConfig(config *Config, environment string) gorm.Dialector {
 	connectionDetails := NewConnectionDetails(config)
-
 	connectionString := connectionDetails.String()
 	driverName := connectionDetails.Dialect
 
@@ -47,22 +33,28 @@ func getDialectorFromConfig(config *Config, environment string) gorm.Dialector {
 		connectionString = testDriverName
 	}
 
-	return mysql.New(mysql.Config{
-		DSN:        connectionString,
-		DriverName: driverName,
-	})
-}
+	serviceName := fmt.Sprintf("%s-mysql", appName)
 
-func databasePoolSettings(gormDB *gorm.DB, config *Config) error {
-	db, err := gormDB.DB()
+	sqltrace.Register(driverName, &mysqldriver.MySQLDriver{}, sqltrace.WithServiceName(serviceName))
+	sqlDB, err := sqltrace.Open(driverName, connectionString)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
+	databasePoolSettings(sqlDB, config)
+
+	dialector := mysql.New(mysql.Config{Conn: sqlDB})
+	db, err := gormtrace.Open(dialector, nil, gormtrace.WithServiceName(serviceName))
+	if err != nil {
+		return nil, err
+	}
+
+	return db, nil
+}
+
+func databasePoolSettings(db *sql.DB, config *Config) {
 	db.SetMaxIdleConns(config.Pool)
 	db.SetMaxOpenConns(config.MaxOpenConnections)
 	db.SetConnMaxIdleTime(config.ConnectionMaxIdleTime)
 	db.SetConnMaxLifetime(config.ConnectionMaxLifetime)
-
-	return nil
 }