diff options
Diffstat (limited to 'vendor/github.com/uptrace/bun/migrate/auto.go')
-rw-r--r-- | vendor/github.com/uptrace/bun/migrate/auto.go | 429 |
1 files changed, 429 insertions, 0 deletions
diff --git a/vendor/github.com/uptrace/bun/migrate/auto.go b/vendor/github.com/uptrace/bun/migrate/auto.go new file mode 100644 index 000000000..e56fa23a0 --- /dev/null +++ b/vendor/github.com/uptrace/bun/migrate/auto.go @@ -0,0 +1,429 @@ +package migrate + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "os" + "path/filepath" + + "github.com/uptrace/bun" + "github.com/uptrace/bun/internal" + "github.com/uptrace/bun/migrate/sqlschema" + "github.com/uptrace/bun/schema" +) + +type AutoMigratorOption func(m *AutoMigrator) + +// WithModel adds a bun.Model to the scope of migrations. +func WithModel(models ...interface{}) AutoMigratorOption { + return func(m *AutoMigrator) { + m.includeModels = append(m.includeModels, models...) + } +} + +// WithExcludeTable tells the AutoMigrator to ignore a table in the database. +// This prevents AutoMigrator from dropping tables which may exist in the schema +// but which are not used by the application. +// +// Do not exclude tables included via WithModel, as BunModelInspector ignores this setting. +func WithExcludeTable(tables ...string) AutoMigratorOption { + return func(m *AutoMigrator) { + m.excludeTables = append(m.excludeTables, tables...) + } +} + +// WithSchemaName changes the default database schema to migrate objects in. +func WithSchemaName(schemaName string) AutoMigratorOption { + return func(m *AutoMigrator) { + m.schemaName = schemaName + } +} + +// WithTableNameAuto overrides default migrations table name. +func WithTableNameAuto(table string) AutoMigratorOption { + return func(m *AutoMigrator) { + m.table = table + m.migratorOpts = append(m.migratorOpts, WithTableName(table)) + } +} + +// WithLocksTableNameAuto overrides default migration locks table name. +func WithLocksTableNameAuto(table string) AutoMigratorOption { + return func(m *AutoMigrator) { + m.locksTable = table + m.migratorOpts = append(m.migratorOpts, WithLocksTableName(table)) + } +} + +// WithMarkAppliedOnSuccessAuto sets the migrator to only mark migrations as applied/unapplied +// when their up/down is successful. +func WithMarkAppliedOnSuccessAuto(enabled bool) AutoMigratorOption { + return func(m *AutoMigrator) { + m.migratorOpts = append(m.migratorOpts, WithMarkAppliedOnSuccess(enabled)) + } +} + +// WithMigrationsDirectoryAuto overrides the default directory for migration files. +func WithMigrationsDirectoryAuto(directory string) AutoMigratorOption { + return func(m *AutoMigrator) { + m.migrationsOpts = append(m.migrationsOpts, WithMigrationsDirectory(directory)) + } +} + +// AutoMigrator performs automated schema migrations. +// +// It is designed to be a drop-in replacement for some Migrator functionality and supports all existing +// configuration options. +// Similarly to Migrator, it has methods to create SQL migrations, write them to a file, and apply them. +// Unlike Migrator, it detects the differences between the state defined by bun models and the current +// database schema automatically. +// +// Usage: +// 1. Generate migrations and apply them au once with AutoMigrator.Migrate(). +// 2. Create up- and down-SQL migration files and apply migrations using Migrator.Migrate(). +// +// While both methods produce complete, reversible migrations (with entries in the database +// and SQL migration files), prefer creating migrations and applying them separately for +// any non-trivial cases to ensure AutoMigrator detects expected changes correctly. +// +// Limitations: +// - AutoMigrator only supports a subset of the possible ALTER TABLE modifications. +// - Some changes are not automatically reversible. For example, you would need to manually +// add a CREATE TABLE query to the .down migration file to revert a DROP TABLE migration. +// - Does not validate most dialect-specific constraints. For example, when changing column +// data type, make sure the data con be auto-casted to the new type. +// - Due to how the schema-state diff is calculated, it is not possible to rename a table and +// modify any of its columns' _data type_ in a single run. This will cause the AutoMigrator +// to drop and re-create the table under a different name; it is better to apply this change in 2 steps. +// Renaming a table and renaming its columns at the same time is possible. +// - Renaming table/column to an existing name, i.e. like this [A->B] [B->C], is not possible due to how +// AutoMigrator distinguishes "rename" and "unchanged" columns. +// +// Dialect must implement both sqlschema.Inspector and sqlschema.Migrator to be used with AutoMigrator. +type AutoMigrator struct { + db *bun.DB + + // dbInspector creates the current state for the target database. + dbInspector sqlschema.Inspector + + // modelInspector creates the desired state based on the model definitions. + modelInspector sqlschema.Inspector + + // dbMigrator executes ALTER TABLE queries. + dbMigrator sqlschema.Migrator + + table string // Migrations table (excluded from database inspection) + locksTable string // Migration locks table (excluded from database inspection) + + // schemaName is the database schema considered for migration. + schemaName string + + // includeModels define the migration scope. + includeModels []interface{} + + // excludeTables are excluded from database inspection. + excludeTables []string + + // diffOpts are passed to detector constructor. + diffOpts []diffOption + + // migratorOpts are passed to Migrator constructor. + migratorOpts []MigratorOption + + // migrationsOpts are passed to Migrations constructor. + migrationsOpts []MigrationsOption +} + +func NewAutoMigrator(db *bun.DB, opts ...AutoMigratorOption) (*AutoMigrator, error) { + am := &AutoMigrator{ + db: db, + table: defaultTable, + locksTable: defaultLocksTable, + schemaName: db.Dialect().DefaultSchema(), + } + + for _, opt := range opts { + opt(am) + } + am.excludeTables = append(am.excludeTables, am.table, am.locksTable) + + dbInspector, err := sqlschema.NewInspector(db, sqlschema.WithSchemaName(am.schemaName), sqlschema.WithExcludeTables(am.excludeTables...)) + if err != nil { + return nil, err + } + am.dbInspector = dbInspector + am.diffOpts = append(am.diffOpts, withCompareTypeFunc(db.Dialect().(sqlschema.InspectorDialect).CompareType)) + + dbMigrator, err := sqlschema.NewMigrator(db, am.schemaName) + if err != nil { + return nil, err + } + am.dbMigrator = dbMigrator + + tables := schema.NewTables(db.Dialect()) + tables.Register(am.includeModels...) + am.modelInspector = sqlschema.NewBunModelInspector(tables, sqlschema.WithSchemaName(am.schemaName)) + + return am, nil +} + +func (am *AutoMigrator) plan(ctx context.Context) (*changeset, error) { + var err error + + got, err := am.dbInspector.Inspect(ctx) + if err != nil { + return nil, err + } + + want, err := am.modelInspector.Inspect(ctx) + if err != nil { + return nil, err + } + + changes := diff(got, want, am.diffOpts...) + if err := changes.ResolveDependencies(); err != nil { + return nil, fmt.Errorf("plan migrations: %w", err) + } + return changes, nil +} + +// Migrate writes required changes to a new migration file and runs the migration. +// This will create and entry in the migrations table, making it possible to revert +// the changes with Migrator.Rollback(). MigrationOptions are passed on to Migrator.Migrate(). +func (am *AutoMigrator) Migrate(ctx context.Context, opts ...MigrationOption) (*MigrationGroup, error) { + migrations, _, err := am.createSQLMigrations(ctx, false) + if err != nil { + return nil, fmt.Errorf("auto migrate: %w", err) + } + + migrator := NewMigrator(am.db, migrations, am.migratorOpts...) + if err := migrator.Init(ctx); err != nil { + return nil, fmt.Errorf("auto migrate: %w", err) + } + + group, err := migrator.Migrate(ctx, opts...) + if err != nil { + return nil, fmt.Errorf("auto migrate: %w", err) + } + return group, nil +} + +// CreateSQLMigration writes required changes to a new migration file. +// Use migrate.Migrator to apply the generated migrations. +func (am *AutoMigrator) CreateSQLMigrations(ctx context.Context) ([]*MigrationFile, error) { + _, files, err := am.createSQLMigrations(ctx, true) + return files, err +} + +// CreateTxSQLMigration writes required changes to a new migration file making sure they will be executed +// in a transaction when applied. Use migrate.Migrator to apply the generated migrations. +func (am *AutoMigrator) CreateTxSQLMigrations(ctx context.Context) ([]*MigrationFile, error) { + _, files, err := am.createSQLMigrations(ctx, false) + return files, err +} + +func (am *AutoMigrator) createSQLMigrations(ctx context.Context, transactional bool) (*Migrations, []*MigrationFile, error) { + changes, err := am.plan(ctx) + if err != nil { + return nil, nil, fmt.Errorf("create sql migrations: %w", err) + } + + name, _ := genMigrationName(am.schemaName + "_auto") + migrations := NewMigrations(am.migrationsOpts...) + migrations.Add(Migration{ + Name: name, + Up: changes.Up(am.dbMigrator), + Down: changes.Down(am.dbMigrator), + Comment: "Changes detected by bun.AutoMigrator", + }) + + // Append .tx.up.sql or .up.sql to migration name, dependin if it should be transactional. + fname := func(direction string) string { + return name + map[bool]string{true: ".tx.", false: "."}[transactional] + direction + ".sql" + } + + up, err := am.createSQL(ctx, migrations, fname("up"), changes, transactional) + if err != nil { + return nil, nil, fmt.Errorf("create sql migration up: %w", err) + } + + down, err := am.createSQL(ctx, migrations, fname("down"), changes.GetReverse(), transactional) + if err != nil { + return nil, nil, fmt.Errorf("create sql migration down: %w", err) + } + return migrations, []*MigrationFile{up, down}, nil +} + +func (am *AutoMigrator) createSQL(_ context.Context, migrations *Migrations, fname string, changes *changeset, transactional bool) (*MigrationFile, error) { + var buf bytes.Buffer + + if transactional { + buf.WriteString("SET statement_timeout = 0;") + } + + if err := changes.WriteTo(&buf, am.dbMigrator); err != nil { + return nil, err + } + content := buf.Bytes() + + fpath := filepath.Join(migrations.getDirectory(), fname) + if err := os.WriteFile(fpath, content, 0o644); err != nil { + return nil, err + } + + mf := &MigrationFile{ + Name: fname, + Path: fpath, + Content: string(content), + } + return mf, nil +} + +// Func creates a MigrationFunc that applies all operations all the changeset. +func (c *changeset) Func(m sqlschema.Migrator) MigrationFunc { + return func(ctx context.Context, db *bun.DB) error { + return c.apply(ctx, db, m) + } +} + +// GetReverse returns a new changeset with each operation in it "reversed" and in reverse order. +func (c *changeset) GetReverse() *changeset { + var reverse changeset + for i := len(c.operations) - 1; i >= 0; i-- { + reverse.Add(c.operations[i].GetReverse()) + } + return &reverse +} + +// Up is syntactic sugar. +func (c *changeset) Up(m sqlschema.Migrator) MigrationFunc { + return c.Func(m) +} + +// Down is syntactic sugar. +func (c *changeset) Down(m sqlschema.Migrator) MigrationFunc { + return c.GetReverse().Func(m) +} + +// apply generates SQL for each operation and executes it. +func (c *changeset) apply(ctx context.Context, db *bun.DB, m sqlschema.Migrator) error { + if len(c.operations) == 0 { + return nil + } + + for _, op := range c.operations { + if _, isComment := op.(*comment); isComment { + continue + } + + b := internal.MakeQueryBytes() + b, err := m.AppendSQL(b, op) + if err != nil { + return fmt.Errorf("apply changes: %w", err) + } + + query := internal.String(b) + if _, err = db.ExecContext(ctx, query); err != nil { + return fmt.Errorf("apply changes: %w", err) + } + } + return nil +} + +func (c *changeset) WriteTo(w io.Writer, m sqlschema.Migrator) error { + var err error + + b := internal.MakeQueryBytes() + for _, op := range c.operations { + if c, isComment := op.(*comment); isComment { + b = append(b, "/*\n"...) + b = append(b, *c...) + b = append(b, "\n*/"...) + continue + } + + b, err = m.AppendSQL(b, op) + if err != nil { + return fmt.Errorf("write changeset: %w", err) + } + b = append(b, ";\n"...) + } + if _, err := w.Write(b); err != nil { + return fmt.Errorf("write changeset: %w", err) + } + return nil +} + +func (c *changeset) ResolveDependencies() error { + if len(c.operations) <= 1 { + return nil + } + + const ( + unvisited = iota + current + visited + ) + + status := make(map[Operation]int, len(c.operations)) + for _, op := range c.operations { + status[op] = unvisited + } + + var resolved []Operation + var nextOp Operation + var visit func(op Operation) error + + next := func() bool { + for op, s := range status { + if s == unvisited { + nextOp = op + return true + } + } + return false + } + + // visit iterates over c.operations until it finds all operations that depend on the current one + // or runs into cirtular dependency, in which case it will return an error. + visit = func(op Operation) error { + switch status[op] { + case visited: + return nil + case current: + // TODO: add details (circle) to the error message + return errors.New("detected circular dependency") + } + + status[op] = current + + for _, another := range c.operations { + if dop, hasDeps := another.(interface { + DependsOn(Operation) bool + }); another == op || !hasDeps || !dop.DependsOn(op) { + continue + } + if err := visit(another); err != nil { + return err + } + } + + status[op] = visited + + // Any dependent nodes would've already been added to the list by now, so we prepend. + resolved = append([]Operation{op}, resolved...) + return nil + } + + for next() { + if err := visit(nextOp); err != nil { + return err + } + } + + c.operations = resolved + return nil +} |