cobble/migration_graph.go

package main

import (
	"errors"
	"fmt"
	"log"
	"os"
	"path/filepath"
)

type MigrationGraph struct {
	// Reference to the root of the graph.
	Root *Migration
	// Name to struct mapping of all migrations part of the graph.
	Migrations map[string]Migration
	// Mapping of all migrations to their parent, if any.
	parentage map[string]*Migration
}

func NewMigrationGraph() MigrationGraph {
	return MigrationGraph{Migrations: map[string]Migration{}, parentage: map[string]*Migration{}}
}

func NewMigrationGraphFromDirectory(pathRoot string) (MigrationGraph, error) {
	migrationGraph := NewMigrationGraph()

	files, err := os.ReadDir(pathRoot)

	if err != nil {
		log.Printf("Failed to read migration directory: %s", pathRoot)
		return migrationGraph, err
	}
	for _, file := range files {
		filename := file.Name()

		if filepath.Ext(filename) != ".sql" {
			continue
		}

		migrationPath := filepath.Join(pathRoot, file.Name())

		file, _ := os.Open(migrationPath)

		headers := gatherMigrationHeaders(file)

		requirements := ""

		if len(headers.Requirements) > 0 {
			requirements = headers.Requirements[0]
		}

		migration := NewMigration(migrationPath, filepath.Base(file.Name()), requirements)
		migrationGraph.AddMigration(migration)
	}

	return migrationGraph, nil
}

// Adds a migration to the graph.
//
// This also adds the migration to the parentage mappings to link it
// to its parent. If the migration added has no parent, then it's also
// set to be the root of the graph.
func (g *MigrationGraph) AddMigration(m Migration) {
	if m.Requires == "" {
		g.addRoot(m)
	}

	g.Migrations[m.Name] = m
	g.parentage[m.Requires] = &m
}

// Builds the linear history of migrations.
func (g *MigrationGraph) GetLinearHistory() ([]Migration, error) {
	if g.Root == nil {
		return []Migration{}, errors.New("Cannot get linear history, the graph has no root.")
	}

	ordered := []Migration{}
	visited := map[string]bool{}
	unordered := []Migration{*g.Root}

	for len(unordered) > 0 {
		migration := unordered[0]
		unordered = unordered[1:]

		if _, hasVisited := visited[migration.Name]; hasVisited {
			return []Migration{}, errors.New("Cycle detected, cannot generate linear history.")
		}

		child, hasChildren := g.parentage[migration.Name]

		ordered = append(ordered, migration)

		if hasChildren {
			unordered = append(unordered, *child)
		}

		visited[migration.Name] = true
	}

	if len(ordered) != len(g.Migrations) {
		return ordered, errors.New("Not all migrations in the graph are part of the history.")
	}

	return ordered, nil
}

func (g *MigrationGraph) addRoot(m Migration) error {
	if g.Root != nil {
		return errors.New(fmt.Sprintf("Cannot have more than one root, tried to add %#v but already knew about %#v.", m, g.Root))
	}

	g.Root = &m

	return nil
}