chore (replstore): fix registration with multiple sql drivers, again (#90990)

* replstore: fix registration with multiple sql drivers, again
* only compile regex once

pkg/services/sqlstore/database_wrapper.go

@@ -61,7 +61,7 @@ func WrapDatabaseDriverWithHooks(dbType string, tracer tracing.Tracer) string {
 // 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 {
+func WrapDatabaseReplDriverWithHooks(dbType string, index uint, tracer tracing.Tracer) string {
 	drivers := map[string]driver.Driver{
 		migrator.SQLite:   &sqlite3.SQLiteDriver{},
 		migrator.MySQL:    &mysql.MySQLDriver{},
@@ -73,7 +73,7 @@ func WrapDatabaseReplDriverWithHooks(dbType string, tracer tracing.Tracer) strin
 		return dbType
 	}
 
-	driverWithHooks := dbType + "ReplicaWithHooks"
+	driverWithHooks := dbType + fmt.Sprintf("ReplicaWithHooks%d", index)
 	sql.Register(driverWithHooks, sqlhooks.Wrap(d, &databaseQueryWrapper{log: log.New("sqlstore.metrics"), tracer: tracer}))
 	core.RegisterDriver(driverWithHooks, &databaseQueryWrapperDriver{dbType: dbType})
 	return driverWithHooks

pkg/services/sqlstore/migrator/dialect.go

@@ -107,13 +107,14 @@ type LockCfg struct {
 type dialectFunc func() Dialect
 
 var supportedDialects = map[string]dialectFunc{
-	MySQL:                      NewMysqlDialect,
-	SQLite:                     NewSQLite3Dialect,
-	Postgres:                   NewPostgresDialect,
-	MySQL + "WithHooks":        NewMysqlDialect,
-	MySQL + "ReplicaWithHooks": NewMysqlDialect,
-	SQLite + "WithHooks":       NewSQLite3Dialect,
-	Postgres + "WithHooks":     NewPostgresDialect,
+	MySQL:                         NewMysqlDialect,
+	SQLite:                        NewSQLite3Dialect,
+	Postgres:                      NewPostgresDialect,
+	MySQL + "WithHooks":           NewMysqlDialect,
+	MySQL + "ReplicaWithHooks":    NewMysqlDialect,
+	SQLite + "WithHooks":          NewSQLite3Dialect,
+	Postgres + "WithHooks":        NewPostgresDialect,
+	Postgres + "ReplicaWithHooks": NewPostgresDialect,
 }
 
 func NewDialect(driverName string) Dialect {

pkg/services/sqlstore/replstore.go

@@ -2,6 +2,7 @@ package sqlstore
 
 import (
 	"errors"
+	"regexp"
 	"sync/atomic"
 	"time"
 
@@ -43,7 +44,6 @@ func (rs *ReplStore) ReadReplica() *SQLStore {
 		rs.log.Debug("ReadReplica not configured, using main SQLStore")
 		return rs.SQLStore
 	}
-	rs.log.Debug("Using ReadReplica")
 	return rs.nextRepl()
 }
 
@@ -89,6 +89,11 @@ func ProvideServiceWithReadReplica(primary *SQLStore, cfg *setting.Cfg,
 	}
 
 	for i, replCfg := range replCfgs {
+		// If the database_instrument_queries feature is enabled, wrap the driver with hooks.
+		if cfg.DatabaseInstrumentQueries {
+			replCfg.Type = WrapDatabaseReplDriverWithHooks(replCfg.Type, uint(i), tracer)
+		}
+
 		s, err := newReadOnlySQLStore(cfg, replCfg, features, bus, tracer)
 		if err != nil {
 			return nil, err
@@ -124,10 +129,18 @@ func newReadOnlySQLStore(cfg *setting.Cfg, dbCfg *DatabaseConfig, features featu
 	if err != nil {
 		return nil, err
 	}
-	s.dialect = migrator.NewDialect(s.engine.DriverName())
+
+	// When there are multiple read replicas, we append an index to the driver name (ex: mysqlWithHooks11).
+	// Remove the index from the end of the driver name to get the original driver name that xorm and other libraries recognize.
+	driverName := digitsRegexp.ReplaceAllString(s.engine.DriverName(), "")
+
+	s.dialect = migrator.NewDialect(driverName)
 	return s, nil
 }
 
+// digitsRegexp is used to remove the index from the end of the driver name.
+var digitsRegexp = regexp.MustCompile("[0-9]+")
+
 // initReadOnlyEngine initializes ss.engine for read-only operations. The database must be a fully-populated read replica.
 func (ss *SQLStore) initReadOnlyEngine(engine *xorm.Engine) error {
 	if ss.engine != nil {
@@ -135,10 +148,6 @@ func (ss *SQLStore) initReadOnlyEngine(engine *xorm.Engine) error {
 		return nil
 	}
 
-	if ss.cfg.DatabaseInstrumentQueries {
-		ss.dbCfg.Type = WrapDatabaseReplDriverWithHooks(ss.dbCfg.Type, ss.tracer)
-	}
-
 	if engine == nil {
 		var err error
 		engine, err = xorm.NewEngine(ss.dbCfg.Type, ss.dbCfg.ConnectionString)