dendrite/internal/sqlutil/migrate.go

package sqlutil

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

	"github.com/matrix-org/dendrite/internal/config"
	"github.com/pressly/goose"
)

type Migrations struct {
	registeredGoMigrations map[int64]*goose.Migration
}

func NewMigrations() *Migrations {
	return &Migrations{
		registeredGoMigrations: make(map[int64]*goose.Migration),
	}
}

// Copy-pasted from goose directly to store migrations into a map we control

// AddMigration adds a migration.
func (m *Migrations) AddMigration(up func(*sql.Tx) error, down func(*sql.Tx) error) {
	_, filename, _, _ := runtime.Caller(1)
	m.AddNamedMigration(filename, up, down)
}

// AddNamedMigration : Add a named migration.
func (m *Migrations) AddNamedMigration(filename string, up func(*sql.Tx) error, down func(*sql.Tx) error) {
	v, _ := goose.NumericComponent(filename)
	migration := &goose.Migration{Version: v, Next: -1, Previous: -1, Registered: true, UpFn: up, DownFn: down, Source: filename}

	if existing, ok := m.registeredGoMigrations[v]; ok {
		panic(fmt.Sprintf("failed to add migration %q: version conflicts with %q", filename, existing.Source))
	}

	m.registeredGoMigrations[v] = migration
}

// RunDeltas up to the latest version.
func (m *Migrations) RunDeltas(db *sql.DB, props *config.DatabaseOptions) error {
	maxVer := goose.MaxVersion
	minVer := int64(0)
	migrations, err := m.collect(minVer, maxVer)
	if err != nil {
		return fmt.Errorf("RunDeltas: Failed to collect migrations: %w", err)
	}
	if props.ConnectionString.IsPostgres() {
		if err = goose.SetDialect("postgres"); err != nil {
			return err
		}
	} else if props.ConnectionString.IsSQLite() {
		if err = goose.SetDialect("sqlite3"); err != nil {
			return err
		}
	} else {
		return fmt.Errorf("Unknown connection string: %s", props.ConnectionString)
	}
	for {
		current, err := goose.EnsureDBVersion(db)
		if err != nil {
			return fmt.Errorf("RunDeltas: Failed to EnsureDBVersion: %w", err)
		}

		next, err := migrations.Next(current)
		if err != nil {
			if err == goose.ErrNoNextVersion {
				return nil
			}

			return fmt.Errorf("RunDeltas: Failed to load next migration to %+v : %w", next, err)
		}

		if err = next.Up(db); err != nil {
			return fmt.Errorf("RunDeltas: Failed run migration: %w", err)
		}
	}
}

func (m *Migrations) collect(current, target int64) (goose.Migrations, error) {
	var migrations goose.Migrations

	// Go migrations registered via goose.AddMigration().
	for _, migration := range m.registeredGoMigrations {
		v, err := goose.NumericComponent(migration.Source)
		if err != nil {
			return nil, err
		}
		if versionFilter(v, current, target) {
			migrations = append(migrations, migration)
		}
	}

	migrations = sortAndConnectMigrations(migrations)

	return migrations, nil
}

func sortAndConnectMigrations(migrations goose.Migrations) goose.Migrations {
	sort.Sort(migrations)

	// now that we're sorted in the appropriate direction,
	// populate next and previous for each migration
	for i, m := range migrations {
		prev := int64(-1)
		if i > 0 {
			prev = migrations[i-1].Version
			migrations[i-1].Next = m.Version
		}
		migrations[i].Previous = prev
	}

	return migrations
}

func versionFilter(v, current, target int64) bool {

	if target > current {
		return v > current && v <= target
	}

	if target < current {
		return v <= current && v > target
	}

	return false
}