"""数据库接口层（云端控制面单机阶段）。 本模块分两层，调用方按需选用： 1. 模块级函数（历史接口，原样保留，见下文）：``connect`` / ``is_locked_error`` / ``execute_with_retry`` / ``read_all`` / ``transaction``。server.py 与旧测试仍走这条路，行为不变。 2. ``Database`` 接口 + 两个实现：``SQLiteDatabase``（现有 SQLite 行为的类封装）与 ``PostgresDatabase``（真实 Postgres 实现， 见下文的适配层）。 新代码应走接口 + ``create_database`` 工厂，不要再直接拼 sqlite3 连接。 只解决一件事：ThreadingHTTPServer 多线程 + 单个 SQLite 文件时的写冲突。 不碰任何业务逻辑、不定义新 schema、不改 server.py 的接线： 本模块只提供连接与读写 helper，协调员后续把 server 的 _connect 换成这里的 connect 即可（表名/字段与 server.py 的 init_db、 memory.py 的 summaries 完全一致，因为本模块根本不建表）。 使用约定（必须遵守，否则 helper 帮不了你）： - 每个线程用自己的连接：绝不在多线程间共用一个 sqlite3 连接 （sqlite3 默认 check_same_thread=True，共用会直接报错）。 典型用法是每请求/每线程调用一次 connect()，用完 close()。 - 写操作走 execute_with_retry + transaction：一次语句失败重试是 安全的（单条 INSERT/UPDATE 被锁时 SQLite 已回滚该语句）； 不要把多条非幂等语句包在一个重试里——重试只针对单条语句。 单机 SQLite 的上限（诚实边界，见模块末尾注释）： - WAL + busy_timeout + 有限重试只能缓解单机多线程写冲突， 不能变成多机/高并发数据库：写仍是单写者串行，读多写少尚可， 写密集或跨机器时必须换 Postgres（下一阶段，不在本轮假装解决）。 后端选择（环境变量，见 cloud/README.md 数据库一节）： - ``CLOUD_DB_BACKEND``：``sqlite``（默认）或 ``postgres``。 - Postgres 连接参数全部走环境变量，不硬编码地址或密码： ``CLOUD_DB_HOST``、``CLOUD_DB_PORT``、``CLOUD_DB_NAME``、 ``CLOUD_DB_USER``、``CLOUD_DB_PASSWORD``。 - Postgres 真后端：``PostgresDatabase`` 已接通真实 Postgres （psycopg v3 优先，psycopg2 其次，延迟加载）。初始化时 找不到驱动仍立即明确报错（``PostgresDriverMissingError``， 文案含 psycopg），绝不静默回落到 SQLite。调用方仍写 sqlite3 风格 SQL（``?`` 占位符、``row["col"]``、``conn.execute`` / ``fetchall`` / ``commit`` / ``rollback`` / ``close``）， 由本模块的 psycopg 适配层（``PgConnection`` / ``PgCursor`` / ``PgRow``）翻译成 Postgres 方言；psycopg 异常被映射成 本模块的 ``PostgresIntegrityError`` / ``PostgresOperationalError`` / ``PostgresConnectionError``， 绝不把驱动异常原样漏成 sqlite3 异常。模块级 ``connect()`` 按 ``CLOUD_DB_BACKEND`` 分发：``postgres`` 时返回适配连接 （``db_path`` 参数被忽略，连接参数只走环境变量），其余 情况仍是原 SQLite 连接——因此 server / AuthStore / MemoryStore / KBStore / CostStore 等沿用 ``connect()`` 的 调用点无需重构即可走 Postgres；但 ``backup.py`` 的 SQLite 在线备份 API 仍只适用于 SQLite，Postgres 备份请用 pg_dump（见 cloud/README.md）。换生产数据库仍不是“改 一个环境变量就自动完成”：还需要真实 Postgres 部署、驱动 安装、建库与迁移脚本（见 cloud/migrations.py 与 cloud/migrate_sqlite_to_postgres.py）。 """ import os import sqlite3 import time from abc import ABC, abstractmethod from contextlib import contextmanager # 与 server.py 旧 _connect 的 timeout=30s 保持一致的连接级等待； # PRAGMA busy_timeout 是语句级等待，两者不冲突，后者以毫秒为单位 # 在每个连接上单独设置（不随库文件持久化，见 connect 的注释）。 DEFAULT_BUSY_TIMEOUT_MS = 5000 DEFAULT_CONNECT_TIMEOUT_S = 30 DEFAULT_MAX_RETRIES = 3 # 指数退避：0.05s、0.1s、0.2s；重试耗尽即抛原始错误，绝不无限循环。 DEFAULT_BACKOFF_SECONDS = (0.05, 0.1, 0.2) _LOCKED_MARKER = "database is locked" # Postgres 连接参数的环境变量名（名字即文档，见 cloud/README.md）。 POSTGRES_ENV_VARS = { "host": "CLOUD_DB_HOST", "port": "CLOUD_DB_PORT", "dbname": "CLOUD_DB_NAME", "user": "CLOUD_DB_USER", "password": "CLOUD_DB_PASSWORD", } BACKEND_ENV_VAR = "CLOUD_DB_BACKEND" DEFAULT_BACKEND = "sqlite" DEFAULT_SQLITE_PATH = "cloud-control-plane.db" def connect(db_path, busy_timeout_ms=DEFAULT_BUSY_TIMEOUT_MS, timeout=DEFAULT_CONNECT_TIMEOUT_S): """返回一个新连接；后端按 CLOUD_DB_BACKEND 分发。 - ``postgres`` / ``postgresql``：返回 Postgres 适配连接 （``PgConnection``，见下文）；``db_path`` 被忽略，连接参数 只走 ``CLOUD_DB_*`` 环境变量；驱动缺失时抛 ``PostgresDriverMissingError``，绝不回落到 SQLite。 - 其它（默认 ``sqlite``）：原 SQLite 行为不变。 """ backend = os.environ.get(BACKEND_ENV_VAR) or DEFAULT_BACKEND if str(backend).strip().lower() in ("postgres", "postgresql"): return PostgresDatabase().connect() return _connect_sqlite(db_path, busy_timeout_ms, timeout) def _connect_sqlite(db_path, busy_timeout_ms=DEFAULT_BUSY_TIMEOUT_MS, timeout=DEFAULT_CONNECT_TIMEOUT_S): """返回一个已设好 WAL + busy_timeout + Row 的新连接。 - PRAGMA journal_mode=WAL：读不阻塞写、写不阻塞读，是多线程 server 的前提；该模式随库文件持久化（:memory: 库除外， 内存库的 journal_mode 恒为 memory，属 SQLite 本身行为）。 - PRAGMA busy_timeout：写冲突时语句级等待窗口，单位毫秒； 它是每个连接的属性，不随库文件持久化，所以每次 connect 都设。 - row_factory=sqlite3.Row：与 server.py 各 handler 在 _db() 里 事后设置的 row_factory 一致，行可用 row["id"] 按列名访问。 - 不设 check_same_thread=False：跨线程共用连接是错误用法， 保持默认 True 让误用尽早报错，而不是静默出竞态。 """ conn = sqlite3.connect(str(db_path), timeout=timeout) try: conn.execute("PRAGMA journal_mode=WAL") conn.execute("PRAGMA busy_timeout=%d" % int(busy_timeout_ms)) conn.row_factory = sqlite3.Row except Exception: conn.close() raise return conn def is_locked_error(exc): """只有 OperationalError 且文案含 database is locked 才算写锁冲突。 其它 OperationalError（如 no such table、语法错误）重试无用， 必须立即抛出，不能被重试掩盖成“锁问题”。 """ return ( isinstance(exc, sqlite3.OperationalError) and _LOCKED_MARKER in str(exc).lower() ) def execute_with_retry(conn, sql, params=(), max_retries=DEFAULT_MAX_RETRIES, backoff_seconds=DEFAULT_BACKOFF_SECONDS): """执行单条写语句，遇写锁冲突时有限重试，返回 cursor。 重试次数 = max_retries（默认 3 次，即最多尝试 1 + 3 = 4 次）， 每次重试前按 backoff_seconds 依次 sleep（默认 0.05/0.1/0.2s， 退避表短于重试次数时取最后一项，不越界）。 重试耗尽后把最后一次的原始 OperationalError 原样抛出，不吞掉、 不包装成别的异常类型，调用方看到的仍是 sqlite3.OperationalError。 非锁错误一次都不重试。 """ attempts = 0 while True: try: return conn.execute(sql, params) except sqlite3.OperationalError as exc: if not is_locked_error(exc) or attempts >= max_retries: raise delay = backoff_seconds[ min(attempts, len(backoff_seconds) - 1) ] time.sleep(delay) attempts += 1 def read_all(conn, sql, params=()): """读操作 helper：不重试，只要求连接本身已走 WAL + busy_timeout。 读在 WAL 下不阻塞写，一般不需要重试；如果调用方传的是旧式 连接（非本模块 connect 建的），读结果仍正确，只是没有 WAL/ busy_timeout 保障——这属于调用方误用，本函数不静默替它建连接。 """ return conn.execute(sql, params).fetchall() @contextmanager def transaction(conn): """单连接事务上下文：成功 commit，异常 rollback 后原样抛出。 与 execute_with_retry 配合的典型写法:: conn = connect(db_path) try: with transaction(conn): execute_with_retry(conn, "INSERT INTO tasks(...) ...", (...)) finally: conn.close() 注意：重试发生在语句级（transaction 内部），事务本身不重试； 若事务内某条语句最终抛错，整个事务回滚，不会留下半写状态。 """ try: yield conn conn.commit() except Exception: try: conn.rollback() except Exception: pass raise # ---------------------------------------------------------------------- # Database 接口层 # ---------------------------------------------------------------------- # 方法名按实际代码为准，只覆盖模块级函数已提供的能力，不 invent # 不存在的方法名：connect（连接 + WAL 设置）、execute_with_retry # （执行）、read_all（查询）、transaction（事务）、is_locked_error # （锁冲突判定）。两个后端的方法签名必须一致，见 test_db_backend.py # 的接口契约测试（按类检查方法名与签名，不依赖驱动是否安装）。 class Database(ABC): """数据库后端接口。实现者不得静默回落到另一个后端。""" @abstractmethod def connect(self): """返回一个新连接（SQLite：已设 WAL + busy_timeout + Row）。""" @abstractmethod def execute_with_retry(self, conn, sql, params=(), max_retries=DEFAULT_MAX_RETRIES, backoff_seconds=DEFAULT_BACKOFF_SECONDS): """执行单条写语句，遇写锁冲突时有限重试，返回 cursor。""" @abstractmethod def read_all(self, conn, sql, params=()): """执行查询并返回全部行。""" @abstractmethod def transaction(self, conn): """返回单连接事务上下文：成功 commit，异常 rollback 后原样抛出。""" @abstractmethod def is_locked_error(self, exc): """判定异常是否为可重试的写锁冲突；非锁错误必须返回 False。""" class SQLiteDatabase(Database): """Database 的 SQLite 实现：原样保留 WAL/busy_timeout/有界重试。 模块级函数仍是同一份实现，本类只是类封装，逐方法委托， 因此旧调用与新接口看到的行为完全一致。 """ def __init__(self, db_path, busy_timeout_ms=DEFAULT_BUSY_TIMEOUT_MS, timeout=DEFAULT_CONNECT_TIMEOUT_S): self.db_path = str(db_path) self.busy_timeout_ms = int(busy_timeout_ms) self.timeout = timeout def __repr__(self): return ("SQLiteDatabase(db_path=%r, busy_timeout_ms=%r, " "timeout=%r)" % (self.db_path, self.busy_timeout_ms, self.timeout)) def connect(self): return connect(self.db_path, busy_timeout_ms=self.busy_timeout_ms, timeout=self.timeout) def execute_with_retry(self, conn, sql, params=(), max_retries=DEFAULT_MAX_RETRIES, backoff_seconds=DEFAULT_BACKOFF_SECONDS): return execute_with_retry(conn, sql, params, max_retries=max_retries, backoff_seconds=backoff_seconds) def read_all(self, conn, sql, params=()): return read_all(conn, sql, params) @contextmanager def transaction(self, conn): with transaction(conn) as cur: yield cur def is_locked_error(self, exc): return is_locked_error(exc) class PostgresDriverMissingError(RuntimeError): """Postgres 驱动缺失：明确报错，绝不静默回落到 SQLite。""" def _load_postgres_driver(): """返回 (驱动名, 模块)；两者都缺时抛 PostgresDriverMissingError。 优先 psycopg（v3），其次 psycopg2；只做 import 探测，不连接、 不读任何真实数据库、不碰任何密码。 """ for name in ("psycopg", "psycopg2"): try: module = __import__(name) except ImportError: continue return name, module raise PostgresDriverMissingError( "Postgres 后端需安装 psycopg 驱动（例如 pip install psycopg），" "当前环境未找到 psycopg/psycopg2；已拒绝静默回落到 SQLite，" "请安装驱动后再选 CLOUD_DB_BACKEND=postgres。" ) # ---------------------------------------------------------------------- # Postgres 适配层：把 sqlite3 风格调用翻译成 psycopg 调用 # ---------------------------------------------------------------------- # 不重构所有 store 的前提：现有代码到处用 ``?`` 占位符、 # ``row["col"]``、``conn.execute().fetchall()``、``commit`` / # ``rollback`` / ``close``、``cursor.lastrowid``、``cursor.rowcount``， # 以及少量 SQLite 方言（``INSERT OR REPLACE``、``INSERT OR # IGNORE``、``PRAGMA table_info``、``AUTOINCREMENT``、FTS5 虚拟表）。 # 适配层在执行前做最小翻译，翻译不了的方言明确报错，绝不静默 # 改语义： # - ``?`` -> ``%s``：只转换单引号字符串字面量之外的 ``?`` # （扫描时跟踪单引号与 ``''`` 转义，字面量里的 ``?`` 不动）； # SQL 里已有的 ``%s``（例如调用方自己拼的 ``IN (%s)`` 占位串 # 在 Python % 格式化之后已变成 ``?``，见 server.py 的 # placeholders 构造）不在此列——适配层只认 ``?``。 # - ``INSERT OR REPLACE`` -> ``ON CONFLICT DO UPDATE``：主键查 # _OR_REPLACE_PK；未知表退化为 ``ON CONFLICT DO NOTHING`` # （语义弱于 REPLACE，调用点只有 project_members 与 # slack_users 两张已知主键表，不会走到退化分支）。 # - ``INSERT OR IGNORE`` -> ``ON CONFLICT DO NOTHING``：无冲突 # 目标，任何唯一冲突都跳过，不改已有行。 # - ``PRAGMA table_info(<表>)`` -> information_schema 查询， # 返回 (cid, name, type, notnull, dflt_value, pk) 六元组， # 与 sqlite 的 PRAGMA table_info 行形状一致，调用方取 row[1] # 即列名；其它 PRAGMA（journal_mode / busy_timeout）是 # SQLite 专属，在 Postgres 连接上直接忽略（返回空结果）。 # - DDL 里的 ``INTEGER PRIMARY KEY AUTOINCREMENT`` -> # ``BIGINT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY``。 # - ``CREATE VIRTUAL TABLE ... USING fts5`` -> 明确抛 # PostgresOperationalError（Postgres 没有 FTS5；KBStore # 捕获后走 LIKE 分支，不要在 Postgres 里建 FTS5 表）。 # - 身份列 INSERT：对 summaries / kb_docs / cost_records / # audit_log 的 INSERT 自动追加 ``RETURNING id``，使 # ``cursor.lastrowid`` 可用；其它 INSERT 不追加。 # # 类型映射（schema 部分，见 POSTGRES_SCHEMA 注释）： # - SQLite REAL（存 epoch 秒浮点）-> DOUBLE PRECISION； # - SQLite INTEGER（存计数、时间戳整数、pid；SQLite INTEGER # 本身就是 64 位）-> BIGINT（计数类）或 INTEGER（小枚举类， # 如 priority / fallback_count / child_pid 的声明见 schema， # 迁移时 Python int 原样传参，Postgres 按列类型收敛）； # - AUTOINCREMENT -> GENERATED BY DEFAULT AS IDENTITY # （BY DEFAULT 而非 ALWAYS：迁移时须能显式插回旧 id， # 正常 INSERT 省略 id 时仍自动生成）； # - UNIQUE / PRIMARY KEY 约束原样保留。 class PostgresDatabaseError(RuntimeError): """Postgres 适配层错误基类：绝不是 sqlite3 异常的别名。""" class PostgresIntegrityError(PostgresDatabaseError): """唯一约束 / 外键违规（对应 psycopg UniqueViolation / ForeignKeyViolation，SQLSTATE 23xxx）。""" class PostgresOperationalError(PostgresDatabaseError): """Postgres 操作错误：语法、表不存在、死锁/锁冲突等。 ``locked`` 为 True 时表示可重试的锁冲突（死锁 40P01、 锁不可用 55P03），供 is_locked_error / execute_with_retry 判定；其它操作错误 locked 为 False，绝不重试。 """ def __init__(self, message, locked=False): super().__init__(message) self.locked = bool(locked) class PostgresConnectionError(PostgresDatabaseError): """连接建立失败（驱动连接错误、SQLSTATE 08xxx 等）。""" def _convert_placeholders(sql): """把 SQL 里字符串字面量之外的 ``?`` 逐个换成 ``%s``。 纯函数，可单测，不碰数据库。单引号字符串用 SQL 规则处理： ``''`` 是转义的单引号，不结束字符串。 """ out = [] in_string = False i = 0 while i < len(sql): ch = sql[i] if in_string: out.append(ch) if ch == "'": if i + 1 < len(sql) and sql[i + 1] == "'": out.append("'") i += 2 continue in_string = False i += 1 else: if ch == "'": in_string = True out.append(ch) i += 1 elif ch == "?": out.append("%s") i += 1 else: out.append(ch) i += 1 return "".join(out) _OR_REPLACE_PK = { "project_members": ("project_id", "user_id"), "slack_users": ("slack_id",), "github_app_config": ("id",), # 自动记忆三张表：主键登记缺失时 Postgres 翻译退化为 # ON CONFLICT DO NOTHING，首写之后的所有更新静默丢弃 # （2026-10-28 生产实测：提取日志永远停在首写原因， # 改设置同理；SQLite 原生 OR REPLACE 测不出来）。 "memory_settings": ("user_id",), "task_memory_flags": ("task_id",), "memory_extraction_log": ("task_id",), } _IDENTITY_TABLES = frozenset( {"summaries", "kb_docs", "cost_records", "audit_log"}) import re as _re _OR_REPLACE_RE = _re.compile( r"^\s*INSERT\s+OR\s+REPLACE\s+INTO\s+(\w+)\s*\(([^)]+)\)\s*" r"VALUES\s*\((.*)\)\s*$", _re.IGNORECASE | _re.DOTALL) _OR_IGNORE_RE = _re.compile( r"^\s*INSERT\s+OR\s+IGNORE\s+INTO\s+(\w+)\s*\(([^)]+)\)\s*" r"VALUES\s*\((.*)\)\s*$", _re.IGNORECASE | _re.DOTALL) _INSERT_TABLE_RE = _re.compile( r"^\s*INSERT\s+INTO\s+(?:\w+\.)?(\w+)", _re.IGNORECASE) _PRAGMA_INFO_RE = _re.compile( r"^\s*PRAGMA\s+table_info\s*\(\s*(\w+)\s*\)\s*;?\s*$", _re.IGNORECASE) _PRAGMA_ANY_RE = _re.compile(r"^\s*PRAGMA\s+", _re.IGNORECASE) _AUTOINCREMENT_RE = _re.compile( r"INTEGER\s+PRIMARY\s+KEY\s+AUTOINCREMENT", _re.IGNORECASE) _FTS5_RE = _re.compile( r"CREATE\s+VIRTUAL\s+TABLE.*USING\s+fts5", _re.IGNORECASE) def _translate_sql(sql): """返回 (translated_sql, kind)：kind 用于游标的特殊处理。 kind 取值：``"normal"`` / ``"pragma_info"``（附表名在 translated_sql 里即表名）/ ``"pragma_ignore"``。 """ stripped = sql.strip() if _FTS5_RE.search(stripped): raise PostgresOperationalError( "Postgres has no FTS5 virtual table; callers must use the" " LIKE search path (KBStore does this).") m = _PRAGMA_INFO_RE.match(stripped) if m: return m.group(1), "pragma_info" if _PRAGMA_ANY_RE.match(stripped): return "", "pragma_ignore" m = _OR_REPLACE_RE.match(stripped) if m: table = m.group(1) cols = [c.strip() for c in m.group(2).split(",")] values = m.group(3) pk = _OR_REPLACE_PK.get(table) if pk: pk_set = set(pk) setters = ", ".join( "%s = EXCLUDED.%s" % (c, c) for c in cols if c not in pk_set) if setters: translated = ( "INSERT INTO %s (%s) VALUES (%s)" " ON CONFLICT (%s) DO UPDATE SET %s" % (table, ", ".join(cols), values, ", ".join(pk), setters)) else: translated = ( "INSERT INTO %s (%s) VALUES (%s)" " ON CONFLICT (%s) DO NOTHING" % (table, ", ".join(cols), values, ", ".join(pk))) else: translated = ( "INSERT INTO %s (%s) VALUES (%s) ON CONFLICT DO NOTHING" % (table, ", ".join(cols), values)) return _convert_placeholders(translated), "normal" m = _OR_IGNORE_RE.match(stripped) if m: # INSERT OR IGNORE（SQLite 方言）-> ON CONFLICT DO NOTHING： # 无冲突目标即可，任何唯一冲突都跳过，不改已有行。2026-10-26 # 实弹教训：GitHub App 安装事件在生产 Postgres 上因这句未 # 翻译直接报语法错，安装记录丢了（delivery 已认领，重试会 # 判 duplicate，事件无法靠重发恢复，只能卸载重装）。 table = m.group(1) cols = m.group(2) values = m.group(3) translated = ( "INSERT INTO %s (%s) VALUES (%s) ON CONFLICT DO NOTHING" % (table, cols, values)) return _convert_placeholders(translated), "normal" translated = _AUTOINCREMENT_RE.sub( "BIGINT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY", sql) translated = _convert_placeholders(translated) # 身份列 INSERT 追加 RETURNING id（供 lastrowid 用）。 m2 = _INSERT_TABLE_RE.match(translated) if m2 and m2.group(1) in _IDENTITY_TABLES \ and "RETURNING" not in translated.upper(): translated = translated.rstrip().rstrip(";") + " RETURNING id" return translated, "normal" def _driver_error_classes(driver): """返回 (integrity类元组, operational类, connection类) 的探测结果。 psycopg（v3）与 psycopg2 的异常都放在各自的 errors 子模块， 且基类名相同；按名探测，探测不到则退化为驱动根模块的 Error / IntegrityError / OperationalError。 """ errors_mod = getattr(driver, "errors", None) if errors_mod is None: driver_name = getattr(driver, "__name__", None) if driver_name: try: errors_mod = __import__( "%s.errors" % driver_name, fromlist=["errors"]) except ImportError: errors_mod = None def _pick(*names): for mod in (errors_mod, driver): if mod is None: continue for name in names: cls = getattr(mod, name, None) if isinstance(cls, type): return cls return None integrity = tuple(c for c in ( _pick("UniqueViolation", "IntegrityError"), _pick("ForeignKeyViolation"), ) if c is not None) operational = _pick("OperationalError") or _pick("Error") return integrity, operational def _map_driver_error(exc, driver): """把驱动异常映射成本模块异常；非驱动异常原样返回。""" integrity, operational = _driver_error_classes(driver) # SQLSTATE / pgcode：psycopg 用 .sqlstate，psycopg2 用 .pgcode。 code = getattr(exc, "sqlstate", None) or getattr(exc, "pgcode", None) if code in ("23505", "23503") or ( integrity and isinstance(exc, integrity)): return PostgresIntegrityError(str(exc)) if code in ("40P01", "55P03"): return PostgresOperationalError(str(exc), locked=True) name = type(exc).__name__ # 按驱动类判定（类名本身不可靠：fake 驱动的类名可能不同， # 但它以 DeadlockDetected / LockNotAvailable 属性挂在驱动上）。 for attr in ("DeadlockDetected", "LockNotAvailable"): cls = getattr(driver, attr, None) if isinstance(cls, type) and isinstance(exc, cls): return PostgresOperationalError(str(exc), locked=True) if name in ("DeadlockDetected", "LockNotAvailable"): return PostgresOperationalError(str(exc), locked=True) if code and str(code).startswith("08"): return PostgresConnectionError(str(exc)) if name in ("OperationalError", "InterfaceError", "DatabaseError") \ or (operational is not None and isinstance(exc, operational)): # 连接类文案单独归类，其余一律按操作错误处理。 msg = str(exc).lower() if "could not connect" in msg or "connection" in msg \ and "refused" in msg: return PostgresConnectionError(str(exc)) return PostgresOperationalError(str(exc)) return exc class PgRow: """Postgres 行：同时支持 row[0] 与 row["col"]、dict(row)。 列名来自游标 description（psycopg3 的 .name、psycopg2 的 元组首项均已归一化）。dict(row) 依赖 keys()，见下文。 """ __slots__ = ("_values", "_columns", "_index") def __init__(self, values, columns): self._values = tuple(values) self._columns = tuple(columns) self._index = {name: i for i, name in enumerate(self._columns)} def __getitem__(self, key): if isinstance(key, str): return self._values[self._index[key]] return self._values[key] def __len__(self): return len(self._values) def __iter__(self): return iter(self._values) def keys(self): return list(self._columns) def __repr__(self): return "PgRow(%r)" % ( {c: v for c, v in zip(self._columns, self._values)},) def _column_names(description): names = [] for col in description or (): if hasattr(col, "name"): names.append(col.name) else: names.append(col[0]) return names class PgCursor: """psycopg 游标的 sqlite3 风格包装。""" def __init__(self, real_cursor, driver, connection): self._cur = real_cursor self._driver = driver self._connection = connection self._fake_rows = None self._fake_columns = () self._returning_row = None self._returning_consumed = False self._lastrowid = None @property def rowcount(self): if self._fake_rows is not None: return len(self._fake_rows) return self._cur.rowcount @property def lastrowid(self): if self._lastrowid is not None: return self._lastrowid if self._returning_row is not None and not self._returning_consumed: self._returning_consumed = True try: row = self._cur.fetchone() except Exception: row = None if row: self._lastrowid = row[0] self._returning_row = None return self._lastrowid def execute(self, sql, params=()): if params is None: params = () try: translated, kind = _translate_sql(sql) except PostgresDatabaseError: raise try: if kind == "pragma_info": # information_schema：列序按 ordinal_position，与 # sqlite PRAGMA table_info 的 (cid, name, ...) 对齐 # 方式是调用方取 row[1] 即 name——fake 行按此形状构造。 info_cur = self._connection._real.cursor() try: info_cur.execute( "SELECT ordinal_position, column_name," " data_type, is_nullable, column_default," " CASE WHEN column_name IN (SELECT kcu.column_name" " FROM information_schema.table_constraints tc" " JOIN information_schema.key_column_usage kcu" " ON tc.constraint_name = kcu.constraint_name" " AND tc.table_schema = kcu.table_schema" " WHERE tc.table_name = %s" " AND tc.constraint_type = 'PRIMARY KEY')" " THEN 1 ELSE 0 END" " FROM information_schema.columns" " WHERE table_name = %s" " ORDER BY ordinal_position", (translated, translated)) raw = info_cur.fetchall() finally: info_cur.close() self._fake_rows = [ PgRow((i, r[1], r[2], 0 if r[3] == "NO" else 1, r[4], r[5]), ("cid", "name", "type", "notnull", "dflt_value", "pk")) for i, r in enumerate(raw)] self._fake_columns = ("cid", "name", "type", "notnull", "dflt_value", "pk") return self if kind == "pragma_ignore": self._fake_rows = [] self._fake_columns = () return self self._fake_rows = None self._returning_row = ( "RETURNING id" in translated.upper() and translated.strip().upper().startswith("INSERT")) self._returning_consumed = False self._lastrowid = None self._cur.execute(translated, tuple(params)) return self except Exception as exc: raise _map_driver_error(exc, self._driver) from exc def executemany(self, sql, seq_of_params): for params in seq_of_params: self.execute(sql, params) return self def fetchall(self): if self._fake_rows is not None: rows, self._fake_rows = self._fake_rows, [] return rows columns = _column_names(self._cur.description) return [PgRow(row, columns) for row in self._cur.fetchall()] def fetchone(self): if self._fake_rows is not None: return self._fake_rows.pop(0) if self._fake_rows else None row = self._cur.fetchone() if row is None: return None return PgRow(row, _column_names(self._cur.description)) def close(self): try: self._cur.close() except Exception: pass class PgConnection: """psycopg 连接的 sqlite3 风格包装。 ``row_factory`` 属性可被调用方事后赋值（如 server._db）， 适配层忽略它——行恒为 PgRow，已支持按名访问。 """ def __init__(self, real_conn, driver): self._real = real_conn self._driver = driver self.row_factory = None def cursor(self): return PgCursor(self._real.cursor(), self._driver, self) def execute(self, sql, params=()): cur = self.cursor() cur.execute(sql, params) return cur def commit(self): try: self._real.commit() except Exception as exc: raise _map_driver_error(exc, self._driver) from exc def rollback(self): try: self._real.rollback() except Exception: pass def close(self): try: self._real.close() except Exception: pass # Postgres schema：与各 store 的 SQLite CREATE TABLE 一一对应， # 表名与列名不改。类型映射见文件头注释；REAL -> DOUBLE PRECISION # （存 epoch 秒浮点），SQLite INTEGER（64 位）-> BIGINT（计数、 # 微秒金额、时间戳整数），小枚举/状态列仍用 INTEGER/ TEXT 原样。 POSTGRES_SCHEMA = ( "CREATE TABLE IF NOT EXISTS users(" "id TEXT PRIMARY KEY, " "username TEXT UNIQUE, " "api_key_hash TEXT, " "created_at DOUBLE PRECISION)", "CREATE TABLE IF NOT EXISTS project_members(" "project_id TEXT, " "user_id TEXT, " "role TEXT, " "PRIMARY KEY(project_id, user_id))", "CREATE TABLE IF NOT EXISTS audit_log(" "id BIGINT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, " "user_id TEXT, " "project_id TEXT, " "action TEXT, " "detail TEXT, " "created_at DOUBLE PRECISION)", # P0 三件：断点与步骤事实源（PG 新库一次建齐；老库走 # server/worker 的 ensure 幂等补表，不丢数据）。 # task_checkpoints：每轮断点快照（worker 镜像线程 + 终局 # 导入写），(task_id, round_no) 唯一；断点丢失时可据此 # 重建 host-state 续跑，不再只信容器里的 host-state.json。 "CREATE TABLE IF NOT EXISTS task_checkpoints(" "task_id TEXT, " "round_no INTEGER, " "thread_id TEXT, " "goal_status TEXT, " "state_sha TEXT, " "verify_out TEXT, " "created_at DOUBLE PRECISION, " "PRIMARY KEY(task_id, round_no))", # task_steps：Run/Step 时间线（worker 侧 claim/resume/final # + host 每轮 round + 镜像导入），timeline 接口直读此表。 "CREATE TABLE IF NOT EXISTS task_steps(" "id TEXT PRIMARY KEY, " "task_id TEXT, " "step_no INTEGER, " "kind TEXT, " "summary TEXT, " "state_sha TEXT, " "created_at DOUBLE PRECISION)", "CREATE INDEX IF NOT EXISTS idx_task_steps_task" " ON task_steps(task_id, step_no)", "CREATE TABLE IF NOT EXISTS projects(" "id TEXT PRIMARY KEY, " "name TEXT, " "owner_id TEXT, " "created_at DOUBLE PRECISION, " "repo_url TEXT, " "auto_pr INTEGER)", # tasks 含 server.migrate_tasks_table 补的列，建表时一次建齐； # 老库迁移仍走 migrate_tasks_table（PRAGMA 翻译后在 PG 上是 no-op）。 "CREATE TABLE IF NOT EXISTS tasks(" "id TEXT PRIMARY KEY, " "project_id TEXT, " "prompt TEXT, " "status TEXT, " "idempotency_key TEXT, " "depends_on TEXT, " "priority INTEGER, " "created_by TEXT, " "created_at DOUBLE PRECISION, " "updated_at DOUBLE PRECISION, " "last_heartbeat DOUBLE PRECISION, " "result_text TEXT, " "error_text TEXT, " "actual_model TEXT, " "started_at DOUBLE PRECISION, " "finished_at DOUBLE PRECISION, " "child_pid BIGINT, " "executor TEXT, " "pr_url TEXT, " "pr_number INTEGER, " "goal TEXT, " "verify_command TEXT, " "goal_max_rounds INTEGER, " "goal_mode INTEGER, " "model TEXT, " "effort TEXT, " "child_model TEXT, " "child_effort TEXT, " "approval TEXT, " "plan INTEGER, " "forked_from TEXT, " "environment_id TEXT, " "resume_message TEXT, " "turn_id TEXT, " "root_turn_id TEXT, " "coordinator_parts TEXT, " "parent_task_id TEXT, " "coordinator_part TEXT, " "goal_resume_pending INTEGER DEFAULT 0, " "UNIQUE(project_id, idempotency_key))", "CREATE INDEX IF NOT EXISTS idx_tasks_status ON tasks(status)", # 附件包与任务附件：与 server._ensure_uploads_tables 同款， # POSTGRES_SCHEMA 在每次初始化时全量执行，新表自动补建。 "CREATE TABLE IF NOT EXISTS uploads(" "id TEXT PRIMARY KEY, " "created_by TEXT, " "created_at DOUBLE PRECISION, " "status TEXT)", "CREATE TABLE IF NOT EXISTS task_attachments(" "task_id TEXT, " "bundle_id TEXT, " "rel_path TEXT, " "size INTEGER, " "PRIMARY KEY(task_id, rel_path))", "CREATE INDEX IF NOT EXISTS idx_attachments_task" " ON task_attachments(task_id)", "CREATE TABLE IF NOT EXISTS task_followups(" "id TEXT PRIMARY KEY, " "task_id TEXT, " "project_id TEXT, " "body TEXT, " "mode TEXT, " "status TEXT, " "created_by TEXT, " "created_at DOUBLE PRECISION, " "updated_at DOUBLE PRECISION)", "CREATE INDEX IF NOT EXISTS idx_followups_task" " ON task_followups(task_id, status)", "CREATE TABLE IF NOT EXISTS environments(" "id TEXT PRIMARY KEY, " "project_id TEXT, " "name TEXT, " "branch TEXT, " "executor TEXT, " "status TEXT, " "error_text TEXT, " "repo_url TEXT, " "created_by TEXT, " "created_at DOUBLE PRECISION, " "updated_at DOUBLE PRECISION)", "CREATE INDEX IF NOT EXISTS idx_env_project" " ON environments(project_id, status)", "CREATE UNIQUE INDEX IF NOT EXISTS idx_env_name" " ON environments(project_id, name)" " WHERE status != 'archived'", "CREATE TABLE IF NOT EXISTS user_quotas(" "user_id TEXT PRIMARY KEY, " "quota_usd_micros BIGINT, " "updated_at DOUBLE PRECISION)", "CREATE TABLE IF NOT EXISTS summaries(" "id BIGINT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, " "project_id TEXT, " "thread_id TEXT, " "goal TEXT, " "conclusion TEXT, " "files_changed TEXT, " "open_issues TEXT, " "status TEXT DEFAULT 'active', " "created_at DOUBLE PRECISION)", # 注意：Postgres 没有 kb_fts（FTS5 虚拟表），KBStore 在 PG 上 # 走 LIKE 分支；不要在这里建任何 FTS 表假装 FTS 存在。 "CREATE TABLE IF NOT EXISTS kb_docs(" "id BIGINT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, " "project_id TEXT, " "path TEXT, " "content TEXT, " "kind TEXT, " "created_at DOUBLE PRECISION)", "CREATE TABLE IF NOT EXISTS cost_records(" "id BIGINT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, " "task_id TEXT, " "project_id TEXT, " "model TEXT, " "input_tokens BIGINT, " "output_tokens BIGINT, " "cached_input_tokens BIGINT DEFAULT 0, " "total_tokens BIGINT, " "usd_micros BIGINT, " "source TEXT, " "fallback_count INTEGER DEFAULT 0, " "completions BIGINT DEFAULT 0, " "day TEXT, " "created_at DOUBLE PRECISION)", "CREATE TABLE IF NOT EXISTS oauth_states(" "state TEXT PRIMARY KEY, " "provider TEXT, " "redirect_uri TEXT, " "created_at DOUBLE PRECISION)", "CREATE TABLE IF NOT EXISTS slack_users(" "slack_id TEXT PRIMARY KEY, " "user_id TEXT, " "created_at DOUBLE PRECISION)", "CREATE TABLE IF NOT EXISTS push_subscriptions(" "id TEXT PRIMARY KEY, " "project_id TEXT NOT NULL, " "user_id TEXT NOT NULL, " "channel TEXT NOT NULL, " "target TEXT NOT NULL, " "preference TEXT NOT NULL, " "created_at DOUBLE PRECISION, " "UNIQUE(project_id, user_id, channel, target))", # System-level identity (cloud/identity.py): TEXT primary keys, # so no IDENTITY sequence to keep in sync in migrations.py. # DDL must match cloud/identity.py CREATE TABLE statements. "CREATE TABLE IF NOT EXISTS user_profiles(" "user_id TEXT PRIMARY KEY, " "display_name TEXT, " "timezone TEXT, " "preferences TEXT, " "updated_at DOUBLE PRECISION)", "CREATE TABLE IF NOT EXISTS identity_memories(" "id TEXT PRIMARY KEY, " "owner_user_id TEXT, " "scope TEXT, " "project_id TEXT, " "kind TEXT, " "content TEXT, " "source_thread_id TEXT, " "status TEXT DEFAULT 'active', " "superseded_by TEXT, " "created_at DOUBLE PRECISION, " "updated_at DOUBLE PRECISION)", # Auto memory (cloud/auto_memory.py): TEXT primary keys, # so no IDENTITY sequence to keep in sync in migrations.py. "CREATE TABLE IF NOT EXISTS memory_settings(" "user_id TEXT PRIMARY KEY, " "generate_memories INTEGER DEFAULT 0, " "use_memories INTEGER DEFAULT 0, " "disable_on_external_context INTEGER DEFAULT 1, " "min_rate_limit_remaining_percent INTEGER DEFAULT 20, " "idle_seconds INTEGER DEFAULT 900, " "min_duration_seconds DOUBLE PRECISION DEFAULT 30, " "min_result_chars INTEGER DEFAULT 200, " "extract_model TEXT, " "consolidation_model TEXT, " "updated_at DOUBLE PRECISION)", "CREATE TABLE IF NOT EXISTS task_memory_flags(" "task_id TEXT PRIMARY KEY, " "use_memories INTEGER, " "generate_memories INTEGER, " "updated_at DOUBLE PRECISION)", "CREATE TABLE IF NOT EXISTS memory_extraction_log(" "task_id TEXT PRIMARY KEY, " "owner_user_id TEXT, " "project_id TEXT, " "status TEXT, " "reason TEXT, " "memory_id TEXT, " "created_at DOUBLE PRECISION, " "processed_at DOUBLE PRECISION)", # kill_requests：跨 worker 真 kill 的意图表。取消/sweep 只改库 # 状态杀不掉 owner 已死时遗留的容器；写一条 pending 请求后， # 任何活着的 worker 在 run loop 的收殓器里按任务派生容器名 # 在本机 docker 杀掉并标 done。同一任务只留一条 pending # （request_kill 的 SELECT 守卫，见 server.py；并发重复写 # 撞上时多出的行同样幂等终结，不会多杀——杀本身按名幂等）。 "CREATE TABLE IF NOT EXISTS kill_requests(" "id TEXT PRIMARY KEY, " "task_id TEXT, " "requested_by TEXT, " "status TEXT, " "requested_at DOUBLE PRECISION, " "done_at DOUBLE PRECISION)", ) def init_postgres_schema(conn): """在给定适配连接上建出全部 Postgres 表（幂等），并提交。 conn 须是 PgConnection（或具同名 execute/commit 的适配连接）； SQLite 连接传进来会在首条 DDL 的 IDENTITY 语法上报错—— 调用方须按后端选对连接，不要把 SQLite 连接混进来。 建表后还给老库 projects 表补 repo_url 列（新库建表时已含）： PRAGMA table_info 经适配层翻译成 information_schema 查询， 调用方取 row[1] 即列名；缺列时 ALTER TABLE 补列，已迁移的 库是 no-op。server.init_db 的 Postgres 分支还会再跑一次 migrate_projects_table，两次补列互不冲突（第二次是 no-op）。 """ for statement in POSTGRES_SCHEMA: conn.execute(statement) cols = {row[1] for row in conn.execute( "PRAGMA table_info(projects)").fetchall()} if "repo_url" not in cols: conn.execute("ALTER TABLE projects ADD COLUMN repo_url TEXT") if "auto_pr" not in cols: conn.execute("ALTER TABLE projects ADD COLUMN auto_pr INTEGER") tcols = {row[1] for row in conn.execute( "PRAGMA table_info(tasks)").fetchall()} for _name, _type in (("executor", "TEXT"), ("pr_url", "TEXT"), ("pr_number", "INTEGER"), ("worktree_host", "TEXT"), ("goal_mode", "INTEGER"), ("environment_id", "TEXT"), ("child_model", "TEXT"), ("child_effort", "TEXT"), ("resume_message", "TEXT")): if _name not in tcols: conn.execute( "ALTER TABLE tasks ADD COLUMN %s %s" % (_name, _type)) conn.commit() class PostgresDatabase(Database): """Database 的 Postgres 实现（已接通真实 Postgres）。 连接参数全部走环境变量（见 POSTGRES_ENV_VARS），显式参数优先 于环境变量；未给出的参数保持 None（传给驱动时省略，不硬编码 任何地址或密码）。初始化时即探测驱动：缺驱动立即抛 PostgresDriverMissingError。驱动分 psycopg（v3，优先）与 psycopg2，连接参数名两者通用（host/port/dbname/user/password）。 本类的 __repr__ 绝不输出 password。 """ def __init__(self, host=None, port=None, dbname=None, user=None, password=None): env = POSTGRES_ENV_VARS self.host = host if host is not None else os.environ.get(env["host"]) raw_port = port if port is not None else os.environ.get(env["port"]) if raw_port is None or raw_port == "": self.port = None else: try: self.port = int(raw_port) except (TypeError, ValueError): raise ValueError( "环境变量 %s 须为整数端口，当前值无法解析" % env["port"]) self.dbname = (dbname if dbname is not None else os.environ.get(env["dbname"])) self.user = (user if user is not None else os.environ.get(env["user"])) # 只保存、不输出、不记录日志；调用方负责传入方式的安全。 self.password = (password if password is not None else os.environ.get(env["password"])) # 驱动探测在初始化时完成：缺驱动这里就报错，不等到首次查询。 self.driver_name, self._driver = _load_postgres_driver() def __repr__(self): return ("PostgresDatabase(host=%r, port=%r, dbname=%r, user=%r)" % (self.host, self.port, self.dbname, self.user)) def _connect_kwargs(self): kwargs = {} if self.host is not None: kwargs["host"] = self.host if self.port is not None: kwargs["port"] = self.port if self.dbname is not None: kwargs["dbname"] = self.dbname if self.user is not None: kwargs["user"] = self.user if self.password is not None: kwargs["password"] = self.password return kwargs def connect(self): try: real = self._driver.connect(**self._connect_kwargs()) except Exception as exc: mapped = _map_driver_error(exc, self._driver) if isinstance(mapped, PostgresDatabaseError): raise mapped from exc raise PostgresConnectionError(str(exc)) from exc return PgConnection(real, self._driver) def execute_with_retry(self, conn, sql, params=(), max_retries=DEFAULT_MAX_RETRIES, backoff_seconds=DEFAULT_BACKOFF_SECONDS): attempts = 0 while True: try: return conn.execute(sql, params) except PostgresOperationalError as exc: if not exc.locked or attempts >= max_retries: raise # Postgres 里失败语句会让事务进入 aborted 状态， # 不 rollback 就无法重试下一条语句。 try: conn.rollback() except Exception: pass delay = backoff_seconds[ min(attempts, len(backoff_seconds) - 1)] time.sleep(delay) attempts += 1 def read_all(self, conn, sql, params=()): return conn.execute(sql, params).fetchall() @contextmanager def transaction(self, conn): try: yield conn conn.commit() except Exception: try: conn.rollback() except Exception: pass raise def is_locked_error(self, exc): if isinstance(exc, PostgresOperationalError): return bool(exc.locked) # 驱动原始异常直接传入时也按名/SQLSTATE 判定（不依赖实例化）。 name = type(exc).__name__ code = getattr(exc, "sqlstate", None) or getattr(exc, "pgcode", None) return name in ("DeadlockDetected", "LockNotAvailable") \ or code in ("40P01", "55P03") def create_database(backend=None, db_path=None, **kwargs): """后端工厂：按 backend（或环境变量）返回 Database 实现。 - backend 为 None 时读 ``CLOUD_DB_BACKEND``，未设则默认 sqlite。 - ``sqlite``：返回 SQLiteDatabase（db_path 未给则用默认文件名， 与 server 旧默认值一致）；kwargs 可传 busy_timeout_ms/timeout。 - ``postgres``/``postgresql``：返回 PostgresDatabase；驱动缺失 时初始化即抛 PostgresDriverMissingError（文案含 psycopg）， 绝不回落到 SQLite；kwargs 可传 host/port/dbname/user/password 覆盖环境变量（测试只能用 fake/test 占位）。 - 其它值抛 ValueError，不静默当成 sqlite。 """ resolved = backend if resolved is None: resolved = os.environ.get(BACKEND_ENV_VAR) or DEFAULT_BACKEND resolved = str(resolved).strip().lower() if resolved == "sqlite": path = db_path if db_path is not None else DEFAULT_SQLITE_PATH return SQLiteDatabase(path, **kwargs) if resolved in ("postgres", "postgresql"): return PostgresDatabase(**kwargs) raise ValueError( "未知数据库后端 %r（%s），仅支持 sqlite/postgres" % (resolved, BACKEND_ENV_VAR)) def database_from_env(db_path=None, **kwargs): """create_database 的环境变量版：backend 一律读 CLOUD_DB_BACKEND。""" return create_database(backend=None, db_path=db_path, **kwargs) __all__ = [ "DEFAULT_BUSY_TIMEOUT_MS", "DEFAULT_CONNECT_TIMEOUT_S", "DEFAULT_MAX_RETRIES", "DEFAULT_BACKOFF_SECONDS", "POSTGRES_ENV_VARS", "BACKEND_ENV_VAR", "DEFAULT_BACKEND", "connect", "is_locked_error", "execute_with_retry", "read_all", "transaction", "Database", "SQLiteDatabase", "PostgresDatabase", "PostgresDriverMissingError", "PostgresDatabaseError", "PostgresIntegrityError", "PostgresOperationalError", "PostgresConnectionError", "PgConnection", "PgCursor", "PgRow", "POSTGRES_SCHEMA", "init_postgres_schema", "create_database", "database_from_env", ]