TaskPool/scripts/migrate-db-brand.go

// 将本地 sqlite 库从 baihu 前缀迁移为 taskpool 前缀，并更新触发类型/站点标题。
//
// 用法（项目根目录）:
//
//	go run ./scripts/migrate-db-brand.go
//	go run ./scripts/migrate-db-brand.go -db data/baihu.db -from baihu_ -to taskpool_
package main

import (
	"database/sql"
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	_ "github.com/glebarez/sqlite"
)

func main() {
	dbPath := flag.String("db", "data/baihu.db", "sqlite database path")
	fromPrefix := flag.String("from", "baihu_", "old table prefix")
	toPrefix := flag.String("to", "taskpool_", "new table prefix")
	outPath := flag.String("out", "data/taskpool.db", "output database path (copy then migrate)")
	dryRun := flag.Bool("dry-run", false, "only print actions")
	flag.Parse()

	if err := run(*dbPath, *outPath, *fromPrefix, *toPrefix, *dryRun); err != nil {
		fmt.Fprintf(os.Stderr, "migrate failed: %v\n", err)
		os.Exit(1)
	}
}

func run(srcDB, outDB, fromPrefix, toPrefix string, dryRun bool) error {
	if _, err := os.Stat(srcDB); err != nil {
		return fmt.Errorf("source db not found: %s", srcDB)
	}

	// 先 checkpoint wal，确保主库完整
	if err := checkpoint(srcDB); err != nil {
		fmt.Printf("warn: checkpoint: %v\n", err)
	}

	absSrc, _ := filepath.Abs(srcDB)
	absOut, _ := filepath.Abs(outDB)
	fmt.Printf("source: %s\noutput: %s\nprefix: %s -> %s\n", absSrc, absOut, fromPrefix, toPrefix)

	if dryRun {
		db, err := sql.Open("sqlite", srcDB)
		if err != nil {
			return err
		}
		defer db.Close()
		tables, err := listTables(db, fromPrefix)
		if err != nil {
			return err
		}
		for _, t := range tables {
			fmt.Printf("would rename %s -> %s\n", t, strings.Replace(t, fromPrefix, toPrefix, 1))
		}
		return nil
	}

	if err := os.MkdirAll(filepath.Dir(outDB), 0o755); err != nil {
		return err
	}

	// 备份源库
	bak := srcDB + ".bak-" + time.Now().Format("20060102-150405")
	if err := copyFile(srcDB, bak); err != nil {
		return fmt.Errorf("backup failed: %w", err)
	}
	fmt.Printf("backup: %s\n", bak)

	// 复制到目标（若目标就是源，则原地改）
	same := filepath.Clean(absSrc) == filepath.Clean(absOut)
	work := outDB
	if !same {
		// 清理目标旁路文件
		_ = os.Remove(outDB)
		_ = os.Remove(outDB + "-wal")
		_ = os.Remove(outDB + "-shm")
		if err := copyFile(srcDB, outDB); err != nil {
			return fmt.Errorf("copy to output failed: %w", err)
		}
		// 若源有 wal，已 checkpoint，通常足够
	} else {
		work = srcDB
	}

	db, err := sql.Open("sqlite", work)
	if err != nil {
		return err
	}
	defer db.Close()

	if _, err := db.Exec(`PRAGMA foreign_keys = OFF`); err != nil {
		return err
	}

	tables, err := listTables(db, fromPrefix)
	if err != nil {
		return err
	}
	if len(tables) == 0 {
		// 也许已经迁过
		tables2, _ := listTables(db, toPrefix)
		if len(tables2) > 0 {
			fmt.Println("tables already use new prefix; continue data fixes")
		} else {
			return fmt.Errorf("no tables with prefix %q found", fromPrefix)
		}
	}

	tx, err := db.Begin()
	if err != nil {
		return err
	}

	for _, oldName := range tables {
		newName := strings.Replace(oldName, fromPrefix, toPrefix, 1)
		if newName == oldName {
			continue
		}
		// 若新表已存在则跳过
		var exists int
		_ = tx.QueryRow(`SELECT COUNT(1) FROM sqlite_master WHERE type='table' AND name=?`, newName).Scan(&exists)
		if exists > 0 {
			fmt.Printf("skip (exists): %s\n", newName)
			continue
		}
		sqlStmt := fmt.Sprintf(`ALTER TABLE "%s" RENAME TO "%s"`, oldName, newName)
		if _, err := tx.Exec(sqlStmt); err != nil {
			_ = tx.Rollback()
			return fmt.Errorf("rename %s -> %s: %w", oldName, newName, err)
		}
		fmt.Printf("renamed: %s -> %s\n", oldName, newName)
	}

	// 索引名里若带旧前缀，SQLite 一般会随表走；对显式旧索引再尝试重命名
	indexes, err := listIndexes(tx, fromPrefix)
	if err != nil {
		_ = tx.Rollback()
		return err
	}
	for _, oldIdx := range indexes {
		newIdx := strings.Replace(oldIdx, fromPrefix, toPrefix, 1)
		if newIdx == oldIdx {
			continue
		}
		var exists int
		_ = tx.QueryRow(`SELECT COUNT(1) FROM sqlite_master WHERE type='index' AND name=?`, newIdx).Scan(&exists)
		if exists > 0 {
			continue
		}
		// SQLite 3.26+ supports ALTER INDEX RENAME; if fails, ignore
		stmt := fmt.Sprintf(`ALTER INDEX "%s" RENAME TO "%s"`, oldIdx, newIdx)
		if _, err := tx.Exec(stmt); err != nil {
			fmt.Printf("warn: index rename %s: %v\n", oldIdx, err)
		} else {
			fmt.Printf("renamed index: %s -> %s\n", oldIdx, newIdx)
		}
	}

	settingsTable := toPrefix + "settings"
	tasksTable := toPrefix + "tasks"

	// 站点标题
	if hasTable(tx, settingsTable) {
		res, err := tx.Exec(`UPDATE "`+settingsTable+`" SET value = ? WHERE section = 'site' AND key = 'title' AND (value = '白虎面板' OR value = 'Baihu Panel' OR value LIKE '%白虎%')`, "任务池")
		if err != nil {
			fmt.Printf("warn: update site title: %v\n", err)
		} else if n, _ := res.RowsAffected(); n > 0 {
			fmt.Printf("updated site title rows: %d\n", n)
		}
	}

	// 触发类型
	if hasTable(tx, tasksTable) {
		res, err := tx.Exec(`UPDATE "`+tasksTable+`" SET trigger_type = 'taskpool_startup' WHERE trigger_type = 'baihu_startup'`)
		if err != nil {
			fmt.Printf("warn: update trigger_type: %v\n", err)
		} else if n, _ := res.RowsAffected(); n > 0 {
			fmt.Printf("updated trigger_type rows: %d\n", n)
		}
	}

	if err := tx.Commit(); err != nil {
		return err
	}

	if _, err := db.Exec(`VACUUM`); err != nil {
		fmt.Printf("warn: vacuum: %v\n", err)
	}
	if _, err := db.Exec(`PRAGMA wal_checkpoint(FULL)`); err != nil {
		fmt.Printf("warn: final checkpoint: %v\n", err)
	}

	fmt.Println("migrate done")
	fmt.Println("next:")
	fmt.Println("  1) update configs/config.ini:")
	fmt.Println("       path = data/taskpool.db")
	fmt.Println("       table_prefix = taskpool_")
	fmt.Println("       dbname = taskpool")
	fmt.Println("  2) restart: bin\\taskpool.exe server")
	return nil
}

