From d988f5c3b064fade6e96511e0024190c22d48e50 Mon Sep 17 00:00:00 2001
From: Kristin Laemmert <mildwonkey@users.noreply.github.com>
Date: Thu, 20 Jun 2024 14:45:39 -0400
Subject: [PATCH] fix: sql instrumentation dual registration error (#89508)

fix dual registration error
---
 pkg/services/sqlstore/database_wrapper.go | 21 +++++++++++++++++++++
 pkg/services/sqlstore/replstore.go        |  2 +-
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/pkg/services/sqlstore/database_wrapper.go b/pkg/services/sqlstore/database_wrapper.go
index fab3664cd1c..7cafa1f712b 100644
--- a/pkg/services/sqlstore/database_wrapper.go
+++ b/pkg/services/sqlstore/database_wrapper.go
@@ -58,6 +58,27 @@ func WrapDatabaseDriverWithHooks(dbType string, tracer tracing.Tracer) string {
 	return driverWithHooks
 }
 
+// WrapDatabaseDriverWithHooks creates a fake database driver that
+// executes pre and post functions which we use to gather metrics about
+// database queries. It also registers the metrics.
+func WrapDatabaseReplDriverWithHooks(dbType string, tracer tracing.Tracer) string {
+	drivers := map[string]driver.Driver{
+		migrator.SQLite:   &sqlite3.SQLiteDriver{},
+		migrator.MySQL:    &mysql.MySQLDriver{},
+		migrator.Postgres: &pq.Driver{},
+	}
+
+	d, exist := drivers[dbType]
+	if !exist {
+		return dbType
+	}
+
+	driverWithHooks := dbType + "ReplicaWithHooks"
+	sql.Register(driverWithHooks, sqlhooks.Wrap(d, &databaseQueryWrapper{log: log.New("sqlstore.metrics"), tracer: tracer}))
+	core.RegisterDriver(driverWithHooks, &databaseQueryWrapperDriver{dbType: dbType})
+	return driverWithHooks
+}
+
 // databaseQueryWrapper satisfies the sqlhook.databaseQueryWrapper interface
 // which allow us to wrap all SQL queries with a `Before` & `After` hook.
 type databaseQueryWrapper struct {
diff --git a/pkg/services/sqlstore/replstore.go b/pkg/services/sqlstore/replstore.go
index c3540d79c4f..050341a3606 100644
--- a/pkg/services/sqlstore/replstore.go
+++ b/pkg/services/sqlstore/replstore.go
@@ -118,7 +118,7 @@ func (ss *SQLStore) initReadOnlyEngine(engine *xorm.Engine) error {
 	ss.dbCfg = dbCfg
 
 	if ss.cfg.DatabaseInstrumentQueries {
-		ss.dbCfg.Type = WrapDatabaseDriverWithHooks(ss.dbCfg.Type, ss.tracer)
+		ss.dbCfg.Type = WrapDatabaseReplDriverWithHooks(ss.dbCfg.Type, ss.tracer)
 	}
 
 	if engine == nil {