migrate/migrate.go

package migrate

import (
	"context"
	"io/fs"
	"log/slog"
	"time"

	"gitea.auvem.com/go-toolkit/app"
	"gitea.auvem.com/go-toolkit/dbx"
	"github.com/pressly/goose/v3"
)

// MigrationsOpts define the options for the migrations module.
type MigrationOpts struct {
	// SQLO is the SQL database handle getter used for migrations. REQUIRED.
	SQLO dbx.SQLOFunc

	// Dialect is the database dialect used for migrations (e.g., "mysql", "postgres").
	// REQUIRED. Must match the dialect used in dbx.
	Dialect goose.Dialect

	// FS is the filesystem where migration files are stored.
	FS fs.FS

	// BasePath is the directory path where migration files are located.
	// Defaults to "." if not set.
	BasePath string
}

const (
	// ModuleMigrationsName is the name of the migrations module.
	ModuleMigrationsName = "migrations"

	// ModuleMigrateUpName is the name of the migrate up module.
	ModuleMigrateUpName = "migrate up"

	// ModuleMigrateBlankName is the name of the migrate blank module.
	ModuleMigrateBlankName = "migrate down to blank"

	// ModuleAutoMigrateName is the name of the auto-migrate module.
	ModuleAutoMigrateName = "auto migrate"
)

var (
	// Migration stores the goose migration provider.
	Migration *goose.Provider

	// migrationsConfig stores the configuration for the migrations module.
	migrationsConfig MigrationOpts

	// migrationsModule is the singleton module instance for the migrations subsystem.
	migrationsModule *app.Module

	// ModuleMigrateUp applies any pending migrations.
	ModuleMigrateUp = app.NewModule(ModuleMigrateUpName, app.ModuleOpts{
		Setup: func(_ *app.Module) error {
			_, err := ApplyPendingMigrations(context.Background(), -1)
			return err
		},
		Depends: []string{ModuleMigrationsName},
	})

	// moduleMigrateBlank resets the database to a blank state, removing all data.
	moduleMigrateBlank *app.Module

	// autoMigrateEnabled controls whether auto-migration is enabled.
	autoMigrateEnabled bool

	// autoMigrateModule is the singleton module instance for the auto-migration subsystem.
	autoMigrateModule *app.Module
)

// MigrationsConfig returns the current configuration for the migrations module.
func MigrationsConfig() MigrationOpts {
	migrationsModule.RequireLoaded() // ensure the migrations module is loaded
	return migrationsConfig
}

// ModuleMigrations returns the migrations module with the provided configuration.
func ModuleMigrations(cfg MigrationOpts) *app.Module {
	if migrationsModule != nil {
		panic("ModuleMigrations initialized multiple times")
	}

	if cfg.SQLO == nil {
		panic("Migration SQL handle (SQLO) must be set in the configuration")
	}
	if cfg.BasePath == "" {
		cfg.BasePath = "." // default base path if not set
	}
	if cfg.FS == nil {
		panic("Migration filesystem (FS) must be set in the configuration")
	}

	migrationsConfig = cfg // store configuration at package level

	migrationsModule = app.NewModule(ModuleMigrationsName, app.ModuleOpts{
		Setup:   setupMigrations,
		Depends: []string{dbx.ModuleDBName},
	})

	return migrationsModule
}

// ModuleMigrateBlank returns the migrate blank module that resets the database to a blank state.
func ModuleMigrateBlank() *app.Module {
	if moduleMigrateBlank != nil {
		panic("ModuleMigrateBlank initialized multiple times")
	}

	moduleMigrateBlank = app.NewModule(ModuleMigrateBlankName, app.ModuleOpts{
		Setup: func(_ *app.Module) error {
			_, err := MigrateToBlank()
			return err
		},
		Depends: []string{ModuleMigrationsName},
	})

	return moduleMigrateBlank
}

// ModuleAutoMigrate returns the auto-migration module with the provided configuration.
func ModuleAutoMigrate(enabled bool) *app.Module {
	if autoMigrateModule != nil {
		panic("ModuleAutoMigrate initialized multiple times")
	}

	autoMigrateEnabled = enabled // store auto-migration state at package level

	autoMigrateModule = app.NewModule(ModuleAutoMigrateName, app.ModuleOpts{
		Setup: func(_ *app.Module) error {
			return AutoMigrate()
		},
		Depends: []string{ModuleMigrationsName},
	})

	return autoMigrateModule
}