func checkpoint(dbPath string) error {
	db, err := sql.Open("sqlite", dbPath)
	if err != nil {
		return err
	}
	defer db.Close()
	_, err = db.Exec(`PRAGMA wal_checkpoint(FULL)`)
	return err
}

func listTables(db *sql.DB, prefix string) ([]string, error) {
	rows, err := db.Query(`SELECT name FROM sqlite_master WHERE type='table' AND name LIKE ? AND name NOT LIKE 'sqlite_%' ORDER BY name`, prefix+"%")
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []string
	for rows.Next() {
		var name string
		if err := rows.Scan(&name); err != nil {
			return nil, err
		}
		out = append(out, name)
	}
	return out, rows.Err()
}

func listIndexes(tx *sql.Tx, prefix string) ([]string, error) {
	rows, err := tx.Query(`SELECT name FROM sqlite_master WHERE type='index' AND name LIKE ? AND name NOT LIKE 'sqlite_%' ORDER BY name`, prefix+"%")
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []string
	for rows.Next() {
		var name string
		if err := rows.Scan(&name); err != nil {
			return nil, err
		}
		out = append(out, name)
	}
	return out, rows.Err()
}

func hasTable(tx *sql.Tx, name string) bool {
	var n int
	_ = tx.QueryRow(`SELECT COUNT(1) FROM sqlite_master WHERE type='table' AND name=?`, name).Scan(&n)
	return n > 0
}

func copyFile(src, dst string) error {
	in, err := os.ReadFile(src)
	if err != nil {
		return err
	}
	return os.WriteFile(dst, in, 0o644)
}