Merge pull request #1510 from c9s/c9s/feature/extra-migration-packages

feature: support extra migration packages

.github/workflows/go.yml

@@ -78,7 +78,7 @@ jobs:
 
     - name: Test
       run: |
-        go test -race -coverprofile coverage.txt -covermode atomic ./pkg/...
+        go test -count 3 -race -coverprofile coverage.txt -covermode atomic ./pkg/...
         sed -i -e '/_requestgen.go/d' coverage.txt
 
     - name: TestDnum

config/environment.yaml

@@ -0,0 +1,7 @@
+environment:
+  disableDefaultKLineSubscription: true
+  disableHistoryKLinePreload: true
+  disableStartupBalanceQuery: true
+  disableSessionTradeBuffer: true
+  disableMarketDataStore: true
+  maxSessionTradeBufferSize: true

go.mod

@@ -10,7 +10,7 @@ require (
 	github.com/adshao/go-binance/v2 v2.4.2
 	github.com/c-bata/goptuna v0.8.1
 	github.com/c9s/requestgen v1.3.6
-	github.com/c9s/rockhopper/v2 v2.0.2
+	github.com/c9s/rockhopper/v2 v2.0.3-0.20240124055428-2473c6221858
 	github.com/cenkalti/backoff/v4 v4.2.0
 	github.com/cheggaaa/pb/v3 v3.0.8
 	github.com/codingconcepts/env v0.0.0-20200821220118-a8fbf8d84482

go.sum

@@ -88,6 +88,8 @@ github.com/c9s/rockhopper/v2 v2.0.1 h1:GlqO5BJX1gA/jsP3OMpZwS4qLGm0CbpIjme+Ekkvz
 github.com/c9s/rockhopper/v2 v2.0.1/go.mod h1:QltQJDqEGppqrAhgXCG98Efm38UVQ0RdnJNEhbnTqYI=
 github.com/c9s/rockhopper/v2 v2.0.2 h1:ccOY4hHry2KbI8XPbFxy841tlCVYcYlhlfYwYGvTx94=
 github.com/c9s/rockhopper/v2 v2.0.2/go.mod h1:QltQJDqEGppqrAhgXCG98Efm38UVQ0RdnJNEhbnTqYI=
+github.com/c9s/rockhopper/v2 v2.0.3-0.20240124055428-2473c6221858 h1:VLvl82XQCtH12axHdroyrfprPYKtA3TueEWicx8yT4A=
+github.com/c9s/rockhopper/v2 v2.0.3-0.20240124055428-2473c6221858/go.mod h1:QltQJDqEGppqrAhgXCG98Efm38UVQ0RdnJNEhbnTqYI=
 github.com/cenkalti/backoff/v4 v4.2.0 h1:HN5dHm3WBOgndBH6E8V0q2jIYIR3s9yglV8k/+MN3u4=
 github.com/cenkalti/backoff/v4 v4.2.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE=
 github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=

pkg/bbgo/bootstrap.go

@@ -34,7 +34,7 @@ func BootstrapEnvironmentLightweight(ctx context.Context, environ *Environment,
 }
 
 func BootstrapEnvironment(ctx context.Context, environ *Environment, userConfig *Config) error {
-	if err := environ.ConfigureDatabase(ctx); err != nil {
+	if err := environ.ConfigureDatabase(ctx, userConfig); err != nil {
 		return err
 	}
 
@@ -66,5 +66,5 @@ func BootstrapEnvironment(ctx context.Context, environ *Environment, userConfig
 }
 
 func BootstrapBacktestEnvironment(ctx context.Context, environ *Environment) error {
-	return environ.ConfigureDatabase(ctx)
+	return environ.ConfigureDatabase(ctx, nil)
 }

pkg/bbgo/config.go

@@ -326,6 +326,13 @@ type ServiceConfig struct {
 	GoogleSpreadSheetService *GoogleSpreadSheetServiceConfig `json:"googleSpreadSheet" yaml:"googleSpreadSheet"`
 }
 
+type DatabaseConfig struct {
+	Driver string `json:"driver"`
+	DSN    string `json:"dsn"`
+
+	ExtraMigrationPackages []string `json:"extraMigrationPackages"`
+}
+
 type EnvironmentConfig struct {
 	DisableDefaultKLineSubscription bool `json:"disableDefaultKLineSubscription"`
 	DisableHistoryKLinePreload      bool `json:"disableHistoryKLinePreload"`
@@ -358,6 +365,8 @@ type Config struct {
 
 	Service *ServiceConfig `json:"services,omitempty" yaml:"services,omitempty"`
 
+	DatabaseConfig *DatabaseConfig `json:"database,omitempty" yaml:"database,omitempty"`
+
 	Environment *EnvironmentConfig `json:"environment,omitempty" yaml:"environment,omitempty"`
 
 	Sessions map[string]*ExchangeSession `json:"sessions,omitempty" yaml:"sessions,omitempty"`

pkg/bbgo/environment.go

@@ -162,34 +162,48 @@ func (environ *Environment) SelectSessions(names ...string) map[string]*Exchange
 	return sessions
 }
 
-func (environ *Environment) ConfigureDatabase(ctx context.Context) error {
+func (environ *Environment) ConfigureDatabase(ctx context.Context, config *Config) error {
 	// configureDB configures the database service based on the environment variable
-	if driver, ok := os.LookupEnv("DB_DRIVER"); ok {
+	var dbDriver string
+	var dbDSN string
+	var extraPkgNames []string
 
-		if dsn, ok := os.LookupEnv("DB_DSN"); ok {
-			return environ.ConfigureDatabaseDriver(ctx, driver, dsn)
-		}
-
-	} else if dsn, ok := os.LookupEnv("SQLITE3_DSN"); ok {
-
-		return environ.ConfigureDatabaseDriver(ctx, "sqlite3", dsn)
+	if config != nil && config.DatabaseConfig != nil {
+		dbDriver = config.DatabaseConfig.Driver
+		dbDSN = config.DatabaseConfig.DSN
+		extraPkgNames = config.DatabaseConfig.ExtraMigrationPackages
+	}
 
-	} else if dsn, ok := os.LookupEnv("MYSQL_URL"); ok {
+	if val, ok := os.LookupEnv("DB_DRIVER"); ok {
+		dbDriver = val
+	}
 
-		return environ.ConfigureDatabaseDriver(ctx, "mysql", dsn)
+	if val, ok := os.LookupEnv("DB_DSN"); ok {
+		dbDSN = val
+	} else if val, ok := os.LookupEnv("SQLITE3_DSN"); ok && (dbDriver == "" || dbDriver == "sqlite3") {
+		dbDSN = val
+		dbDriver = "sqlite3"
+	} else if val, ok := os.LookupEnv("MYSQL_URL"); ok && (dbDriver == "" || dbDriver == "mysql") {
+		dbDSN = val
+		dbDriver = "mysql"
+	}
 
+	if dbDriver == "" {
+		return fmt.Errorf("either env DB_DRIVER or config.Driver is not set")
 	}
 
-	return nil
+	return environ.ConfigureDatabaseDriver(ctx, dbDriver, dbDSN, extraPkgNames...)
 }
 
-func (environ *Environment) ConfigureDatabaseDriver(ctx context.Context, driver string, dsn string) error {
+func (environ *Environment) ConfigureDatabaseDriver(ctx context.Context, driver string, dsn string, extraPkgNames ...string) error {
 	environ.DatabaseService = service.NewDatabaseService(driver, dsn)
 	err := environ.DatabaseService.Connect()
 	if err != nil {
 		return err
 	}
 
+	environ.DatabaseService.AddMigrationPackages(extraPkgNames...)
+
 	if err := environ.DatabaseService.Upgrade(ctx); err != nil {
 		return err
 	}

pkg/cmd/account.go

@@ -38,7 +38,7 @@ var accountCmd = &cobra.Command{
 		}
 
 		environ := bbgo.NewEnvironment()
-		if err := environ.ConfigureDatabase(ctx); err != nil {
+		if err := environ.ConfigureDatabase(ctx, userConfig); err != nil {
 			return err
 		}
 

pkg/cmd/market.go

@@ -44,7 +44,7 @@ var marketCmd = &cobra.Command{
 		}
 
 		environ := bbgo.NewEnvironment()
-		if err := environ.ConfigureDatabase(ctx); err != nil {
+		if err := environ.ConfigureDatabase(ctx, userConfig); err != nil {
 			return err
 		}
 

pkg/cmd/pnl.go

@@ -88,7 +88,7 @@ var PnLCmd = &cobra.Command{
 
 		environ := bbgo.NewEnvironment()
 
-		if err := environ.ConfigureDatabase(ctx); err != nil {
+		if err := environ.ConfigureDatabase(ctx, userConfig); err != nil {
 			return err
 		}
 

pkg/cmd/sync.go

@@ -49,7 +49,7 @@ var SyncCmd = &cobra.Command{
 		}
 
 		environ := bbgo.NewEnvironment()
-		if err := environ.ConfigureDatabase(ctx); err != nil {
+		if err := environ.ConfigureDatabase(ctx, userConfig); err != nil {
 			return err
 		}
 

pkg/service/database.go

@@ -19,6 +19,8 @@ type DatabaseService struct {
 	Driver string
 	DSN    string
 	DB     *sqlx.DB
+
+	migrationPackages []string
 }
 
 func NewDatabaseService(driver, dsn string) *DatabaseService {
@@ -35,7 +37,6 @@ func NewDatabaseService(driver, dsn string) *DatabaseService {
 		Driver: driver,
 		DSN:    dsn,
 	}
-
 }
 
 func (s *DatabaseService) Connect() error {
@@ -50,6 +51,10 @@ func (s *DatabaseService) Insert(record interface{}) error {
 	return err
 }
 
+func (s *DatabaseService) AddMigrationPackages(pkgNames ...string) {
+	s.migrationPackages = append(s.migrationPackages, pkgNames...)
+}
+
 func (s *DatabaseService) Close() error {
 	return s.DB.Close()
 }
@@ -77,23 +82,12 @@ func (s *DatabaseService) Upgrade(ctx context.Context) error {
 		return err
 	}
 
-	migrations = migrations.FilterPackage([]string{"main"}).SortAndConnect()
 	if len(migrations) == 0 {
 		return nil
 	}
 
-	_, lastAppliedMigration, err := rh.FindLastAppliedMigration(ctx, migrations)
-	if err != nil {
-		return err
-	}
-
-	if lastAppliedMigration != nil {
-		return rockhopper.Up(ctx, rh, lastAppliedMigration.Next, 0)
-	}
-
-	// TODO: use align in the next major version
-	// return rockhopper.Align(ctx, rh, 20231123125402, migrations)
-	return rockhopper.Up(ctx, rh, migrations.Head(), 0)
+	pkgNames := append([]string{rockhopper.DefaultPackageName}, s.migrationPackages...)
+	return rockhopper.Upgrade(ctx, rh, migrations.FilterPackage(pkgNames))
 }
 
 func ReformatMysqlDSN(dsn string) (string, error) {