// setupMigrations initializes the goose migration provider.
func setupMigrations(_ *app.Module) error {
	var err error

	if err := goose.SetDialect(string(migrationsConfig.Dialect)); err != nil {
		migrationsModule.Logger().Error("Couldn't set database dialect for goose", "err", err)
		return err
	}

	// Set base filesystem for goose migrations
	goose.SetBaseFS(migrationsConfig.FS)

	// Initialize the goose migration provider
	Migration, err = goose.NewProvider(migrationsConfig.Dialect, migrationsConfig.SQLO(), migrationsConfig.FS)
	if err != nil {
		migrationsModule.Logger().Error("Couldn't initialize goose migration provider", "err", err)
		return err
	}

	return nil
}

// ApplyPendingMigrations applies all pending migrations. Returns the number of
// migrations applied and an error if any occurred. The number of migrations that
// will be applied are specified by `pendingCount`. If `pendingCount` is 0, no migrations
// are applied. If `pendingCount` is negative, the number of pending migrations
// is fetched from the database.
func ApplyPendingMigrations(ctx context.Context, pendingCount int64) (int64, error) {
	migrationsModule.RequireLoaded() // ensure the migrations module is loaded

	if pendingCount == 0 {
		return 0, nil
	}

	if pendingCount < 0 {
		curr, target, err := Migration.GetVersions(ctx)
		if err != nil {
			return 0, err
		}
		pendingCount = target - curr
	}

	var count int64
	for range pendingCount {
		res, err := Migration.UpByOne(ctx)
		if err := handleMigrationResults(migrationsModule.Logger(), res, err); err != nil {
			return count, err
		}

		if res.Error == nil {
			count++
		}
	}
	return count, nil
}

// MigrateToBlank resets the database to a blank state, removing all data and
// running all down migrations. Returns number of migrations applied and an error
// if any occurred.
func MigrateToBlank() (int64, error) {
	migrationsModule.RequireLoaded() // ensure the migrations module is loaded

	ctx := context.Background()
	current, target, err := Migration.GetVersions(ctx)
	if err != nil {
		return 0, err
	}
	moduleMigrateBlank.Logger().Info("Database versions", "current", current, "target", target)

	var count int64
	for current > 0 {
		res, err := Migration.Down(ctx)
		if err := handleMigrationResults(moduleMigrateBlank.Logger(), res, err); err != nil {
			return count, err
		}

		count++
		current--
	}
	return count, nil
}

// AutoMigrate applies any pending migrations if auto-migration is enabled.
func AutoMigrate() error {
	migrationsModule.RequireLoaded() // ensure the migrations module is loaded

	// Check if there are any pending migrations
	migrationCtx := context.Background()
	migrationCurrent, migrationTarget, err := Migration.GetVersions(migrationCtx)
	if err != nil {
		migrationsModule.Logger().Error("Couldn't check for pending migrations", "err", err)
		return err
	}

	migrationFields := []any{
		"current", migrationCurrent,
		"target", migrationTarget,
	}

	if migrationCurrent >= migrationTarget {
		migrationsModule.Logger().Info("No pending migrations", "version", migrationCurrent)
	} else if !autoMigrateEnabled {
		migrationsModule.Logger().Error(
			"Pending migrations detected, but auto-migration is disabled. Please run `acrm migrate up` to apply them.",
			migrationFields...,
		)
		return err
	} else {
		migrationsModule.Logger().Info("Pending migrations detected, applying them...", migrationFields...)
		now := time.Now()
		count, err := ApplyPendingMigrations(migrationCtx, migrationTarget-migrationCurrent)
		if err != nil {
			migrationsModule.Logger().Error("Couldn't apply pending migrations", "current", migrationCurrent+count, "target", migrationTarget, "err", err)
			return err
		}
		migrationsModule.Logger().Info("Applied pending migrations", "current", migrationTarget, "appliedCount", count, "duration", time.Since(now))
	}

	return nil
}

// handleMigrationResults is a helper function that prints various responses
// based on a *goose.MigrationResult.
func handleMigrationResults(logger *slog.Logger, res *goose.MigrationResult, err error) error {
	if err != nil {
		return err
	}

	fields := []any{
		"dir", res.Direction,
		"version", res.Source.Version,
		"source", res.Source.Path,
		"duration", res.Duration,
	}

	if res.Error != nil {
		fields = append(fields, "err", res.Error)
		logger.Error("Couldn't apply migration", fields...)
		return res.Error
	} else if res.Empty {
		logger.Warn("Applied empty migration", fields...)
	} else {
		logger.Info("Applied migration", fields...)
	}

	return nil
}