lychiyu/bbgo_origin
1
0
Fork 0
mirror of https://github.com/c9s/bbgo.git synced 2024-11-21 22:43:52 +00:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

support extra migration packages

Browse Source
This commit is contained in:
c9s 2024-01-24 15:18:31 +08:00
parent ea70544b38
commit 59713fa532
No known key found for this signature in database
GPG Key ID: 7385E7E464CB0A54
10 changed files with 56 additions and 37 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
go.mod
View File

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

2
go.sum
View File

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

4
pkg/bbgo/bootstrap.go
View File

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

9
pkg/bbgo/config.go
View File

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

46
pkg/bbgo/environment.go
View File

@ -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 {
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)
} else if dsn, ok := os.LookupEnv("MYSQL_URL"); ok {
return environ.ConfigureDatabaseDriver(ctx, "mysql", dsn)
var dbDriver string
var dbDSN string
var extraPkgNames []string
if config != nil && config.DatabaseConfig != nil {
dbDriver = config.DatabaseConfig.Driver
dbDSN = config.DatabaseConfig.DSN
extraPkgNames = config.DatabaseConfig.ExtraMigrationPackages
}
return nil
if val, ok := os.LookupEnv("DB_DRIVER"); ok {
dbDriver = val
}
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 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
}

2
pkg/cmd/account.go
View File

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

2
pkg/cmd/market.go
View File

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

2
pkg/cmd/pnl.go
View File

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

2
pkg/cmd/sync.go
View File

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

22
pkg/service/database.go
View File

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