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

* replstore: fix registration with multiple sql drivers, again
* only compile regex once
This commit is contained in:
Kristin Laemmert 2024-07-25 15:13:36 -04:00 committed by GitHub
parent 9852513c65
commit daedb358dd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 25 additions and 15 deletions

View File

@ -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

View File

@ -114,6 +114,7 @@ var supportedDialects = map[string]dialectFunc{
MySQL + "ReplicaWithHooks": NewMysqlDialect,
SQLite + "WithHooks": NewSQLite3Dialect,
Postgres + "WithHooks": NewPostgresDialect,
Postgres + "ReplicaWithHooks": NewPostgresDialect,
}
func NewDialect(driverName string) Dialect {

View File

@ -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)