slogan/scripts/split_db/main.go

// 一次性迁移工具：把单文件 slogan.db 拆分为 4 个 SQLite 库
//
//	slogan.db      主库（用户域 + 低频配置）
//	slogan_plan.db 穿搭方案域
//	slogan_pay.db  会员/支付域
//	slogan_cps.db  CPS 联盟域
//
// 用法：在 slogan-agent 目录执行 `go run ./scripts/split_db`
package main

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

	_ "github.com/gogf/gf/contrib/drivers/sqlite/v2"
)

const mainDB = "slogan.db"

var groups = []struct {
	file   string
	tables []string
}{
	{"slogan_plan.db", []string{
		"slogan_outfit_generation_task", "slogan_outfit_plan", "slogan_plan_outfit_item",
		"slogan_plan_effect_image", "slogan_plan_review", "slogan_hairstyle_asset",
	}},
	{"slogan_pay.db", []string{
		"slogan_member_plan", "slogan_user_member", "slogan_payment_order",
		"slogan_pay_notify_log", "slogan_ad_reward_log",
	}},
	{"slogan_cps.db", []string{
		"slogan_cps_category", "slogan_cps_product", "slogan_cps_click_log",
		"slogan_scene_category_map",
	}},
}

func main() {
	dir, err := os.Getwd()
	if err != nil {
		fatal("getwd: %v", err)
	}
	mainPath := filepath.Join(dir, mainDB)
	if _, err := os.Stat(mainPath); err != nil {
		fatal("slogan.db 不存在（请在 slogan-agent 目录执行）: %v", err)
	}
	backup := filepath.Join(dir, "slogan_backup_"+time.Now().Format("20060102_150405")+".db")
	if err := copyFile(mainPath, backup); err != nil {
		fatal("备份失败: %v", err)
	}
	fmt.Printf("已备份 -> %s\n", backup)

	src, err := sql.Open("sqlite", mainPath)
	if err != nil {
		fatal("open %s: %v", mainDB, err)
	}
	src.SetMaxOpenConns(1)
	defer src.Close()

	ddl, err := loadDDL(src)
	if err != nil {
		fatal("读取 DDL: %v", err)
	}
	srcAbs, err := filepath.Abs(mainPath)
	if err != nil {
		fatal("abs: %v", err)
	}

	failed := false
	for _, g := range groups {
		if err := migrateGroup(dir, src, srcAbs, g.file, g.tables, ddl); err != nil {
			fmt.Printf("❌ %s 迁移失败: %v\n", g.file, err)
			failed = true
		}
	}
	if failed {
		os.Exit(1)
	}

	// 主库删除已迁出的表（连带索引），清理自增序列残留
	for _, g := range groups {
		for _, t := range g.tables {
			if _, err := src.Exec("DROP TABLE IF EXISTS " + t); err != nil {
				fatal("DROP %s: %v", t, err)
			}
		}
		names := quoteList(g.tables)
		if _, err := src.Exec("DELETE FROM sqlite_sequence WHERE name IN (" + names + ")"); err != nil {
			fmt.Printf("⚠ 清理 sqlite_sequence 失败（可忽略）: %v\n", err)
		}
	}
	fmt.Printf("✅ %s 主库已清理，剩余表：\n", mainDB)
	if err := listTables(src, mainDB); err != nil {
		fatal("list: %v", err)
	}
	fmt.Println("✅ 拆分完成")
}

// migrateGroup 新建目标库文件并拷贝表 + 索引 + 校验行数
func migrateGroup(dir string, src *sql.DB, srcAbs, file string, tables []string, ddl map[string][]string) error {
	dstPath := filepath.Join(dir, file)
	_ = os.Remove(dstPath) // 覆盖上次失败残留
	dst, err := sql.Open("sqlite", dstPath)
	if err != nil {
		return err
	}
	dst.SetMaxOpenConns(1)
	defer dst.Close()

	if _, err := dst.Exec(fmt.Sprintf("ATTACH DATABASE %q AS src", srcAbs)); err != nil {
		return fmt.Errorf("attach: %w", err)
	}
	defer dst.Exec("DETACH DATABASE src")

	for _, t := range tables {
		createDDL, ok := firstByType(ddl[t], "table")
		if !ok {
			return fmt.Errorf("表 %s 未找到建表 DDL", t)
		}
		if _, err := dst.Exec(createDDL); err != nil {
			return fmt.Errorf("create %s: %w", t, err)
		}
		if _, err := dst.Exec(fmt.Sprintf("INSERT INTO %s SELECT * FROM src.%s", t, t)); err != nil {
			return fmt.Errorf("copy %s: %w", t, err)
		}
	}
	for _, t := range tables {
		for _, idx := range ddl[t] {
			if !strings.HasPrefix(idx, "CREATE INDEX") && !strings.HasPrefix(idx, "CREATE UNIQUE INDEX") {
				continue
			}
			if _, err := dst.Exec(idx); err != nil {
				return fmt.Errorf("index %s: %w", idx, err)
			}
		}
	}
	// 校验行数
	for _, t := range tables {
		srcN, dstN, err := countPair(src, dst, t)
		if err != nil {
			return err
		}
		if srcN != dstN {
			return fmt.Errorf("%s 行数不一致: src=%d dst=%d", t, srcN, dstN)
		}
		fmt.Printf("✅ %-24s %8d 行\n", t, dstN)
	}
	return nil
}

func loadDDL(db *sql.DB) (map[string][]string, error) {
	rows, err := db.Query("SELECT type, name, sql FROM sqlite_master WHERE sql IS NOT NULL AND name NOT LIKE 'sqlite_%' AND type IN ('table','index') ORDER BY type DESC")
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := map[string][]string{}
	for rows.Next() {
		var typ, name, sqlText string
		if err := rows.Scan(&typ, &name, &sqlText); err != nil {
			return nil, err
		}
		if strings.HasPrefix(sqlText, "CREATE TABLE") {
			out[name] = append([]string{sqlText}, out[name]...) // table 放最前
		} else {
			out[name] = append(out[name], sqlText)
		}
	}
	return out, rows.Err()
}

func firstByType(ddls []string, prefix string) (string, bool) {
	for _, d := range ddls {
		if strings.HasPrefix(d, "CREATE TABLE") {
			return d, true
		}
	}
	return "", false
}

func countPair(src, dst *sql.DB, table string) (int, int, error) {
	srcN, err := count(src, table)
	if err != nil {
		return 0, 0, err
	}
	dstN, err := count(dst, table)
	if err != nil {
		return 0, 0, err
	}
	return srcN, dstN, nil
}

func count(db *sql.DB, table string) (int, error) {
	var n int
	err := db.QueryRow("SELECT COUNT(*) FROM " + table).Scan(&n)
	return n, err
}

func listTables(db *sql.DB, file string) error {
	rows, err := db.Query("SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY name")
	if err != nil {
		return err
	}
	defer rows.Close()
	for rows.Next() {
		var name string
		if err := rows.Scan(&name); err != nil {
			return err
		}
		fmt.Printf("   %s\n", name)
	}
	return rows.Err()
}

func quoteList(items []string) string {
	q := make([]string, len(items))
	for i, s := range items {
		q[i] = "'" + s + "'"
	}
	return strings.Join(q, ",")
}

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

func fatal(format string, args ...any) {
	fmt.Printf("❌ "+format+"\n", args...)
	os.Exit(1)
}