131 lines
3.3 KiB
Go
131 lines
3.3 KiB
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
|
|
}
|