MailUI4Agents/gateway/internal/db/migrate.go

package db

import (
	"context"
	_ "embed"
	"fmt"
	"strings"
)

//go:embed migrations/init.sql
var initSQLPostgres string

//go:embed migrations/init_sqlite.sql
var initSQLSQLite string

// Migrate 建表建索引。两种方言各有一份 schema，语义保持一致。
func Migrate(ctx context.Context) error {
	switch D {
	case Postgres:
		// PG 侧含 DO $$ … $$ 迁移块，必须整体提交
		if _, err := DB.ExecContext(ctx, initSQLPostgres); err != nil {
			return fmt.Errorf("migrate postgres: %w", err)
		}
	case SQLite:
		// modernc.org/sqlite 的 Exec 不接受多语句，逐条执行
		for i, stmt := range splitStatements(initSQLSQLite) {
			if _, err := DB.ExecContext(ctx, stmt); err != nil {
				return fmt.Errorf("migrate sqlite (语句 #%d: %.60s): %w", i+1, stmt, err)
			}
		}
		// CREATE TABLE IF NOT EXISTS 不会给**已存在**的表补列，而 SQLite 又没有
		// ADD COLUMN IF NOT EXISTS。已部署的库靠这一步补齐新列。
		if err := addMissingColumns(ctx); err != nil {
			return err
		}
	default:
		return fmt.Errorf("migrate: 未初始化的方言")
	}

	fmt.Printf("数据库迁移完成（%s）\n", D)
	return nil
}

// splitStatements 按分号切分 SQL 脚本并剔除注释行。
// 本项目的 SQLite schema 只有 CREATE 语句，不含字符串字面量里的分号，
// 因此按分号朴素切分是安全的；若将来加入含分号的字面量需改用真正的词法切分。
func splitStatements(script string) []string {
	var out []string
	for _, raw := range strings.Split(script, ";") {
		var lines []string
		for _, line := range strings.Split(raw, "\n") {
			if t := strings.TrimSpace(line); t == "" || strings.HasPrefix(t, "--") {
				continue
			}
			lines = append(lines, line)
		}
		if stmt := strings.TrimSpace(strings.Join(lines, "\n")); stmt != "" {
			out = append(out, stmt)
		}
	}
	return out
}

// sqliteAddColumns 声明 SQLite 侧需要在已存在的表上补齐的列。
//
// 新库由 init_sqlite.sql 的 CREATE TABLE 一次建全，这里只服务**已部署的库**。
// PG 侧用 ALTER TABLE ... ADD COLUMN IF NOT EXISTS 就够，SQLite 没有这个语法，
// 只能先查 pragma 再决定加不加。
//
// 新增列时同时改两处：init_sqlite.sql 的 CREATE TABLE（给新库）与这张表（给老库）。
var sqliteAddColumns = []struct{ table, column, ddl string }{
	{"mails", "rename_alias", "ALTER TABLE mails ADD COLUMN rename_alias TEXT"},
	{"mails", "rename_reason", "ALTER TABLE mails ADD COLUMN rename_reason TEXT"},
	{"sessions", "rename_dismissed", "ALTER TABLE sessions ADD COLUMN rename_dismissed TEXT"},
	{"sessions", "alias_source", "ALTER TABLE sessions ADD COLUMN alias_source TEXT NOT NULL DEFAULT 'platform'"},
	// 会话级往返预算（0 = 不限）。旧库默认 0：引入预算不应该把已在进行的会话卡死。
	{"sessions", "max_rounds", "ALTER TABLE sessions ADD COLUMN max_rounds INTEGER NOT NULL DEFAULT 0"},
	{"sessions", "used_rounds", "ALTER TABLE sessions ADD COLUMN used_rounds INTEGER NOT NULL DEFAULT 0"},
	// 会话所属的工作目录。旧库默认空串：历史会话的 workspace 无法可靠反推
	// （Agent 回信的 from_workspace 存的是 Agent 名而不是路径），强行回填只会
	// 造出一批看起来有值实际是错的数据。
	{"sessions", "workspace", "ALTER TABLE sessions ADD COLUMN workspace TEXT NOT NULL DEFAULT ''"},
	// 派给该 Agent 的新任务默认多少个来回。
	// 旧库也给 20：之前的 max_rounds 默认是 10 但那是终身额度，语义不同，
	// 不能直接搬过来当单任务预算。
	{"agents", "default_rounds", "ALTER TABLE agents ADD COLUMN default_rounds INTEGER NOT NULL DEFAULT 20"},
}

// sqliteAddIndexes 是建表后才能建的索引（依赖上面补的列）。
// CREATE INDEX IF NOT EXISTS 天然幂等，直接执行即可。
var sqliteAddIndexes = []string{
	// 人类决策后要按 mail_id 反查上游 permission id
	"CREATE INDEX IF NOT EXISTS idx_relayed_mail ON relayed_mails(mail_id)",
}

func addMissingColumns(ctx context.Context) error {
	for _, c := range sqliteAddColumns {
		has, err := columnExists(ctx, c.table, c.column)
		if err != nil {
			return fmt.Errorf("migrate sqlite: 检查 %s.%s: %w", c.table, c.column, err)
		}
		if has {
			continue
		}
		if _, err := DB.ExecContext(ctx, c.ddl); err != nil {
			return fmt.Errorf("migrate sqlite: 补列 %s.%s: %w", c.table, c.column, err)
		}
		fmt.Printf("补列 %s.%s\n", c.table, c.column)
	}
	for _, ddl := range sqliteAddIndexes {
		if _, err := DB.ExecContext(ctx, ddl); err != nil {
			return fmt.Errorf("migrate sqlite: 建索引 %.60s: %w", ddl, err)
		}
	}
	return nil
}

func columnExists(ctx context.Context, table, column string) (bool, error) {
	// pragma_table_info 是表函数形式的 PRAGMA，可以直接当表查（比解析 PRAGMA 输出干净）。
	// table 与 column 都来自上面的硬编码常量表，不存在注入面。
	var n int
	err := DB.QueryRowContext(ctx,
		`SELECT COUNT(*) FROM pragma_table_info(?) WHERE name = ?`,
		table, column).Scan(&n)
	return n > 0, err
}