From fc38460bc5346162cffd3533d2e64a6d3cc2b4b8 Mon Sep 17 00:00:00 2001
From: Daniel Bedrood <danielb@scribd.com>
Date: Thu, 20 Jun 2024 16:08:31 +0200
Subject: [PATCH] feat: Add MySQL driver tracing

---
 pkg/database/gorm.go | 46 ++++++++++++++++++--------------------------
 1 file changed, 19 insertions(+), 27 deletions(-)

diff --git a/pkg/database/gorm.go b/pkg/database/gorm.go
index 4dab3a8..2959a16 100644
--- a/pkg/database/gorm.go
+++ b/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
 }