status-im/status-go

View on GitHub
sqlite/migrate.go

Summary

Maintainability
A
0 mins
Test Coverage
D
69%
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
}