bumps uptrace/bun dependencies to v1.2.6 (#3569)

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