codex-app-server/cloud/migrations.py

"""SQLite -> Postgres 迁移（控制面状态库）。

只做一件事：把一个 SQLite 库文件里的行原样插进 Postgres，
表名与列名不改。幂等：Postgres 侧用 ``ON CONFLICT DO NOTHING``，
跑两次不重复插（第二次各表 inserted 为 0）。

- schema 来自 cloud.db.POSTGRES_SCHEMA / init_postgres_schema，
  本模块不另发明表结构；迁移前先建表（若目标库已有表，
  CREATE TABLE IF NOT EXISTS 是 no-op）。
- P2-ops 对齐：MIGRATION_TABLES 必须等于 POSTGRES_SCHEMA 的
  表集合（见 migration_coverage 与 test_p2_migrations.py）。
  以前漏掉的 uploads、task_attachments、task_followups、
  environments 已补进清单；补进时发现 task_followups 在
  SQLite 侧比 PG schema 多一列 sender_task_id（peer_messages
  的运行时 ensure_followups_table 会在 PG 上 ALTER 补列，
  POSTGRES_SCHEMA 里没有它），迁移时按同样方式先补列再插，
  否则整批迁移会在该表上撞未知列失败。
- kb_fts（SQLite FTS5 虚拟表）不迁移：Postgres 没有 FTS5，
  KBStore 在 PG 上走 LIKE 分支，kb_docs 的 content 已足够。
- sqlite_sequence（AUTOINCREMENT 内部表）不迁移。
- UNMIGRATED_TABLES：SQLite 各 store 建出、但 POSTGRES_SCHEMA
  里没有的表（api_keys、评审、自动化、GitHub App、网关 key
  等）。它们不是“被跳过”，而是本模块无法迁移——迁移报告
  （migration_coverage / CLI 输出）逐项列出，调用方不得
  把“迁移成功”读成“全部数据都过去了”。
- 不打印、不返回任何密码或密钥：函数只返回 {表: inserted}，
  CLI 只打印表名与行数。
- Postgres 连接参数只走 CLOUD_DB_* 环境变量（见 cloud.db），
  本模块不读、不拼连接串、不接受命令行密码参数。
"""

import os
import re
import sqlite3

from cloud.db import (
    POSTGRES_SCHEMA,
    PostgresDatabase,
    init_postgres_schema,
)

# 迁移顺序：无外键约束，顺序只影响日志可读性；用户表在前，
# 审计/成本等追加表在后。kb_fts 与 sqlite_sequence 明确排除。
MIGRATION_TABLES = (
    "users",
    "projects",
    "project_members",
    "tasks",
    "audit_log",
    "summaries",
    "kb_docs",
    "cost_records",
    "oauth_states",
    "slack_users",
    "push_subscriptions",
    # P0 安全批：配额行必须随库迁移，否则迁后用户变回无行
    # （本批后无行按默认额度处理，不再是无限额，但显式
    # 管理员额度仍需原样迁过去，不能丢）。
    "user_quotas",
    # P2-ops 对齐：以下 5 张表在 POSTGRES_SCHEMA 里有、
    # 原清单里没有，补进后清单与 schema 表集合一致。
    # uploads / task_attachments：附件包与任务附件；
    # task_followups：任务间互发消息的后续消息队列；
    # environments：项目环境（分支/executor）。
    # kill_requests：跨 worker 真 kill 的意图表（残留批次 B）。
    # 五张表主键均为 TEXT 或复合主键，没有身份列，
    # 不需要序列拨动（见 SEQUENCE_TABLES）。
    "uploads",
    "task_attachments",
    "task_followups",
    "environments",
    "kill_requests",
    # Cloud identity batch: TEXT primary keys, no sequence to sync.
    "user_profiles",
    "identity_memories",
    # Auto memory batch: same TEXT-primary-key pattern.
    "memory_settings",
    "task_memory_flags",
    "memory_extraction_log",
    # P0 三件：断点与步骤表（TEXT 主键 / 复合主键，无身份列，
    # 不需要序列拨动；sqlite 侧由 server/worker ensure 建出，
    # 老库没有时迁移记 0，不中断其余表）。
    "task_checkpoints",
    "task_steps",
    # 常驻监听（P0 S3：与 POSTGRES_SCHEMA 同款两表，TEXT 主键，
    # 无身份列；sqlite 侧由 resident.ensure 建出，老库没有时
    # 记 0，不中断其余表；failure_reason 列两端 schema 一致）。
    "resident_watches",
    "resident_polls",
    # 尝试账本（v6）：同款复合主键、无身份列。
    "task_attempts",
)

