NixMsg/internal/store/migrate.go

package store

import (
	"database/sql"
	"embed"
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"time"
)

//go:embed migrations/*.sql
var migrationFS embed.FS

// Migrate 应用尚未执行的嵌入迁移。
// 若 dbExisted 为 true（调用 Open/Migrate 前已有 nixmsg.db）且存在未应用版本，
// 先 VACUUM INTO <data_dir>/backup/pre-migrate-v{当前}-to-v{目标}.db，再迁移。
// 备份文件已存在则复用，避免迁移稳定失败时反复全量复制。
// 任一版本失败则返回错误，调用方不得继续带半新半旧库提供服务。
func Migrate(db *sql.DB, dataDir string, dbExisted bool) error {
	if err := ensureMigrationsTable(db); err != nil {
		return err
	}

	pending, err := pendingMigrations(db)
	if err != nil {
		return err
	}
	return applyPending(db, dataDir, dbExisted, pending)
}

func ensureMigrationsTable(db *sql.DB) error {
	if _, err := db.Exec(`
CREATE TABLE IF NOT EXISTS schema_migrations (
  version INTEGER PRIMARY KEY,
  applied_at INTEGER NOT NULL
)`); err != nil {
		return fmt.Errorf("ensure schema_migrations: %w", err)
	}
	return nil
}

func applyPending(db *sql.DB, dataDir string, dbExisted bool, pending []migrationFile) error {
	if len(pending) == 0 {
		return nil
	}

	if dbExisted {
		from, err := currentSchemaVersion(db)
		if err != nil {
			return err
		}
		to := pending[len(pending)-1].version
		if err := backupBeforeMigrate(db, dataDir, from, to); err != nil {
			return err
		}
	}

	for _, m := range pending {
		if err := applyMigration(db, m); err != nil {
			return err
		}
	}
	return nil
}

type migrationFile struct {
	version int
	name    string
	body    string
}

func pendingMigrations(db *sql.DB) ([]migrationFile, error) {
	entries, err := fs.ReadDir(migrationFS, "migrations")
	if err != nil {
		return nil, fmt.Errorf("read migrations: %w", err)
	}
	var names []string
	for _, e := range entries {
		if e.IsDir() || !strings.HasSuffix(e.Name(), ".sql") {
			continue
		}
		names = append(names, e.Name())
	}
	sort.Strings(names)

	var pending []migrationFile
	for _, name := range names {
		version, err := parseMigrationVersion(name)
		if err != nil {
			return nil, err
		}
		var exists int
		err = db.QueryRow(`SELECT 1 FROM schema_migrations WHERE version = ?`, version).Scan(&exists)
		if err == nil {
			continue
		}
		if err != sql.ErrNoRows {
			return nil, fmt.Errorf("check migration %d: %w", version, err)
		}
		body, err := migrationFS.ReadFile("migrations/" + name)
		if err != nil {
			return nil, fmt.Errorf("read migration %s: %w", name, err)
		}
		pending = append(pending, migrationFile{
			version: version,
			name:    name,
			body:    strings.TrimSpace(string(body)),
		})
	}
	return pending, nil
}

func currentSchemaVersion(db *sql.DB) (int, error) {
	var v sql.NullInt64
	if err := db.QueryRow(`SELECT MAX(version) FROM schema_migrations`).Scan(&v); err != nil {
		return 0, fmt.Errorf("current schema version: %w", err)
	}
	if !v.Valid {
		return 0, nil
	}
	return int(v.Int64), nil
}

func backupBeforeMigrate(db *sql.DB, dataDir string, fromVer, toVer int) error {
	backupDir := filepath.Join(dataDir, "backup")
	if err := os.MkdirAll(backupDir, 0o755); err != nil {
		return fmt.Errorf("mkdir backup: %w", err)
	}
	backupPath := filepath.Join(backupDir, fmt.Sprintf("pre-migrate-v%d-to-v%d.db", fromVer, toVer))
	exists, err := fileExists(backupPath)
	if err != nil {
		return err
	}
	if exists {
		return nil
	}
	if err := ensureDiskSpace(backupDir, backupNeedBytes(dataDir)); err != nil {
		return err
	}
	// SQLite VACUUM INTO 需要字面量路径；统一用斜杠，并对单引号转义。
	quoted := strings.ReplaceAll(filepath.ToSlash(backupPath), "'", "''")
	if _, err := db.Exec("VACUUM INTO '" + quoted + "'"); err != nil {
		return fmt.Errorf("vacuum into backup %s: %w", backupPath, err)
	}
	return nil
}

func backupNeedBytes(dataDir string) int64 {
	var n int64
	for _, name := range []string{DBFileName, DBFileName + "-wal", DBFileName + "-shm"} {
		st, err := os.Stat(filepath.Join(dataDir, name))
		if err == nil {
			n += st.Size()
		}
	}
	return n
}

func ensureDiskSpace(dir string, need int64) error {
	if need < 0 {
		need = 0
	}
	avail, err := availableBytes(dir)
	if err != nil {
		// 探测失败不阻断迁移，VACUUM INTO 自身会因空间不足报错。
		return nil
	}
	const margin = 1 << 20
	want := uint64(need) + margin
	if avail < want {
		return fmt.Errorf("insufficient disk space for migrate backup: have %d bytes, need ~%d", avail, want)
	}
	return nil
}

func applyMigration(db *sql.DB, m migrationFile) error {
	tx, err := db.Begin()
	if err != nil {
		return fmt.Errorf("begin migration %d: %w", m.version, err)
	}
	if m.body != "" {
		if _, err := tx.Exec(m.body); err != nil {
			_ = tx.Rollback()
			return fmt.Errorf("apply migration %d (%s): %w", m.version, m.name, err)
		}
	}
	if _, err := tx.Exec(
		`INSERT INTO schema_migrations(version, applied_at) VALUES(?, ?)`,
		m.version,
		time.Now().UnixMilli(),
	); err != nil {
		_ = tx.Rollback()
		return fmt.Errorf("record migration %d: %w", m.version, err)
	}
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("commit migration %d: %w", m.version, err)
	}
	return nil
}

func parseMigrationVersion(name string) (int, error) {
	base := strings.TrimSuffix(name, ".sql")
	parts := strings.SplitN(base, "_", 2)
	if len(parts) < 1 {
		return 0, fmt.Errorf("invalid migration name: %s", name)
	}
	v, err := strconv.Atoi(parts[0])
	if err != nil {
		return 0, fmt.Errorf("invalid migration version in %s: %w", name, err)
	}
	return v, nil
}