migration.go

package migration

import (
	"bytes"
	"fmt"
	"io"
	"regexp"
	"sort"
	"strconv"

	"github.com/Boostport/migration/parser"
)

// Direction type up/down
type Direction int

// String returns a string representation of the direction
func (d Direction) String() string {
	switch d {
	case Up:
		return "up"
	case Down:
		return "down"
	default:
		return "directionless"
	}
}

// Constants for direction
const (
	Up Direction = iota
	Down
)

var numberPrefixRegex = regexp.MustCompile(`^(\d+).*$`)

// Migration represents a migration, containing statements for migrating up and down.
type Migration struct {
	ID   string
	Up   *parser.ParsedMigration
	Down *parser.ParsedMigration
}

// PlannedMigration is a migration with a direction defined. This allows the driver to
// work out how to apply the migration.
type PlannedMigration struct {
	*Migration
	Direction Direction
}

// Less compares two migrations to determine how they should be ordered.
func (m Migration) Less(other *Migration) bool {
	switch {
	case m.isNumeric() && other.isNumeric() && m.VersionInt() != other.VersionInt():
		return m.VersionInt() < other.VersionInt()
	case m.isNumeric() && !other.isNumeric():
		return true
	case !m.isNumeric() && other.isNumeric():
		return false
	default:
		return m.ID < other.ID
	}
}

func (m Migration) isNumeric() bool {
	return len(m.NumberPrefixMatches()) > 0
}

// NumberPrefixMatches returns a list of string matches
func (m Migration) NumberPrefixMatches() []string {
	return numberPrefixRegex.FindStringSubmatch(m.ID)
}

// VersionInt converts the migration version to an 64-bit integer.
func (m Migration) VersionInt() int64 {
	v := m.NumberPrefixMatches()[1]
	value, err := strconv.ParseInt(v, 10, 64)
	if err != nil {
		panic(fmt.Sprintf("Could not parse %q into int64: %s", v, err))
	}
	return value
}

type byID []*Migration

func (b byID) Len() int           { return len(b) }
func (b byID) Swap(i, j int)      { b[i], b[j] = b[j], b[i] }
func (b byID) Less(i, j int) bool { return b[i].Less(b[j]) }

// Migrate runs a migration using a given driver and MigrationSource. The direction defines whether
// the migration is up or down, and max is the maximum number of migrations to apply. If max is set to 0,
// then there is no limit on the number of migrations to apply.
func Migrate(driver Driver, migrations Source, direction Direction, max int) (int, error) {
	count := 0

	m, err := getMigrations(migrations)
	if err != nil {
		return count, err
	}

	appliedMigrations, err := driver.Versions()
	if err != nil {
		return count, err
	}

	migrationsToApply := planMigrations(m, appliedMigrations, direction, max)
	for _, plannedMigration := range migrationsToApply {
		logPrintf("Applying migration (%s) named '%s'...", direction.String(), plannedMigration.ID)

		err = driver.Migrate(plannedMigration)
		if err != nil {
			errorMessage := "Error while running migration " + plannedMigration.ID

			if plannedMigration.Direction == Up {
				errorMessage += " (up)"
			} else {
				errorMessage += " (down)"
			}
			return count, fmt.Errorf(errorMessage+": %s", err)
		}

		logPrintf("Applied migration (%s) named '%s'", direction.String(), plannedMigration.ID)
		count++
	}

	err = driver.Close()
	return count, err
}

func getMigrations(migrations Source) ([]*Migration, error) {
	var m []*Migration
	tempMigrations := map[string]*Migration{}

	files, err := migrations.ListMigrationFiles()
	if err != nil {
		return m, err
	}

	regex := regexp.MustCompile(`(\d*_.*)\.(up|down)\..*`)

	for _, file := range files {
		matches := regex.FindStringSubmatch(file)

		if len(matches) > 0 && file == matches[0] {
			id := matches[1]
			direction := matches[2]

			if _, ok := tempMigrations[id]; !ok {
				tempMigrations[id] = &Migration{
					ID: id,
				}
			}

			reader, err := migrations.GetMigrationFile(file)
			if err != nil {
				return m, fmt.Errorf("Error getting migrations: %s", err)
			}

			contents, err := io.ReadAll(reader)
			if err != nil {
				return m, fmt.Errorf("Error getting migration content: %s", err)
			}

			parsed, err := parser.Parse(bytes.NewReader(contents))
			if err != nil {
				return m, fmt.Errorf("Error parsing migration %s: %s", id, err)
			}

			if direction == "up" {
				tempMigrations[id].Up = parsed
			} else {
				tempMigrations[id].Down = parsed
			}
		}
	}

	for _, migration := range tempMigrations {
		m = append(m, migration)
	}

	sort.Sort(byID(m))

	return m, nil
}

func planMigrations(migrations []*Migration, appliedMigrations []string, direction Direction, max int) []*PlannedMigration {
	var applied []*Migration

	for _, appliedMigration := range appliedMigrations {
		applied = append(applied, &Migration{
			ID: appliedMigration,
		})
	}

	sort.Sort(byID(applied))

	// Get last migration that was run
	record := &Migration{}

	if len(applied) > 0 {
		record = applied[len(applied)-1]
	}

	var result []*PlannedMigration

	// Add missing migrations up to the last run migration.
	// This can happen for example when merges happened.
	if len(applied) > 0 {
		result = append(result, toCatchup(migrations, applied, record)...)
	}

	// Figure out which migrations to apply
	toApply := toApply(migrations, record.ID, direction)
	toApplyCount := len(toApply)

	if max > 0 && max < toApplyCount {
		toApplyCount = max
	}

	for _, v := range toApply[0:toApplyCount] {
		result = append(result, &PlannedMigration{
			Migration: v,
			Direction: direction,
		})
	}

	return result
}

// Filter a slice of migrations into ones that should be applied.
func toApply(migrations []*Migration, current string, direction Direction) []*Migration {
	var index = -1

	if current != "" {
		for index < len(migrations)-1 {
			index++
			if migrations[index].ID == current {
				break
			}
		}
	}

	if direction == Up {
		return migrations[index+1:]
	} else if direction == Down {
		if index == -1 {
			return []*Migration{}
		}

		// Add in reverse order
		toApply := make([]*Migration, index+1)
		for i := 0; i < index+1; i++ {
			toApply[index-i] = migrations[i]
		}
		return toApply
	}

	panic("Not possible")
}

// Get migrations that we need to apply regardless of whether the direction is up or down. This is
// because there may be migration "holes" due to merges.
func toCatchup(migrations, existingMigrations []*Migration, lastRun *Migration) []*PlannedMigration {
	var missing []*PlannedMigration

	for _, migration := range migrations {
		found := false

		for _, existing := range existingMigrations {
			if existing.ID == migration.ID {
				found = true
				break
			}
		}

		if !found && migration.Less(lastRun) {
			missing = append(missing, &PlannedMigration{Migration: migration, Direction: Up})
		}
	}

	return missing
}