You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
99 lines
2.3 KiB
99 lines
2.3 KiB
package db
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"ims/util/backoff"
|
|
|
|
"gorm.io/gorm"
|
|
"zestack.dev/env"
|
|
)
|
|
|
|
type mysqlSchemaHelper struct{}
|
|
|
|
func (s *mysqlSchemaHelper) PublicSchema() string {
|
|
return env.String("DB_NAME")
|
|
}
|
|
|
|
func (s *mysqlSchemaHelper) TenantSchema(tenantId uint) string {
|
|
if tenantId == 0 {
|
|
return s.PublicSchema()
|
|
}
|
|
return fmt.Sprintf(
|
|
"%s%d%s",
|
|
env.String("DB_TENANT_PREFIX", "tenant_"),
|
|
tenantId,
|
|
env.String("DB_TENANT_SUFFIX", ""),
|
|
)
|
|
}
|
|
|
|
func (s *mysqlSchemaHelper) CurrentSchema(tx *gorm.DB) string {
|
|
var schema string
|
|
tx.Raw("SELECT DATABASE()").Row().Scan(&schema)
|
|
return schema
|
|
}
|
|
|
|
func (s *mysqlSchemaHelper) UseSchema(tx *gorm.DB, schema string) (func() error, error) {
|
|
if schema == "" {
|
|
err := errors.New("schema name is empty")
|
|
tx.AddError(err)
|
|
return nil, err
|
|
}
|
|
|
|
currentSchema := s.CurrentSchema(tx)
|
|
publicSchema := s.PublicSchema()
|
|
|
|
// 当前 schema 与目标 schema 相同，无需切换
|
|
if schema == currentSchema {
|
|
reset := func() error { return nil }
|
|
return reset, nil
|
|
}
|
|
|
|
// 不支持租户切换
|
|
if currentSchema != publicSchema && schema != publicSchema {
|
|
err := fmt.Errorf(
|
|
"failed to switch schema %s from current schema %s: %w",
|
|
schema,
|
|
currentSchema,
|
|
ErrSwitchSchema,
|
|
)
|
|
tx.AddError(err)
|
|
return nil, err
|
|
}
|
|
|
|
sqlstr := "USE " + tx.Statement.Quote(schema)
|
|
if execErr := tx.Exec(sqlstr).Error; execErr != nil {
|
|
err := fmt.Errorf("failed to set database %q: %w", schema, execErr)
|
|
tx.AddError(err)
|
|
return nil, err
|
|
}
|
|
|
|
if schema == publicSchema {
|
|
reset := func() error { return nil }
|
|
return reset, nil
|
|
}
|
|
|
|
reset := func() error {
|
|
return tx.Exec("USE " + tx.Statement.Quote(publicSchema)).Error
|
|
}
|
|
|
|
return reset, nil
|
|
}
|
|
|
|
func (s *mysqlSchemaHelper) LockSchema(tx *gorm.DB, schema string, retry *backoff.Options) (func() error, error) {
|
|
return AcquireLock(tx, schema, retry)
|
|
}
|
|
|
|
func (s *mysqlSchemaHelper) CreateSchema(tx *gorm.DB, schema string) error {
|
|
sql := fmt.Sprintf("CREATE DATABASE IF NOT EXISTS %s", tx.Statement.Quote(schema))
|
|
return tx.Exec(sql).Error
|
|
}
|
|
|
|
func (s *mysqlSchemaHelper) DropSchema(tx *gorm.DB, schema string) error {
|
|
sql := fmt.Sprintf("DROP DATABASE IF EXISTS %s", tx.Statement.Quote(schema))
|
|
return tx.Exec(sql).Error
|
|
}
|
|
|
|
func (s *mysqlSchemaHelper) String() string {
|
|
return "mysql"
|
|
}
|
|
|