host
44e0f4f1d6
- B-A: migrations.py tasks backfill loop now types audit_loop as INTEGER (other 6 columns stay TEXT), matching db.py backfill, server/worker ensure lists; TEXT readback '0' is truthy and would silently enable the audit loop for audit_loop=0 rows. - B-C: gateway preparation failure early-return now calls _running_clear(task_id) next to _cleanup_task_seed, matching missing-key/repo/capability/tool gates; previously the capacity slot leaked until worker exit (default capacity 1 blocks reclaim). - Tests: cloud/tests/test_audit_batch1.py (3 tests): migration ALTER types, INTEGER 0 readback not enabled via real worker row mapping, gateway failure clears slot and capacity reclaim. _claim doc 'three gates' list left unchanged (full early-return gate enumeration not verified this batch; residual wording). - No other audit items touched (I-1..I-14, N-1..N-12).
306 lines
13 KiB
Python
306 lines
13 KiB
Python
"""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",
|
||
# 审计循环状态（audit loop）：task_id 单行主键、无身份列。
|
||
"task_audit_loop_state",
|
||
)
|
||
|
||
_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, _coltype in (
|
||
("agent_name", "TEXT"),
|
||
("agent_version", "TEXT"),
|
||
("agent_instructions", "TEXT"),
|
||
("agent_tools", "TEXT"),
|
||
("agent_tool_approvals", "TEXT"),
|
||
# 审计循环接线：开关列必须与 db.py 回填清单、
|
||
# server/worker ensure 清单一致用 INTEGER；
|
||
# 若建成 TEXT，PG 读回的 "0" 为真值，worker
|
||
# 的 `or 0` 兜底拦不住，audit_loop=0 的
|
||
# 历史任务会被悄悄当成启用（Audit batch 1）。
|
||
("audit_loop", "INTEGER"),
|
||
("audit_loop_packet", "TEXT")):
|
||
if _col not in _existing:
|
||
try:
|
||
pg_conn.execute(
|
||
"ALTER TABLE tasks ADD COLUMN %s %s"
|
||
% (_col, _coltype))
|
||
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",
|
||
]
|