feat: use XORM EngineGroup instead of single Engine connection (go-gitea#7212)

Resolves go-gitea#7207

Add new configuration to make XORM work with a main and replicas database instances. The follow configuration parameters were added:

- `HOST_PRIMARY`
- `HOST_REPLICAS`
- `LOAD_BALANCE_POLICY`. Options:
    - `"WeightRandom"` -> `xorm.WeightRandomPolicy`
    - `"WeightRoundRobin`  -> `WeightRoundRobinPolicy`
    - `"LeastCon"` -> `LeastConnPolicy`
    - `"RoundRobin"` -> `xorm.RoundRobinPolicy()`
    - default: `xorm.RandomPolicy()`
- `LOAD_BALANCE_WEIGHTS`

Co-authored-by: pat-s <patrick.schratz@gmail.com@>
Reviewed-on: https://codeberg.org/forgejo/forgejo/pulls/7212
Reviewed-by: Gusted <gusted@noreply.codeberg.org>
Co-authored-by: pat-s <patrick.schratz@gmail.com>
Co-committed-by: pat-s <patrick.schratz@gmail.com>

Author: 2 people

cmd/doctor.go

@@ -20,7 +20,6 @@ import (
 	"forgejo.org/services/doctor"
 
 	"github.com/urfave/cli/v2"
-	"xorm.io/xorm"
 )
 
 // CmdDoctor represents the available doctor sub-command.
@@ -120,7 +119,7 @@ func runRecreateTable(ctx *cli.Context) error {
 
 	args := ctx.Args()
 	names := make([]string, 0, ctx.NArg())
-	for i := 0; i < ctx.NArg(); i++ {
+	for i := range ctx.NArg() {
 		names = append(names, args.Get(i))
 	}
 
@@ -130,11 +129,17 @@ func runRecreateTable(ctx *cli.Context) error {
 	}
 	recreateTables := migrate_base.RecreateTables(beans...)
 
-	return db.InitEngineWithMigration(stdCtx, func(x *xorm.Engine) error {
-		if err := migrations.EnsureUpToDate(x); err != nil {
+	return db.InitEngineWithMigration(stdCtx, func(x db.Engine) error {
+		engine, err := db.GetMasterEngine(x)
+		if err != nil {
 			return err
 		}
-		return recreateTables(x)
+
+		if err := migrations.EnsureUpToDate(engine); err != nil {
+			return err
+		}
+
+		return recreateTables(engine)
 	})
 }
 

cmd/migrate.go

@@ -36,7 +36,13 @@ func runMigrate(ctx *cli.Context) error {
 	log.Info("Log path: %s", setting.Log.RootPath)
 	log.Info("Configuration file: %s", setting.CustomConf)
 
-	if err := db.InitEngineWithMigration(context.Background(), migrations.Migrate); err != nil {
+	if err := db.InitEngineWithMigration(context.Background(), func(dbEngine db.Engine) error {
+		masterEngine, err := db.GetMasterEngine(dbEngine)
+		if err != nil {
+			return err
+		}
+		return migrations.Migrate(masterEngine)
+	}); err != nil {
 		log.Fatal("Failed to initialize ORM engine: %v", err)
 		return err
 	}

cmd/migrate_storage.go

@@ -23,6 +23,7 @@ import (
 	"forgejo.org/modules/storage"
 
 	"github.com/urfave/cli/v2"
+	"xorm.io/xorm"
 )
 
 // CmdMigrateStorage represents the available migrate storage sub-command.
@@ -195,7 +196,9 @@ func runMigrateStorage(ctx *cli.Context) error {
 	log.Info("Log path: %s", setting.Log.RootPath)
 	log.Info("Configuration file: %s", setting.CustomConf)
 
-	if err := db.InitEngineWithMigration(context.Background(), migrations.Migrate); err != nil {
+	if err := db.InitEngineWithMigration(context.Background(), func(e db.Engine) error {
+		return migrations.Migrate(e.(*xorm.Engine))
+	}); err != nil {
 		log.Fatal("Failed to initialize ORM engine: %v", err)
 		return err
 	}

models/db/engine.go

@@ -95,87 +95,132 @@ func init() {
 	}
 }
 
-// newXORMEngine returns a new XORM engine from the configuration
-func newXORMEngine() (*xorm.Engine, error) {
-	connStr, err := setting.DBConnStr()
+// newXORMEngineGroup creates an xorm.EngineGroup (with one master and one or more slaves).
+// It assumes you have separate master and slave DSNs defined via the settings package.
+func newXORMEngineGroup() (Engine, error) {
+	// Retrieve master DSN from settings.
+	masterConnStr, err := setting.DBMasterConnStr()
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("failed to determine master DSN: %w", err)
 	}
 
-	var engine *xorm.Engine
-
+	var masterEngine *xorm.Engine
+	// For PostgreSQL: if a schema is provided, we use the special "postgresschema" driver.
 	if setting.Database.Type.IsPostgreSQL() && len(setting.Database.Schema) > 0 {
-		// OK whilst we sort out our schema issues - create a schema aware postgres
 		registerPostgresSchemaDriver()
-		engine, err = xorm.NewEngine("postgresschema", connStr)
+		masterEngine, err = xorm.NewEngine("postgresschema", masterConnStr)
 	} else {
-		engine, err = xorm.NewEngine(setting.Database.Type.String(), connStr)
+		masterEngine, err = xorm.NewEngine(setting.Database.Type.String(), masterConnStr)
 	}
-
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("failed to create master engine: %w", err)
 	}
 	if setting.Database.Type.IsMySQL() {
-		engine.Dialect().SetParams(map[string]string{"rowFormat": "DYNAMIC"})
+		masterEngine.Dialect().SetParams(map[string]string{"rowFormat": "DYNAMIC"})
 	}
-	engine.SetSchema(setting.Database.Schema)
-	return engine, nil
+	masterEngine.SetSchema(setting.Database.Schema)
+
+	slaveConnStrs, err := setting.DBSlaveConnStrs()
+	if err != nil {
+		return nil, fmt.Errorf("failed to load slave DSNs: %w", err)
+	}
+
+	var slaveEngines []*xorm.Engine
+	// Iterate over all slave DSNs and create engines
+	for _, dsn := range slaveConnStrs {
+		slaveEngine, err := xorm.NewEngine(setting.Database.Type.String(), dsn)
+		if err != nil {
+			return nil, fmt.Errorf("failed to create slave engine for dsn %q: %w", dsn, err)
+		}
+		if setting.Database.Type.IsMySQL() {
+			slaveEngine.Dialect().SetParams(map[string]string{"rowFormat": "DYNAMIC"})
+		}
+		slaveEngine.SetSchema(setting.Database.Schema)
+		slaveEngines = append(slaveEngines, slaveEngine)
+	}
+
+	policy := setting.BuildLoadBalancePolicy(&setting.Database, slaveEngines)
+
+	// Create the EngineGroup using the selected policy
+	group, err := xorm.NewEngineGroup(masterEngine, slaveEngines, policy)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create engine group: %w", err)
+	}
+	return engineGroupWrapper{group}, nil
+}
+
+type engineGroupWrapper struct {
+	*xorm.EngineGroup
+}
+
+func (w engineGroupWrapper) AddHook(hook contexts.Hook) bool {
+	w.EngineGroup.AddHook(hook)
+	return true
 }
 
-// SyncAllTables sync the schemas of all tables, is required by unit test code
+// SyncAllTables sync the schemas of all tables
 func SyncAllTables() error {
 	_, err := x.StoreEngine("InnoDB").SyncWithOptions(xorm.SyncOptions{
 		WarnIfDatabaseColumnMissed: true,
 	}, tables...)
 	return err
 }
 
-// InitEngine initializes the xorm.Engine and sets it as db.DefaultContext
+// InitEngine initializes the xorm EngineGroup and sets it as db.DefaultContext
 func InitEngine(ctx context.Context) error {
-	xormEngine, err := newXORMEngine()
+	xormEngine, err := newXORMEngineGroup()
 	if err != nil {
 		return fmt.Errorf("failed to connect to database: %w", err)
 	}
+	// Try to cast to the concrete type to access diagnostic methods
+	if eng, ok := xormEngine.(engineGroupWrapper); ok {
+		eng.SetMapper(names.GonicMapper{})
+		// WARNING: for serv command, MUST remove the output to os.Stdout,
+		// so use a log file instead of printing to stdout.
+		eng.SetLogger(NewXORMLogger(setting.Database.LogSQL))
+		eng.ShowSQL(setting.Database.LogSQL)
+		eng.SetMaxOpenConns(setting.Database.MaxOpenConns)
+		eng.SetMaxIdleConns(setting.Database.MaxIdleConns)
+		eng.SetConnMaxLifetime(setting.Database.ConnMaxLifetime)
+		eng.SetConnMaxIdleTime(setting.Database.ConnMaxIdleTime)
+		eng.SetDefaultContext(ctx)
+
+		if setting.Database.SlowQueryThreshold > 0 {
+			eng.AddHook(&SlowQueryHook{
+				Treshold: setting.Database.SlowQueryThreshold,
+				Logger:   log.GetLogger("xorm"),
+			})
+		}
 
-	xormEngine.SetMapper(names.GonicMapper{})
-	// WARNING: for serv command, MUST remove the output to os.stdout,
-	// so use log file to instead print to stdout.
-	xormEngine.SetLogger(NewXORMLogger(setting.Database.LogSQL))
-	xormEngine.ShowSQL(setting.Database.LogSQL)
-	xormEngine.SetMaxOpenConns(setting.Database.MaxOpenConns)
-	xormEngine.SetMaxIdleConns(setting.Database.MaxIdleConns)
-	xormEngine.SetConnMaxLifetime(setting.Database.ConnMaxLifetime)
-	xormEngine.SetConnMaxIdleTime(setting.Database.ConnMaxIdleTime)
-	xormEngine.SetDefaultContext(ctx)
-
-	if setting.Database.SlowQueryThreshold > 0 {
-		xormEngine.AddHook(&SlowQueryHook{
-			Treshold: setting.Database.SlowQueryThreshold,
-			Logger:   log.GetLogger("xorm"),
-		})
-	}
-
-	errorLogger := log.GetLogger("xorm")
-	if setting.IsInTesting {
-		errorLogger = log.GetLogger(log.DEFAULT)
-	}
+		errorLogger := log.GetLogger("xorm")
+		if setting.IsInTesting {
+			errorLogger = log.GetLogger(log.DEFAULT)
+		}
 
-	xormEngine.AddHook(&ErrorQueryHook{
-		Logger: errorLogger,
-	})
+		eng.AddHook(&ErrorQueryHook{
+			Logger: errorLogger,
+		})
 
-	xormEngine.AddHook(&TracingHook{})
+		eng.AddHook(&TracingHook{})
 
-	SetDefaultEngine(ctx, xormEngine)
+		SetDefaultEngine(ctx, eng)
+	} else {
+		// Fallback: if type assertion fails, set default engine without extended diagnostics
+		SetDefaultEngine(ctx, xormEngine)
+	}
 	return nil
 }
 
-// SetDefaultEngine sets the default engine for db
-func SetDefaultEngine(ctx context.Context, eng *xorm.Engine) {
-	x = eng
+// SetDefaultEngine sets the default engine for db.
+func SetDefaultEngine(ctx context.Context, eng Engine) {
+	masterEngine, err := GetMasterEngine(eng)
+	if err == nil {
+		x = masterEngine
+	}
+
 	DefaultContext = &Context{
 		Context: ctx,
-		e:       x,
+		e:       eng,
 	}
 }
 
@@ -191,12 +236,12 @@ func UnsetDefaultEngine() {
 	DefaultContext = nil
 }
 
-// InitEngineWithMigration initializes a new xorm.Engine and sets it as the db.DefaultContext
+// InitEngineWithMigration initializes a new xorm EngineGroup, runs migrations, and sets it as db.DefaultContext
 // This function must never call .Sync() if the provided migration function fails.
 // When called from the "doctor" command, the migration function is a version check
 // that prevents the doctor from fixing anything in the database if the migration level
 // is different from the expected value.
-func InitEngineWithMigration(ctx context.Context, migrateFunc func(*xorm.Engine) error) (err error) {
+func InitEngineWithMigration(ctx context.Context, migrateFunc func(Engine) error) (err error) {
 	if err = InitEngine(ctx); err != nil {
 		return err
 	}
@@ -230,14 +275,14 @@ func InitEngineWithMigration(ctx context.Context, migrateFunc func(*xorm.Engine)
 	return nil
 }
 
-// NamesToBean return a list of beans or an error
+// NamesToBean returns a list of beans given names
 func NamesToBean(names ...string) ([]any, error) {
 	beans := []any{}
 	if len(names) == 0 {
 		beans = append(beans, tables...)
 		return beans, nil
 	}
-	// Need to map provided names to beans...
+	// Map provided names to beans
 	beanMap := make(map[string]any)
 	for _, bean := range tables {
 		beanMap[strings.ToLower(reflect.Indirect(reflect.ValueOf(bean)).Type().Name())] = bean
@@ -259,7 +304,7 @@ func NamesToBean(names ...string) ([]any, error) {
 	return beans, nil
 }
 
-// DumpDatabase dumps all data from database according the special database SQL syntax to file system.
+// DumpDatabase dumps all data from database using special SQL syntax to the file system.
 func DumpDatabase(filePath, dbType string) error {
 	var tbs []*schemas.Table
 	for _, t := range tables {
@@ -295,29 +340,33 @@ func MaxBatchInsertSize(bean any) int {
 	return 999 / len(t.ColumnsSeq())
 }
 
-// IsTableNotEmpty returns true if table has at least one record
+// IsTableNotEmpty returns true if the table has at least one record
 func IsTableNotEmpty(beanOrTableName any) (bool, error) {
 	return x.Table(beanOrTableName).Exist()
 }
 
-// DeleteAllRecords will delete all the records of this table
+// DeleteAllRecords deletes all records in the given table.
 func DeleteAllRecords(tableName string) error {
 	_, err := x.Exec(fmt.Sprintf("DELETE FROM %s", tableName))
 	return err
 }
 
-// GetMaxID will return max id of the table
+// GetMaxID returns the maximum id in the table
 func GetMaxID(beanOrTableName any) (maxID int64, err error) {
 	_, err = x.Select("MAX(id)").Table(beanOrTableName).Get(&maxID)
 	return maxID, err
 }
 
 func SetLogSQL(ctx context.Context, on bool) {
-	e := GetEngine(ctx)
-	if x, ok := e.(*xorm.Engine); ok {
-		x.ShowSQL(on)
-	} else if sess, ok := e.(*xorm.Session); ok {
+	ctxEngine := GetEngine(ctx)
+
+	if sess, ok := ctxEngine.(*xorm.Session); ok {
 		sess.Engine().ShowSQL(on)
+	} else if wrapper, ok := ctxEngine.(engineGroupWrapper); ok {
+		// Handle engineGroupWrapper directly
+		wrapper.ShowSQL(on)
+	} else if masterEngine, err := GetMasterEngine(ctxEngine); err == nil {
+		masterEngine.ShowSQL(on)
 	}
 }
 
@@ -374,3 +423,18 @@ func (h *ErrorQueryHook) AfterProcess(c *contexts.ContextHook) error {
 	}
 	return nil
 }
+
+// GetMasterEngine extracts the master xorm.Engine from the provided xorm.Engine.
+// This handles both direct xorm.Engine cases and engines that implement a Master() method.
+func GetMasterEngine(x Engine) (*xorm.Engine, error) {
+	if getter, ok := x.(interface{ Master() *xorm.Engine }); ok {
+		return getter.Master(), nil
+	}
+
+	engine, ok := x.(*xorm.Engine)
+	if !ok {
+		return nil, fmt.Errorf("unsupported engine type: %T", x)
+	}
+
+	return engine, nil
+}

models/db/index_test.go

@@ -33,10 +33,11 @@ func getCurrentResourceIndex(ctx context.Context, tableName string, groupID int6
 
 func TestSyncMaxResourceIndex(t *testing.T) {
 	require.NoError(t, unittest.PrepareTestDatabase())
-	xe := unittest.GetXORMEngine()
+	xe, err := unittest.GetXORMEngine()
+	require.NoError(t, err)
 	require.NoError(t, xe.Sync(&TestIndex{}))
 
-	err := db.SyncMaxResourceIndex(db.DefaultContext, "test_index", 10, 51)
+	err = db.SyncMaxResourceIndex(db.DefaultContext, "test_index", 10, 51)
 	require.NoError(t, err)
 
 	// sync new max index
@@ -88,7 +89,8 @@ func TestSyncMaxResourceIndex(t *testing.T) {
 
 func TestGetNextResourceIndex(t *testing.T) {
 	require.NoError(t, unittest.PrepareTestDatabase())
-	xe := unittest.GetXORMEngine()
+	xe, err := unittest.GetXORMEngine()
+	require.NoError(t, err)
 	require.NoError(t, xe.Sync(&TestIndex{}))
 
 	// create a new record

models/db/iterate_test.go

@@ -17,7 +17,8 @@ import (
 
 func TestIterate(t *testing.T) {
 	require.NoError(t, unittest.PrepareTestDatabase())
-	xe := unittest.GetXORMEngine()
+	xe, err := unittest.GetXORMEngine()
+	require.NoError(t, err)
 	require.NoError(t, xe.Sync(&repo_model.RepoUnit{}))
 
 	cnt, err := db.GetEngine(db.DefaultContext).Count(&repo_model.RepoUnit{})