github.com/status-im/status-go@v1.1.0/sqlite/migrate.go

package sqlite

import (
	"database/sql"
	"fmt"
	"sort"

	"github.com/status-im/migrate/v4"
	"github.com/status-im/migrate/v4/database/sqlcipher"
	bindata "github.com/status-im/migrate/v4/source/go_bindata"
)

type CustomMigrationFunc func(tx *sql.Tx) error

type PostStep struct {
	Version         uint
	CustomMigration CustomMigrationFunc
	RollBackVersion uint
}

var migrationTable = "status_go_" + sqlcipher.DefaultMigrationsTable

// Migrate database with option to augment the migration steps with additional processing using the customSteps
// parameter. For each PostStep entry in customSteps the CustomMigration will be called after the migration step
// with the matching Version number has been executed. If the CustomMigration returns an error, the migration process
// is aborted. In case the custom step failures the migrations are run down to RollBackVersion if > 0.
//
// The recommended way to create a custom migration is by providing empty and versioned run/down sql files as markers.
// Then running all the SQL code inside the same transaction to transform and commit provides the possibility
// to completely rollback the migration in case of failure, avoiding to leave the DB in an inconsistent state.
//
// Marker migrations can be created by using PostStep structs with specific Version numbers and a callback function,
// even when no accompanying SQL migration is needed. This can be used to trigger Go code at specific points
// during the migration process.
//
// Caution: This mechanism should be used as a last resort. Prefer data migration using SQL migration files
// whenever possible to ensure consistency and compatibility with standard migration tools.
//
// untilVersion, for testing purposes optional parameter, can be used to limit the migration to a specific version.
// Pass nil to migrate to the latest available version.
func Migrate(db *sql.DB, resources *bindata.AssetSource, customSteps []*PostStep, untilVersion *uint) error {
	source, err := bindata.WithInstance(resources)
	if err != nil {
		return fmt.Errorf("failed to create bindata migration source: %w", err)
	}

	driver, err := sqlcipher.WithInstance(db, &sqlcipher.Config{
		MigrationsTable: migrationTable,
	})
	if err != nil {
		return fmt.Errorf("failed to create sqlcipher driver: %w", err)
	}

	m, err := migrate.NewWithInstance("go-bindata", source, "sqlcipher", driver)
	if err != nil {
		return fmt.Errorf("failed to create migration instance: %w", err)
	}

	if len(customSteps) == 0 {
		return runRemainingMigrations(m, untilVersion)
	}

	sort.Slice(customSteps, func(i, j int) bool {
		return customSteps[i].Version < customSteps[j].Version
	})

	lastVersion, err := getCurrentVersion(m, db)
	if err != nil {
		return err
	}

	customIndex := 0
	// ignore processed versions
	for customIndex < len(customSteps) && customSteps[customIndex].Version <= lastVersion {
		customIndex++
	}

	if err := runCustomMigrations(m, db, customSteps, customIndex, untilVersion); err != nil {
		return err
	}

	return runRemainingMigrations(m, untilVersion)
}

// runCustomMigrations performs source migrations from current to each custom steps, then runs custom migration callback
// until it executes all custom migrations or an error occurs and it tries to rollback to RollBackVersion if > 0.
func runCustomMigrations(m *migrate.Migrate, db *sql.DB, customSteps []*PostStep, customIndex int, untilVersion *uint) error {
	for customIndex < len(customSteps) && (untilVersion == nil || customSteps[customIndex].Version <= *untilVersion) {
		customStep := customSteps[customIndex]

		if err := m.Migrate(customStep.Version); err != nil && err != migrate.ErrNoChange {
			return fmt.Errorf("failed to migrate to version %d: %w", customStep.Version, err)
		}

		if err := runCustomMigrationStep(db, customStep, m); err != nil {
			return err
		}

		customIndex++
	}
	return nil
}

func runCustomMigrationStep(db *sql.DB, customStep *PostStep, m *migrate.Migrate) error {

	sqlTx, err := db.Begin()
	if err != nil {
		return fmt.Errorf("failed to begin transaction: %w", err)
	}

	if err := customStep.CustomMigration(sqlTx); err != nil {
		_ = sqlTx.Rollback()
		return rollbackCustomMigration(m, customStep, err)
	}

	if err := sqlTx.Commit(); err != nil {
		return fmt.Errorf("failed to commit transaction: %w", err)
	}
	return nil
}

func rollbackCustomMigration(m *migrate.Migrate, customStep *PostStep, customErr error) error {
	if customStep.RollBackVersion > 0 {
		err := m.Migrate(customStep.RollBackVersion)
		newV, _, _ := m.Version()
		if err != nil {
			return fmt.Errorf("failed to rollback migration to version %d: %w", customStep.RollBackVersion, err)
		}
		return fmt.Errorf("custom migration step failed for version %d. Successfully rolled back migration to version %d: %w", customStep.Version, newV, customErr)
	}
	return fmt.Errorf("custom migration step failed for version %d: %w", customStep.Version, customErr)
}

func runRemainingMigrations(m *migrate.Migrate, untilVersion *uint) error {
	if untilVersion != nil {
		if err := m.Migrate(*untilVersion); err != nil && err != migrate.ErrNoChange {
			return fmt.Errorf("failed to migrate to version %d: %w", *untilVersion, err)
		}
	} else {
		if err := m.Up(); err != nil && err != migrate.ErrNoChange {
			ver, _, _ := m.Version()
			return fmt.Errorf("failed to migrate up: %w, current version: %d", err, ver)
		}
	}
	return nil
}

func getCurrentVersion(m *migrate.Migrate, db *sql.DB) (uint, error) {
	lastVersion, dirty, err := m.Version()
	if err != nil && err != migrate.ErrNilVersion {
		return 0, fmt.Errorf("failed to get migration version: %w", err)
	}
	if dirty {
		return 0, fmt.Errorf("DB is dirty after migration version %d", lastVersion)
	}
	if err == migrate.ErrNilVersion {
		lastVersion, _, err = GetLastMigrationVersion(db)
		return lastVersion, err
	}
	return lastVersion, nil
}

// GetLastMigrationVersion returns the last migration version stored in the migration table.
// Returns 0 for version in case migrationTableExists is true
func GetLastMigrationVersion(db *sql.DB) (version uint, migrationTableExists bool, err error) {
	// Check if the migration table exists
	row := db.QueryRow("SELECT exists(SELECT name FROM sqlite_master WHERE type='table' AND name=?)", migrationTable)
	migrationTableExists = false
	err = row.Scan(&migrationTableExists)
	if err != nil && err != sql.ErrNoRows {
		return 0, false, err
	}

	var lastMigration uint64 = 0
	if migrationTableExists {
		row = db.QueryRow("SELECT version FROM status_go_schema_migrations")
		err = row.Scan(&lastMigration)
		if err != nil && err != sql.ErrNoRows {
			return 0, true, err
		}
	}
	return uint(lastMigration), migrationTableExists, nil
}