_SKIP_TABLES = frozenset({"kb_fts", "sqlite_sequence"})

# 身份列的表：POSTGRES_SCHEMA 里唯一用
# ``BIGINT GENERATED BY DEFAULT AS IDENTITY`` 的 4 张表。
# 迁移时显式插 id 不会推序列，迁完必须 setval 到 MAX(id)，
# 否则后续省略 id 的 INSERT 会撞主键。新增身份列表必须
# 同步改这里，test_p2_migrations.py 会拿 schema 实测对齐。
SEQUENCE_TABLES = (
    "audit_log",
    "summaries",
    "kb_docs",
    "cost_records",
)

# SQLite 各 store 建出、但 POSTGRES_SCHEMA 里没有的表：
# 本迁移模块无法把它们迁到 Postgres（目标库无此表），
# 不是静默跳过——migration_coverage() 与 CLI 逐项报告。
# 其中 rate_limit_hits / workers 是易失的限流与心跳数据，
# 丢了只影响限流窗口与 worker 可见性，不影响业务数据；
# 其余（api_keys、task_reviews、automations、GitHub App、
# user_gateway_keys 等）丢了就是真丢，迁库前必须另行
# 处理（重建/重绑/重发 key），不能只跑本迁移就宣称迁完。
UNMIGRATED_TABLES = (
    "api_keys",
    "task_approvals",
    "task_reviews",
    "automations",
    "automation_runs",
    "github_app_setup_states",
    "github_app_config",
    "github_app_installations",
    "github_fix_review_dedup",
    "external_identities",
    "pr_jobs",
    "task_handoffs",
    "task_worktree_snapshots",
    "task_refs",
    "automation_bindings",
    "rate_limit_hits",
    "user_gateway_keys",
    "webhook_deliveries",
    "workers",
)


def schema_tables():
    """POSTGRES_SCHEMA 里的表名集合（只认 CREATE TABLE 语句）。"""
    names = set()
    for statement in POSTGRES_SCHEMA:
        m = re.search(r"CREATE TABLE IF NOT EXISTS (\w+)",
                      statement)
        if m:
            names.add(m.group(1))
    return names


def migration_coverage():
    """迁移覆盖报告：清单、schema、未迁移项三者一次给全。

    migrated：本次迁移模块会迁的表（= MIGRATION_TABLES）；
    schema_tables：POSTGRES_SCHEMA 里的表；
    missing_from_migration：schema 里有、清单里没有的表
        （本批后应为空；非空即对齐回归，调用方应报错
        而不是继续跑迁移）；
    unmigrated_no_pg_schema：UNMIGRATED_TABLES（目标库
        无此表，无法迁移，须另行处理）；
    skipped_by_design：kb_fts / sqlite_sequence（设计上不迁）。
    不含任何密钥或行内容，只含表名。
    """
    schema = schema_tables()
    return {
        "migrated": tuple(MIGRATION_TABLES),
        "schema_tables": tuple(sorted(schema)),
        "missing_from_migration": tuple(
            sorted(schema - set(MIGRATION_TABLES))),
        "unmigrated_no_pg_schema": tuple(UNMIGRATED_TABLES),
        "skipped_by_design": tuple(sorted(_SKIP_TABLES)),
    }


def migrate_sqlite_to_postgres(sqlite_path, pg_conn=None, tables=None):
    """迁移并返回 {表名: 本次插入行数}。

    pg_conn 为 None 时用 PostgresDatabase().connect() 新建连接
    （驱动缺失时抛 PostgresDriverMissingError，绝不回落）；
    传入连接时调用方负责关闭，本函数不关闭它。tables 为 None
    时迁 MIGRATION_TABLES 全量；显式传表名时只迁其中在
    SQLite 里存在的表，不存在的表在结果里记 0（不报错，
    因为老库可能还没建出某张表，例如未用过的 push 表）。
    """
    wanted = tuple(tables) if tables is not None else MIGRATION_TABLES
    src = sqlite3.connect(str(sqlite_path))
    src.row_factory = sqlite3.Row
    own_conn = pg_conn is None
    if own_conn:
        pg_conn = PostgresDatabase().connect()
    try:
        init_postgres_schema(pg_conn)
        result = {}
        for table in wanted:
            if table in _SKIP_TABLES:
                result[table] = 0
                continue
            try:
                rows = src.execute(
                    'SELECT * FROM "%s"' % table).fetchall()
            except sqlite3.OperationalError:
                # 老库里没有这张表：记 0，不中断其余表。
                result[table] = 0
                continue
            if not rows:
                result[table] = 0
                continue
            columns = [d[0] for d in src.execute(
                'SELECT * FROM "%s" LIMIT 0' % table).description]
            if table == "tasks":
                # Audit fix 3（I6）：task_followups.sender_task_id
                # 同款特殊回填。SQLite 侧 tasks 已有 agent 列，
                # POSTGRES_SCHEMA 新库建表时一次建齐，但老
                # PG 库的回填清单此前没有这 5 列；直插前先
                # 查现有列、缺了再 ALTER 补列，已有列是
                # no-op。补列失败不吞掉迁移：INSERT 时若列
                # 仍缺失会直接报错（大声失败，绝不静默丢列）。
                try:
                    _existing = {row[1] for row in pg_conn.execute(
                        "PRAGMA table_info(tasks)").fetchall()}
                except Exception:
                    _existing = set()
                for _col in ("agent_name", "agent_version",
                             "agent_instructions", "agent_tools",
                             "agent_tool_approvals"):
                    if _col not in _existing:
                        try:
                            pg_conn.execute(
                                "ALTER TABLE tasks ADD COLUMN %s TEXT"
                                % _col)
                        except Exception:
                            pass
            if table == "task_followups" and "sender_task_id" in columns:
                # SQLite 侧有、POSTGRES_SCHEMA 里没有的列：
                # 与 peer_messages.ensure_followups_table 同款，
                # 先查现有列、缺了再 ALTER 补列，已有列是 no-op。
                # PRAGMA 经适配层翻译成 information_schema 查询，
                # 取 row[1] 即列名（与 init_postgres_schema 同款）。
                # 补列失败也不在这里吞掉迁移：INSERT 时若列仍
                # 缺失会直接报错（大声失败，绝不静默丢列）。
                try:
                    existing = {row[1] for row in pg_conn.execute(
                        "PRAGMA table_info(task_followups)").fetchall()}
                except Exception:
                    existing = set()
                if "sender_task_id" not in existing:
                    try:
                        pg_conn.execute(
                            "ALTER TABLE task_followups"
                            " ADD COLUMN sender_task_id TEXT")
                    except Exception:
                        pass
            placeholders = ", ".join("?" for _ in columns)
            col_list = ", ".join('"%s"' % c for c in columns)
            inserted = 0
            for row in rows:
                cur = pg_conn.execute(
                    'INSERT INTO "%s" (%s) VALUES (%s)'
                    " ON CONFLICT DO NOTHING"
                    % (table, col_list, placeholders),
                    tuple(row[c] for c in columns))
                if cur.rowcount:
                    inserted += cur.rowcount
            # 身份列的序列不会因显式插 id 自动前进：迁完把序列
            # 拨到当前最大 id，否则后续省略 id 的 INSERT 会撞主键。
            if table in SEQUENCE_TABLES and "id" in columns:
                try:
                    pg_conn.execute(
                        "SELECT setval(pg_get_serial_sequence(?, ?),"
                        " COALESCE((SELECT MAX(id) FROM \"%s\"), 1), true)"
                        % table,
                        (table, "id"))
                except Exception:
                    pass
            pg_conn.commit()
            result[table] = inserted
        return result
    finally:
        src.close()
        if own_conn:
            try:
                pg_conn.close()
            except Exception:
                pass


__all__ = [
    "MIGRATION_TABLES",
    "SEQUENCE_TABLES",
    "UNMIGRATED_TABLES",
    "migration_coverage",
    "migrate_sqlite_to_postgres",
    "schema_tables",
]