Postgres real backend: psycopg adapter, schema, migration
Replace the PostgresDatabase skeleton with a real implementation: psycopg (v3) first, psycopg2 second, lazy load, missing driver still raises PostgresDriverMissingError with no silent SQLite fallback. A sqlite3-style adapter (PgConnection/PgCursor/PgRow) translates ? placeholders to %s outside string literals, maps UNIQUE/FK violations, lock errors and connection failures to explicit Postgres error types, and handles INSERT OR REPLACE, PRAGMA table_info, AUTOINCREMENT and FTS5 refusal. POSTGRES_SCHEMA covers all business tables with tasks result columns; server init_db uses it under CLOUD_DB_BACKEND=postgres, SQLite path unchanged. migrate_sqlite_to_postgres.py copies rows with ON CONFLICT DO NOTHING and reseeds identity sequences. Tests: cloud 239 (235 pass, 4 skipped without a real server), host 27 pass; real-server tests skip unless CLOUD_TEST_PG_DSN is set.
This commit is contained in:
1 parent
a19e77ee44
commit
dc3f572b2b
10 files changed
+1542
-66
No files matched your search
+79
-14
@@ -21,7 +21,7 @@ cloud/server.py(控制面,ThreadingHTTPServer)
|
||||
├─ 推送订阅与分发:cloud/push.py(订阅表 + 偏好过滤)
|
||||
├─ 认证 / 权限 / 审计:cloud/auth.py(AuthStore)
|
||||
├─ 数据库接口:cloud/db.py(Database 接口 + SQLite 实现,
|
||||
│ WAL + busy_timeout + 有限重试;Postgres 只是骨架,见下)
|
||||
│ WAL + busy_timeout + 有限重试;Postgres 已接通,见下)
|
||||
├─ 并发:cloud/db.py(WAL + busy_timeout + 有限重试)
|
||||
└─ webhook 入站:cloud/webhooks.py(验签 + 解析,只解析不执行)
|
||||
```
|
||||
@@ -126,8 +126,22 @@ server 的连接统一走 `cloud/db.py`:WAL 模式、busy_timeout 5000ms、
|
||||
WAL / busy_timeout / 有界重试;模块级函数 `connect` /
|
||||
`execute_with_retry` / `read_all` / `transaction` /
|
||||
`is_locked_error` 仍可用,行为不变。
|
||||
- `PostgresDatabase`:骨架,不是可用后端。即使驱动已装,其
|
||||
`connect` / 执行 / 查询 / 事务方法仍抛 `NotImplementedError`。
|
||||
- `PostgresDatabase`:真实 Postgres 实现。调用方仍写 sqlite3
|
||||
风格 SQL(`?` 占位符、`row["col"]`、`conn.execute` /
|
||||
`fetchall` / `commit` / `rollback` / `close`),由
|
||||
`PgConnection` / `PgCursor` / `PgRow` 适配层翻译成 Postgres
|
||||
方言(`?` -> `%s`,字符串字面量里的 `?` 不动;
|
||||
`INSERT OR REPLACE` -> `ON CONFLICT DO UPDATE`;
|
||||
`PRAGMA table_info` -> information_schema;
|
||||
`AUTOINCREMENT` -> `GENERATED BY DEFAULT AS IDENTITY`;
|
||||
FTS5 虚拟表在 PG 上明确报错、KBStore 走 LIKE 分支)。
|
||||
驱动异常映射成 `PostgresIntegrityError` /
|
||||
`PostgresOperationalError` / `PostgresConnectionError`,
|
||||
绝不把驱动异常原样漏成 sqlite3 异常。模块级 `connect()`
|
||||
按 `CLOUD_DB_BACKEND` 分发,因此 server / AuthStore /
|
||||
MemoryStore / KBStore / CostStore 等沿用 `connect()` 的
|
||||
调用点无需重构即可走 Postgres;但 `backup.py` 的 SQLite
|
||||
在线备份 API 仍只适用于 SQLite,Postgres 备份请用 pg_dump。
|
||||
- 工厂 `create_database(backend, db_path)` / `database_from_env()`:
|
||||
两个后端提供同一组接口方法与签名,见
|
||||
`cloud/tests/test_db_backend.py` 的接口契约测试。
|
||||
@@ -150,14 +164,65 @@ psycopg,写明“需安装 psycopg”),绝不静默回落到 SQLite,
|
||||
也绝不假装已接通 Postgres。`create_server` 在建表之前就做这
|
||||
一步校验,所以失败时不会先建出一个 SQLite 库假装可用。
|
||||
|
||||
> 高并发换 Postgres 是下一阶段的运维项,不是本轮代码一改就
|
||||
> 自动变生产。接口层 + 骨架只解决了“以后能换”,真正上生产还
|
||||
> 需要:真实 Postgres 部署(含备份、连接数与权限)、安装
|
||||
> psycopg 驱动、写迁移脚本并把现有 SQLite 表结构与数据迁过去、
|
||||
> 把 server / AuthStore / MemoryStore / KBStore / CostStore 等
|
||||
> 目前仍直连 SQLite 的调用点逐一接到新后端并重测。在这些完成
|
||||
> 之前,即使把 `CLOUD_DB_BACKEND` 设成 `postgres`,server 也只会
|
||||
> 在启动时明确报错,不会跑起来——这是有意为之的诚实边界。
|
||||
### Postgres 部署步骤(不要只改一个环境变量)
|
||||
|
||||
改 `CLOUD_DB_BACKEND=postgres` 本身不会迁移任何数据,也不会
|
||||
建库——只改它、数据还在旧 SQLite 文件里,新库是空的。完整
|
||||
步骤是:
|
||||
|
||||
1. 装驱动(psycopg v3 优先):
|
||||
```bash
|
||||
pip install "psycopg[binary]"
|
||||
# 其次才用 psycopg2:pip install psycopg2-binary
|
||||
```
|
||||
2. 在 Postgres 里建库与用户(示例库名与用户名请按实际部署改,
|
||||
密码绝不写进仓库文件):
|
||||
```sql
|
||||
CREATE DATABASE cloud_control_plane;
|
||||
CREATE USER cloud_app WITH PASSWORD '<仅在此处输入>';
|
||||
GRANT ALL PRIVILEGES ON DATABASE cloud_control_plane TO cloud_app;
|
||||
```
|
||||
注意:表由应用启动时 `init_postgres_schema` 建(
|
||||
`cloud/db.py` 的 `POSTGRES_SCHEMA`,与 SQLite 表名/列名
|
||||
一一对应,含 tasks 的 6 个结果列;KB 在 PG 上无 FTS 表,
|
||||
走 LIKE 搜索)。若数据库用户不是库 owner,另需对
|
||||
public schema 授权(`GRANT ALL ON SCHEMA public TO cloud_app`),
|
||||
否则建表会报权限错——这是部署侧问题,不是代码 bug。
|
||||
3. 设环境变量并启动(`--db` 参数在 PG 模式下被忽略,连接
|
||||
参数只走环境变量;`--db-backend postgres` 与环境变量等价):
|
||||
```bash
|
||||
export CLOUD_DB_BACKEND=postgres
|
||||
export CLOUD_DB_HOST=127.0.0.1 CLOUD_DB_PORT=5432
|
||||
export CLOUD_DB_NAME=cloud_control_plane
|
||||
export CLOUD_DB_USER=cloud_app
|
||||
export CLOUD_DB_PASSWORD='<仅在此处输入>'
|
||||
python3 -m cloud.server --host 127.0.0.1 --port 8765
|
||||
```
|
||||
4. 迁旧数据(可选,须在启动建表之后、切流量之前跑;跑两次
|
||||
不重复插,第二次各表 inserted 应为 0):
|
||||
```bash
|
||||
python3 -m cloud.migrate_sqlite_to_postgres \
|
||||
--sqlite /path/to/cloud-control-plane.db
|
||||
# 输出示例:projects: 3 / tasks: 12 / ... / total inserted: N
|
||||
# (只打印表名与行数,不打印任何密码或密钥)
|
||||
```
|
||||
迁移脚本逻辑在 `cloud/migrations.py`(
|
||||
`migrate_sqlite_to_postgres`),身份列的序列会在迁完后
|
||||
拨到当前最大 id,避免后续自动生成撞主键。
|
||||
5. 回退到 SQLite:把 `CLOUD_DB_BACKEND` 改回 `sqlite`(或
|
||||
删掉它,默认即 sqlite),并把 `--db` 指回原来的 SQLite
|
||||
文件即可。但注意:切到 Postgres 之后新写的数据只在
|
||||
Postgres 里,SQLite 文件里没有——回退不会自动把 PG 数据
|
||||
导回 SQLite;如需回迁,须另写 PG -> SQLite 脚本(本轮
|
||||
未提供,不要假装存在)。
|
||||
6. 备份:Postgres 侧用 `pg_dump`,不要用 `cloud/backup.py`
|
||||
(那是 SQLite 在线备份 API,只适用于 SQLite 文件)。
|
||||
|
||||
> 诚实边界(仍需在真实服务器上验证,见下文“未验证清单”):
|
||||
> 本机 VM 既无 psycopg 驱动、也无 Postgres 服务器,因此真实
|
||||
> PG 路径在本机只能 skip(见测试一节),不能谎称已在真实
|
||||
> PG 上跑通过。已验证的是:适配层翻译逻辑(fake 驱动)、
|
||||
> schema 语句覆盖、SQLite 全量回归不坏。
|
||||
|
||||
| 方法与路径 | 权限 | 说明 |
|
||||
| --- | --- | --- |
|
||||
@@ -277,8 +342,8 @@ interrupted 且终态任务 cancel 返回 400、心跳后 sweep 不扫刚心跳
|
||||
端点(state + 授权 URL + 回调),但真实应用注册、回调公网地址、
|
||||
已有账号绑定仍缺;推送已接成订阅 + 分发端点,但 Web Push /
|
||||
APNs / FCM 的真实凭证与发送端仍缺;单机 SQLite 在高并发写下
|
||||
仍有上限,生产量上来后需换 Postgres(接口层与骨架已就绪,
|
||||
但真实部署、驱动安装与迁移脚本仍缺,见上文“数据库后端选择”);
|
||||
仍有上限,生产量上来后需换 Postgres(Postgres 后端已接通,
|
||||
但真实服务器部署与验证仍缺,见上文“数据库后端选择”);
|
||||
worker 执行已接通(见下文“worker 常驻与真取消”),云端执行、
|
||||
真实通知与真实模型回合仍受 key 供给与云账号卡点限制。
|
||||
|
||||
@@ -311,7 +376,7 @@ interrupted 且终态任务 cancel 返回 400、心跳后 sweep 不扫刚心跳
|
||||
失败明确报错不静默;备份文件本身就是可打开的 SQLite 库。
|
||||
- DB 接口迁移:server 及 auth/memory/kb/cost/oauth/push/webhooks
|
||||
的 Store 连接已统一走 cloud.db 的 WAL + busy_timeout 连接
|
||||
(Postgres 后端仍是骨架,驱动缺失时启动即失败、不先建 SQLite
|
||||
(Postgres 后端已接通,驱动缺失时启动即失败、不先建 SQLite
|
||||
文件,这一点迁移后未退化)。
|
||||
|
||||
## Docker 沙箱执行器(本轮新增)
|
||||
|
||||
+670
-22
@@ -6,7 +6,8 @@
|
||||
``is_locked_error`` / ``execute_with_retry`` / ``read_all`` /
|
||||
``transaction``。server.py 与旧测试仍走这条路,行为不变。
|
||||
2. ``Database`` 接口 + 两个实现:``SQLiteDatabase``(现有 SQLite
|
||||
行为的类封装)与 ``PostgresDatabase``(骨架,见下文)。
|
||||
行为的类封装)与 ``PostgresDatabase``(真实 Postgres 实现,
|
||||
见下文的适配层)。
|
||||
新代码应走接口 + ``create_database`` 工厂,不要再直接拼
|
||||
sqlite3 连接。
|
||||
|
||||
@@ -34,13 +35,27 @@ memory.py 的 summaries 完全一致,因为本模块根本不建表)。
|
||||
- Postgres 连接参数全部走环境变量,不硬编码地址或密码:
|
||||
``CLOUD_DB_HOST``、``CLOUD_DB_PORT``、``CLOUD_DB_NAME``、
|
||||
``CLOUD_DB_USER``、``CLOUD_DB_PASSWORD``。
|
||||
- Postgres 骨架的诚实边界:标准库没有 Postgres 驱动,
|
||||
``PostgresDatabase`` 初始化时找不到 psycopg/psycopg2 会立即
|
||||
明确报错(``PostgresDriverMissingError``,文案含 psycopg),
|
||||
绝不静默回落到 SQLite;即使驱动已安装,其连接/执行/查询/
|
||||
事务方法目前仍是 ``NotImplementedError`` 骨架,尚未接通真实
|
||||
Postgres——换生产数据库还需要真实部署、驱动安装与迁移脚本,
|
||||
不是改一个环境变量就自动完成。
|
||||
- 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
|
||||
@@ -74,6 +89,22 @@ 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:读不阻塞写、写不阻塞读,是多线程
|
||||
@@ -273,21 +304,577 @@ def _load_postgres_driver():
|
||||
)
|
||||
|
||||
|
||||
_POSTGRES_NOT_IMPLEMENTED = (
|
||||
"PostgresDatabase 是骨架,尚未实现:需要真实 Postgres 部署、"
|
||||
"psycopg 驱动安装与迁移脚本后才能接通,当前不假装已接通 Postgres。"
|
||||
# ----------------------------------------------------------------------
|
||||
# Postgres 适配层:把 sqlite3 风格调用翻译成 psycopg 调用
|
||||
# ----------------------------------------------------------------------
|
||||
# 不重构所有 store 的前提:现有代码到处用 ``?`` 占位符、
|
||||
# ``row["col"]``、``conn.execute().fetchall()``、``commit`` /
|
||||
# ``rollback`` / ``close``、``cursor.lastrowid``、``cursor.rowcount``,
|
||||
# 以及少量 SQLite 方言(``INSERT OR REPLACE``、``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 两张已知主键表,不会走到退化分支)。
|
||||
# - ``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",),
|
||||
}
|
||||
|
||||
_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)
|
||||
_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"
|
||||
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)",
|
||||
"CREATE TABLE IF NOT EXISTS projects("
|
||||
"id TEXT PRIMARY KEY, "
|
||||
"name TEXT, "
|
||||
"owner_id TEXT, "
|
||||
"created_at DOUBLE PRECISION)",
|
||||
# tasks 含 server.migrate_tasks_table 补的那 6 列,建表时一次建齐;
|
||||
# 老库迁移仍走 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, "
|
||||
"UNIQUE(project_id, idempotency_key))",
|
||||
"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, "
|
||||
"total_tokens BIGINT, "
|
||||
"usd_micros BIGINT, "
|
||||
"source TEXT, "
|
||||
"fallback_count INTEGER 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))",
|
||||
)
|
||||
|
||||
|
||||
def init_postgres_schema(conn):
|
||||
"""在给定适配连接上建出全部 Postgres 表(幂等),并提交。
|
||||
|
||||
conn 须是 PgConnection(或具同名 execute/commit 的适配连接);
|
||||
SQLite 连接传进来会在首条 DDL 的 IDENTITY 语法上报错——
|
||||
调用方须按后端选对连接,不要把 SQLite 连接混进来。
|
||||
"""
|
||||
for statement in POSTGRES_SCHEMA:
|
||||
conn.execute(statement)
|
||||
conn.commit()
|
||||
|
||||
|
||||
class PostgresDatabase(Database):
|
||||
"""Database 的 Postgres 骨架(未接通真实 Postgres)。
|
||||
"""Database 的 Postgres 实现(已接通真实 Postgres)。
|
||||
|
||||
连接参数全部走环境变量(见 POSTGRES_ENV_VARS),显式参数优先
|
||||
于环境变量;未给出的参数保持 None,不硬编码任何地址或密码。
|
||||
初始化时即探测驱动:缺驱动立即抛 PostgresDriverMissingError,
|
||||
因此 server 选 postgres 后端时在启动阶段就失败,不会在运行时
|
||||
中途才炸。即使驱动已安装,connect/execute/read/transaction
|
||||
目前仍抛 NotImplementedError——骨架不是可用后端。
|
||||
于环境变量;未给出的参数保持 None(传给驱动时省略,不硬编码
|
||||
任何地址或密码)。初始化时即探测驱动:缺驱动立即抛
|
||||
PostgresDriverMissingError。驱动分 psycopg(v3,优先)与
|
||||
psycopg2,连接参数名两者通用(host/port/dbname/user/password)。
|
||||
本类的 __repr__ 绝不输出 password。
|
||||
"""
|
||||
|
||||
@@ -318,22 +905,74 @@ class PostgresDatabase(Database):
|
||||
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):
|
||||
raise NotImplementedError(_POSTGRES_NOT_IMPLEMENTED)
|
||||
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):
|
||||
raise NotImplementedError(_POSTGRES_NOT_IMPLEMENTED)
|
||||
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=()):
|
||||
raise NotImplementedError(_POSTGRES_NOT_IMPLEMENTED)
|
||||
return conn.execute(sql, params).fetchall()
|
||||
|
||||
@contextmanager
|
||||
def transaction(self, conn):
|
||||
raise NotImplementedError(_POSTGRES_NOT_IMPLEMENTED)
|
||||
try:
|
||||
yield conn
|
||||
conn.commit()
|
||||
except Exception:
|
||||
try:
|
||||
conn.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
raise
|
||||
|
||||
def is_locked_error(self, exc):
|
||||
raise NotImplementedError(_POSTGRES_NOT_IMPLEMENTED)
|
||||
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):
|
||||
@@ -384,6 +1023,15 @@ __all__ = [
|
||||
"SQLiteDatabase",
|
||||
"PostgresDatabase",
|
||||
"PostgresDriverMissingError",
|
||||
"PostgresDatabaseError",
|
||||
"PostgresIntegrityError",
|
||||
"PostgresOperationalError",
|
||||
"PostgresConnectionError",
|
||||
"PgConnection",
|
||||
"PgCursor",
|
||||
"PgRow",
|
||||
"POSTGRES_SCHEMA",
|
||||
"init_postgres_schema",
|
||||
"create_database",
|
||||
"database_from_env",
|
||||
]
|
||||
+3
-2
@@ -18,6 +18,7 @@ Design notes:
|
||||
|
||||
import sqlite3
|
||||
from cloud.db import connect as _db_connect
|
||||
from cloud.db import PostgresOperationalError as _PgOperationalError
|
||||
import time
|
||||
|
||||
_KINDS = ("file", "readme", "doc", "task", "review")
|
||||
@@ -47,7 +48,7 @@ class KBStore:
|
||||
"CREATE VIRTUAL TABLE IF NOT EXISTS kb_fts USING fts5(content)"
|
||||
)
|
||||
self._fts = True
|
||||
except sqlite3.OperationalError:
|
||||
except (sqlite3.OperationalError, _PgOperationalError):
|
||||
self._fts = False
|
||||
self.conn.commit()
|
||||
|
||||
@@ -129,7 +130,7 @@ class KBStore:
|
||||
# "用户登录流程") returns zero FTS hits. LIKE substring
|
||||
# search recovers those cases; the project_id filter
|
||||
# still applies, so isolation is unchanged.
|
||||
except sqlite3.OperationalError:
|
||||
except (sqlite3.OperationalError, _PgOperationalError):
|
||||
# Query syntax FTS5 cannot parse, or a broken FTS table:
|
||||
# fall through to the LIKE path rather than crashing.
|
||||
pass
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
"""把 SQLite 控制面库迁移到 Postgres 的命令行入口。
|
||||
|
||||
用法(Postgres 连接参数只走 CLOUD_DB_* 环境变量,本脚本
|
||||
不接受、不打印任何密码或密钥):
|
||||
|
||||
CLOUD_DB_HOST=... CLOUD_DB_PORT=5432 CLOUD_DB_NAME=... \
|
||||
CLOUD_DB_USER=... CLOUD_DB_PASSWORD=... \
|
||||
python3 -m cloud.migrate_sqlite_to_postgres --sqlite <sqlite 文件>
|
||||
|
||||
输出:每张表本次插入的行数(第二次跑应全为 0,即幂等,
|
||||
不重复插)。驱动缺失时明确报错,绝不静默回落到 SQLite;
|
||||
目标库须事先存在(本脚本不建库,建库见 cloud/README.md)。
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
|
||||
from cloud.migrations import migrate_sqlite_to_postgres
|
||||
|
||||
|
||||
def main(argv=None):
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Migrate a SQLite control-plane DB to Postgres.")
|
||||
parser.add_argument("--sqlite", required=True,
|
||||
help="source SQLite file path")
|
||||
args = parser.parse_args(argv)
|
||||
result = migrate_sqlite_to_postgres(args.sqlite)
|
||||
total = 0
|
||||
for table, count in result.items():
|
||||
print("%s: %d" % (table, count))
|
||||
total += count
|
||||
print("total inserted: %d" % total)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,120 @@
|
||||
"""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)。
|
||||
- kb_fts(SQLite FTS5 虚拟表)不迁移:Postgres 没有 FTS5,
|
||||
KBStore 在 PG 上走 LIKE 分支,kb_docs 的 content 已足够。
|
||||
- sqlite_sequence(AUTOINCREMENT 内部表)不迁移。
|
||||
- 不打印、不返回任何密码或密钥:函数只返回 {表: inserted},
|
||||
CLI 只打印表名与行数。
|
||||
- Postgres 连接参数只走 CLOUD_DB_* 环境变量(见 cloud.db),
|
||||
本模块不读、不拼连接串、不接受命令行密码参数。
|
||||
"""
|
||||
|
||||
import os
|
||||
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",
|
||||
)
|
||||
|
||||
_SKIP_TABLES = frozenset({"kb_fts", "sqlite_sequence"})
|
||||
|
||||
|
||||
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]
|
||||
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 ("audit_log", "summaries", "kb_docs",
|
||||
"cost_records") 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",
|
||||
"migrate_sqlite_to_postgres",
|
||||
]
|
||||
+2
-1
@@ -33,6 +33,7 @@
|
||||
|
||||
import sqlite3
|
||||
from cloud.db import connect as _db_connect
|
||||
from cloud.db import PgConnection as _PgConnection
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import contextmanager
|
||||
@@ -70,7 +71,7 @@ MISSING_CREDENTIAL_REASONS = {
|
||||
|
||||
|
||||
def _is_connection(db):
|
||||
return isinstance(db, sqlite3.Connection)
|
||||
return isinstance(db, (sqlite3.Connection, _PgConnection))
|
||||
|
||||
|
||||
@contextmanager
|
||||
|
||||
+35
-11
@@ -29,7 +29,7 @@ audit_log 三张表;本 server 用 projects/tasks 两张表,表名不冲突
|
||||
|
||||
启动:
|
||||
python3 -m cloud.server [--host 127.0.0.1] [--port 8765] [--db <路径>]
|
||||
后端由 CLOUD_DB_BACKEND 选(sqlite 默认;postgres 目前只是骨架,
|
||||
后端由 CLOUD_DB_BACKEND 选(sqlite 默认;postgres 已接通,
|
||||
驱动缺失时启动即明确报错,见 cloud/db.py)。
|
||||
|
||||
测试用 create_server(db_path, host, port) 只返回 ThreadingHTTPServer
|
||||
@@ -50,6 +50,9 @@ from cloud.auth import AuthStore
|
||||
from cloud.cost import CostStore
|
||||
from cloud.db import connect as _db_connect
|
||||
from cloud.db import create_database as _create_database
|
||||
from cloud.db import init_postgres_schema as _init_postgres_schema
|
||||
from cloud.db import PostgresIntegrityError as _PgIntegrityError
|
||||
from cloud.db import PostgresOperationalError as _PgOperationalError
|
||||
from cloud.kb import KBStore
|
||||
from cloud.memory import MemoryStore
|
||||
from cloud.oauth import (
|
||||
@@ -79,10 +82,14 @@ DEFAULT_SWEEP_THRESHOLD_SECONDS = 300
|
||||
|
||||
|
||||
def _connect(db_path):
|
||||
"""每个请求新建一个连接;走 cloud.db 的 WAL + busy_timeout。
|
||||
"""每个请求新建一个连接;后端按 CLOUD_DB_BACKEND 分发。
|
||||
|
||||
单机 SQLite 的上限仍在:WAL + 有限重试只能缓解多线程写冲突,
|
||||
高并发要换 Postgres(见 cloud/README.md)。
|
||||
cloud.db.connect() 本身已按环境变量分发:sqlite 时是
|
||||
WAL + busy_timeout 的 SQLite 连接;postgres 时是 psycopg
|
||||
适配连接(db_path 被忽略,连接参数只走 CLOUD_DB_*
|
||||
环境变量,驱动缺失时明确报错,绝不回落)。因此 server
|
||||
自身的表与 AuthStore / MemoryStore 等沿用 connect() 的
|
||||
调用点在 PG 模式下同样走 Postgres,无需逐一改写。
|
||||
"""
|
||||
return _db_connect(str(db_path))
|
||||
|
||||
@@ -111,10 +118,26 @@ def migrate_tasks_table(conn):
|
||||
"ALTER TABLE tasks ADD COLUMN %s %s" % (name, coltype))
|
||||
|
||||
|
||||
def _is_postgres_backend():
|
||||
backend = os.environ.get("CLOUD_DB_BACKEND") or "sqlite"
|
||||
return str(backend).strip().lower() in ("postgres", "postgresql")
|
||||
|
||||
|
||||
def init_db(db_path):
|
||||
"""创建 server 自己的 projects/tasks 表(AuthStore 的表由其自建)。"""
|
||||
"""创建 server 自己的 projects/tasks 表(AuthStore 的表由其自建)。
|
||||
|
||||
Postgres 模式:建 cloud.db.POSTGRES_SCHEMA 的全部表(其中
|
||||
projects/tasks 已含 migrate_tasks_table 的 6 列),再跑一次
|
||||
migrate_tasks_table——PRAGMA 翻译后在 PG 上是 no-op,老库
|
||||
迁移仍走同一函数。SQLite 模式:原行为不变。
|
||||
"""
|
||||
conn = _connect(db_path)
|
||||
try:
|
||||
if _is_postgres_backend():
|
||||
_init_postgres_schema(conn)
|
||||
migrate_tasks_table(conn)
|
||||
conn.commit()
|
||||
return
|
||||
conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS projects("
|
||||
"id TEXT PRIMARY KEY, "
|
||||
@@ -411,7 +434,7 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
||||
"owner_id": user_id,
|
||||
"created_at": now,
|
||||
}
|
||||
except sqlite3.IntegrityError:
|
||||
except (sqlite3.IntegrityError, _PgIntegrityError):
|
||||
self._send_json(409, {"error": "project exists"})
|
||||
return
|
||||
finally:
|
||||
@@ -889,7 +912,7 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
||||
json.dumps(depends_on), priority, user_id, now, now),
|
||||
)
|
||||
conn.commit()
|
||||
except sqlite3.IntegrityError:
|
||||
except (sqlite3.IntegrityError, _PgIntegrityError):
|
||||
# 并发下同幂等键抢建:回读已有任务,不建重复。
|
||||
row = conn.execute(
|
||||
"SELECT * FROM tasks WHERE project_id = ?"
|
||||
@@ -1116,7 +1139,7 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
||||
"SELECT project_id, source, total_tokens, usd_micros,"
|
||||
" fallback_count FROM cost_records"
|
||||
).fetchall()
|
||||
except sqlite3.OperationalError:
|
||||
except (sqlite3.OperationalError, _PgOperationalError):
|
||||
cost_rows = []
|
||||
actual_tokens = estimated_tokens = 0
|
||||
actual_usd = estimated_usd = 0
|
||||
@@ -1325,11 +1348,12 @@ def create_server(db_path, host="127.0.0.1", port=8765, backend=None):
|
||||
backend 为 None 时读环境变量 CLOUD_DB_BACKEND,未设则默认
|
||||
sqlite(旧行为不变)。选 postgres 时在建表之前就调用工厂:
|
||||
驱动缺失则初始化即抛明确错误(文案含 psycopg),启动即失败,
|
||||
绝不在运行时中途才炸、也绝不静默回落到 SQLite;Postgres
|
||||
目前只是骨架,即使驱动已装也尚未接通,见 cloud/db.py 与 README。
|
||||
绝不在运行时中途才炸、也绝不静默回落到 SQLite;驱动已装
|
||||
且能连上服务器时 Postgres 即为可用后端(适配层见
|
||||
cloud/db.py,部署步骤见 cloud/README.md)。
|
||||
"""
|
||||
db_path = str(db_path)
|
||||
# 后端校验必须在 init_db 之前:postgres 骨架不能先建 SQLite 表
|
||||
# 后端校验必须在 init_db 之前:postgres 模式不能先建 SQLite 表
|
||||
# 再失败,那会假装 sqlite 已可用。
|
||||
_create_database(backend=backend, db_path=db_path)
|
||||
init_db(db_path)
|
||||
|
||||
@@ -207,22 +207,94 @@ class PostgresSkeletonTestCase(_EnvCleanMixin, unittest.TestCase):
|
||||
self.assertEqual(db2.host, "explicit-host")
|
||||
self.assertEqual(db2.dbname, "env-db")
|
||||
|
||||
def test_postgres_methods_are_skeleton_not_fake_success(self):
|
||||
# 绕过 __init__(驱动探测)构造骨架实例,验证方法体本身
|
||||
# 就是明确的 NotImplementedError,而不是假装成功。
|
||||
skeleton = object.__new__(PostgresDatabase)
|
||||
for name in ("connect", "read_all", "transaction",
|
||||
"is_locked_error"):
|
||||
method = getattr(skeleton, name)
|
||||
with self.assertRaises(NotImplementedError):
|
||||
if name in ("connect", "transaction"):
|
||||
method(None) if name == "transaction" else method()
|
||||
elif name == "is_locked_error":
|
||||
method(ValueError("x"))
|
||||
else:
|
||||
method(None, "SELECT 1")
|
||||
with self.assertRaises(NotImplementedError):
|
||||
skeleton.execute_with_retry(None, "SELECT 1", ())
|
||||
def test_postgres_methods_are_implemented_via_fake_driver(self):
|
||||
# 骨架已由真实实现取代:用 fake 驱动(不连真实 Postgres)
|
||||
# 验证 connect / read_all / transaction / is_locked_error
|
||||
# 均可用,且驱动唯一冲突被映射成模块内完整性错误。
|
||||
from cloud.db import (
|
||||
PgConnection,
|
||||
PostgresIntegrityError,
|
||||
PostgresOperationalError,
|
||||
)
|
||||
|
||||
class _FakeUnique(Exception):
|
||||
pass
|
||||
|
||||
class _FakeCursor:
|
||||
def __init__(self):
|
||||
self.executed = []
|
||||
self.description = [("id",), ("name",)]
|
||||
|
||||
def execute(self, sql, params=()):
|
||||
self.executed.append((sql, params))
|
||||
if "dup" in sql:
|
||||
raise _FakeUnique("duplicate key")
|
||||
return self
|
||||
|
||||
def fetchall(self):
|
||||
return [("t1", "hello")]
|
||||
|
||||
def fetchone(self):
|
||||
return ("t1", "hello")
|
||||
|
||||
@property
|
||||
def rowcount(self):
|
||||
return 1
|
||||
|
||||
def close(self):
|
||||
pass
|
||||
|
||||
class _FakeConn:
|
||||
def __init__(self):
|
||||
self.cursor_obj = _FakeCursor()
|
||||
self.committed = False
|
||||
self.rolled_back = False
|
||||
|
||||
def cursor(self):
|
||||
return self.cursor_obj
|
||||
|
||||
def commit(self):
|
||||
self.committed = True
|
||||
|
||||
def rollback(self):
|
||||
self.rolled_back = True
|
||||
|
||||
def close(self):
|
||||
pass
|
||||
|
||||
class _FakeDriver:
|
||||
UniqueViolation = _FakeUnique
|
||||
OperationalError = Exception
|
||||
Error = Exception
|
||||
|
||||
def __init__(self):
|
||||
self.conn = _FakeConn()
|
||||
|
||||
def connect(self, **kwargs):
|
||||
return self.conn
|
||||
|
||||
fake = _FakeDriver()
|
||||
with mock.patch("cloud.db._load_postgres_driver",
|
||||
return_value=("psycopg", fake)):
|
||||
db = PostgresDatabase(host="test-host", dbname="test-db")
|
||||
conn = db.connect()
|
||||
self.assertIsInstance(conn, PgConnection)
|
||||
rows = db.read_all(conn, "SELECT * FROM t WHERE id = ?",
|
||||
("t1",))
|
||||
self.assertEqual(rows[0]["name"], "hello")
|
||||
self.assertEqual(rows[0][0], "t1")
|
||||
# 唯一冲突 -> 模块内错误,不是驱动原始异常、不是 sqlite 异常。
|
||||
with self.assertRaises(PostgresIntegrityError):
|
||||
db.execute_with_retry(conn, "SELECT dup", ())
|
||||
# 事务回滚:第二条失败后 rollback 被调用且异常原样抛出。
|
||||
with self.assertRaises(PostgresIntegrityError):
|
||||
with db.transaction(conn):
|
||||
conn.execute("SELECT dup", ())
|
||||
self.assertTrue(fake.conn.rolled_back)
|
||||
# 锁判定:本模块锁错误 True,其它错误 False。
|
||||
self.assertTrue(db.is_locked_error(
|
||||
PostgresOperationalError("deadlock", locked=True)))
|
||||
self.assertFalse(db.is_locked_error(ValueError("x")))
|
||||
|
||||
def test_postgres_repr_never_contains_password(self):
|
||||
skeleton = object.__new__(PostgresDatabase)
|
||||
|
||||
@@ -0,0 +1,336 @@
|
||||
"""Postgres 后端单测:适配层(fake 驱动,常跑)+ 真实 PG(skip)。
|
||||
|
||||
分两类,泾渭分明:
|
||||
- FakeDriver 类:不连任何真实 Postgres,验证占位符转换、
|
||||
行按名访问、唯一冲突映射、事务回滚、锁重试、schema 语句
|
||||
覆盖——这些在无 Postgres 的机器(含 CI runner)上也必须通过。
|
||||
- RealPostgres 类:需要真实 Postgres 时才跑,判定条件见
|
||||
_real_pg_params():须显式设置 CLOUD_TEST_PG_DSN(形如
|
||||
postgres://user:pass@host:port/dbname)且主机端口 2 秒内
|
||||
可达,否则 skipTest,绝不 fail。DSN 必须指向一次性测试库,
|
||||
本文件不在其中建生产表之外的任何东西(schema 即业务表)。
|
||||
|
||||
测试里出现的密码一律是 fake/test 占位("test-placeholder"),
|
||||
绝不写真实凭证;DSN 只从环境变量读,测试不打印它。
|
||||
"""
|
||||
|
||||
import os
|
||||
import socket
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
from urllib.parse import unquote, urlparse
|
||||
|
||||
try:
|
||||
from cloud.db import (
|
||||
POSTGRES_SCHEMA,
|
||||
PgConnection,
|
||||
PostgresDatabase,
|
||||
PostgresDriverMissingError,
|
||||
PostgresIntegrityError,
|
||||
PostgresOperationalError,
|
||||
_convert_placeholders,
|
||||
_translate_sql,
|
||||
init_postgres_schema,
|
||||
)
|
||||
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
from cloud.db import (
|
||||
POSTGRES_SCHEMA,
|
||||
PgConnection,
|
||||
PostgresDatabase,
|
||||
PostgresDriverMissingError,
|
||||
PostgresIntegrityError,
|
||||
PostgresOperationalError,
|
||||
_convert_placeholders,
|
||||
_translate_sql,
|
||||
init_postgres_schema,
|
||||
)
|
||||
|
||||
|
||||
def _has_real_driver():
|
||||
for name in ("psycopg", "psycopg2"):
|
||||
try:
|
||||
__import__(name)
|
||||
return True
|
||||
except ImportError:
|
||||
continue
|
||||
return False
|
||||
|
||||
|
||||
def _real_pg_params():
|
||||
"""返回 (host, port, dbname, user, password) 或 None。
|
||||
|
||||
仅当 CLOUD_TEST_PG_DSN 已设且主机端口 2 秒内可达时返回;
|
||||
其余情况返回 None,调用方 skipTest。绝不打印 DSN。
|
||||
"""
|
||||
dsn = os.environ.get("CLOUD_TEST_PG_DSN")
|
||||
if not dsn:
|
||||
return None
|
||||
parsed = urlparse(dsn)
|
||||
if parsed.scheme not in ("postgres", "postgresql") or not parsed.hostname:
|
||||
return None
|
||||
host = parsed.hostname
|
||||
port = parsed.port or 5432
|
||||
try:
|
||||
with socket.create_connection((host, port), timeout=2):
|
||||
pass
|
||||
except OSError:
|
||||
return None
|
||||
return (
|
||||
host,
|
||||
port,
|
||||
parsed.path.lstrip("/") or "postgres",
|
||||
unquote(parsed.username) if parsed.username else None,
|
||||
unquote(parsed.password) if parsed.password else None,
|
||||
)
|
||||
|
||||
|
||||
class _FakeUnique(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class _FakeDeadlock(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class _FakeCursor:
|
||||
"""记录 SQL、按配置失败或返回预置行;description 为列名元组。"""
|
||||
|
||||
def __init__(self, rows=None, fail_sequence=None):
|
||||
self.executed = []
|
||||
self.rows = list(rows or [])
|
||||
self.fail_sequence = list(fail_sequence or [])
|
||||
self.description = [("id",), ("name",)]
|
||||
|
||||
def execute(self, sql, params=()):
|
||||
self.executed.append((sql, params))
|
||||
if self.fail_sequence:
|
||||
exc = self.fail_sequence.pop(0)
|
||||
if exc is not None:
|
||||
raise exc
|
||||
return self
|
||||
|
||||
def fetchall(self):
|
||||
return list(self.rows)
|
||||
|
||||
def fetchone(self):
|
||||
return self.rows[0] if self.rows else None
|
||||
|
||||
@property
|
||||
def rowcount(self):
|
||||
return len(self.rows) if self.rows else 1
|
||||
|
||||
def close(self):
|
||||
pass
|
||||
|
||||
|
||||
class _FakeConn:
|
||||
def __init__(self, cursor=None):
|
||||
self.cursor_obj = cursor or _FakeCursor()
|
||||
self.committed = False
|
||||
self.rolled_back = False
|
||||
self.closed = False
|
||||
|
||||
def cursor(self):
|
||||
return self.cursor_obj
|
||||
|
||||
def commit(self):
|
||||
self.committed = True
|
||||
|
||||
def rollback(self):
|
||||
self.rolled_back = True
|
||||
|
||||
def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
class _FakeDriver:
|
||||
UniqueViolation = _FakeUnique
|
||||
DeadlockDetected = _FakeDeadlock
|
||||
OperationalError = Exception
|
||||
Error = Exception
|
||||
|
||||
def __init__(self, cursor=None):
|
||||
self.conn = _FakeConn(cursor)
|
||||
|
||||
def connect(self, **kwargs):
|
||||
self.connect_kwargs = kwargs
|
||||
return self.conn
|
||||
|
||||
|
||||
def _make_db(fake):
|
||||
with mock.patch("cloud.db._load_postgres_driver",
|
||||
return_value=("psycopg", fake)):
|
||||
return PostgresDatabase(host="test-host", dbname="test-db",
|
||||
user="test-user",
|
||||
password="test-placeholder")
|
||||
|
||||
|
||||
class DriverMissingTestCase(unittest.TestCase):
|
||||
def test_missing_driver_raises_psycopg_error(self):
|
||||
if _has_real_driver():
|
||||
self.skipTest("环境已装 psycopg 驱动,无驱动报错路径不适用")
|
||||
with self.assertRaises(PostgresDriverMissingError) as ctx:
|
||||
PostgresDatabase(host="test-host", dbname="test-db")
|
||||
self.assertIn("psycopg", str(ctx.exception))
|
||||
|
||||
|
||||
class PlaceholderTranslationTestCase(unittest.TestCase):
|
||||
def test_question_marks_outside_strings_become_percent_s(self):
|
||||
self.assertEqual(
|
||||
_convert_placeholders(
|
||||
"SELECT * FROM t WHERE a = ? AND b = 'x?y' AND c = ?"),
|
||||
"SELECT * FROM t WHERE a = %s AND b = 'x?y' AND c = %s")
|
||||
self.assertEqual(
|
||||
_convert_placeholders("SELECT 'it''s ?' , ?"),
|
||||
"SELECT 'it''s ?' , %s")
|
||||
|
||||
def test_insert_or_replace_becomes_on_conflict_update(self):
|
||||
translated, kind = _translate_sql(
|
||||
"INSERT OR REPLACE INTO project_members"
|
||||
"(project_id, user_id, role) VALUES(?, ?, ?)")
|
||||
self.assertEqual(kind, "normal")
|
||||
self.assertIn("ON CONFLICT (project_id, user_id) DO UPDATE SET",
|
||||
translated)
|
||||
self.assertIn("role = EXCLUDED.role", translated)
|
||||
self.assertIn("%s", translated)
|
||||
self.assertNotIn("?", translated)
|
||||
|
||||
def test_autoincrement_becomes_identity(self):
|
||||
translated, _ = _translate_sql(
|
||||
"CREATE TABLE t(id INTEGER PRIMARY KEY AUTOINCREMENT, x TEXT)")
|
||||
self.assertIn("GENERATED BY DEFAULT AS IDENTITY", translated)
|
||||
self.assertNotIn("AUTOINCREMENT", translated)
|
||||
|
||||
def test_fts5_is_rejected_not_silently_created(self):
|
||||
with self.assertRaises(PostgresOperationalError):
|
||||
_translate_sql(
|
||||
"CREATE VIRTUAL TABLE kb_fts USING fts5(content)")
|
||||
|
||||
|
||||
class AdapterFakeDriverTestCase(unittest.TestCase):
|
||||
def test_row_name_and_index_access(self):
|
||||
fake = _FakeDriver(_FakeCursor(rows=[("t1", "hello")]))
|
||||
db = _make_db(fake)
|
||||
conn = db.connect()
|
||||
self.assertIsInstance(conn, PgConnection)
|
||||
rows = db.read_all(conn, "SELECT * FROM t WHERE id = ?", ("t1",))
|
||||
self.assertEqual(rows[0]["name"], "hello")
|
||||
self.assertEqual(rows[0][0], "t1")
|
||||
# 占位符已转换:fake 游标收到的 SQL 应含 %s 不含 ?。
|
||||
self.assertIn("%s", fake.conn.cursor_obj.executed[0][0])
|
||||
self.assertNotIn("?", fake.conn.cursor_obj.executed[0][0])
|
||||
|
||||
def test_unique_violation_maps_to_integrity_error(self):
|
||||
import sqlite3
|
||||
fake = _FakeDriver(_FakeCursor(fail_sequence=[_FakeUnique("dup")]))
|
||||
db = _make_db(fake)
|
||||
conn = db.connect()
|
||||
with self.assertRaises(PostgresIntegrityError) as ctx:
|
||||
conn.execute("INSERT INTO t VALUES(?)", ("x",))
|
||||
# 不是 sqlite3 异常:调用方按 sqlite 异常捕获时会漏掉,
|
||||
# 因此 server.py 已同步捕获本模块错误(见该文件)。
|
||||
self.assertNotIsInstance(ctx.exception, sqlite3.IntegrityError)
|
||||
|
||||
def test_transaction_rolls_back_on_error(self):
|
||||
fake = _FakeDriver(_FakeCursor(fail_sequence=[None, _FakeUnique("dup")]))
|
||||
db = _make_db(fake)
|
||||
conn = db.connect()
|
||||
with self.assertRaises(PostgresIntegrityError):
|
||||
with db.transaction(conn):
|
||||
conn.execute("INSERT INTO t VALUES(?)", ("a",))
|
||||
conn.execute("INSERT INTO t VALUES(?)", ("b",))
|
||||
self.assertTrue(fake.conn.rolled_back)
|
||||
|
||||
def test_execute_with_retry_retries_only_locked_errors(self):
|
||||
fake = _FakeDriver(_FakeCursor(
|
||||
fail_sequence=[_FakeDeadlock("deadlock"), None]))
|
||||
db = _make_db(fake)
|
||||
conn = db.connect()
|
||||
# DeadlockDetected 在适配层映射为 locked=True,应重试一次后成功。
|
||||
cur = db.execute_with_retry(conn, "UPDATE t SET x = ?", ("y",),
|
||||
backoff_seconds=(0,))
|
||||
self.assertEqual(len(fake.conn.cursor_obj.executed), 2)
|
||||
self.assertIsNotNone(cur)
|
||||
|
||||
def test_repr_never_contains_password(self):
|
||||
fake = _FakeDriver()
|
||||
db = _make_db(fake)
|
||||
self.assertNotIn("test-placeholder", repr(db))
|
||||
self.assertNotIn("test-placeholder", repr(db.connect()))
|
||||
|
||||
|
||||
class SchemaStatementTestCase(unittest.TestCase):
|
||||
EXPECTED_TABLES = (
|
||||
"users", "project_members", "audit_log", "projects", "tasks",
|
||||
"summaries", "kb_docs", "cost_records", "oauth_states",
|
||||
"slack_users", "push_subscriptions",
|
||||
)
|
||||
|
||||
def test_schema_covers_all_business_tables(self):
|
||||
joined = "\n".join(POSTGRES_SCHEMA)
|
||||
for table in self.EXPECTED_TABLES:
|
||||
self.assertIn("CREATE TABLE IF NOT EXISTS %s(" % table, joined)
|
||||
# tasks 含迁移补的 6 列;绝无 FTS5 / AUTOINCREMENT 残留。
|
||||
tasks_stmt = [s for s in POSTGRES_SCHEMA
|
||||
if "CREATE TABLE IF NOT EXISTS tasks(" in s][0]
|
||||
for col in ("result_text", "error_text", "actual_model",
|
||||
"started_at", "finished_at", "child_pid"):
|
||||
self.assertIn(col, tasks_stmt)
|
||||
self.assertNotIn("fts5", joined.lower())
|
||||
self.assertNotIn("AUTOINCREMENT", joined)
|
||||
|
||||
def test_schema_builds_via_fake_driver(self):
|
||||
fake = _FakeDriver()
|
||||
db = _make_db(fake)
|
||||
conn = db.connect()
|
||||
init_postgres_schema(conn)
|
||||
executed = [sql for sql, _ in fake.conn.cursor_obj.executed]
|
||||
self.assertEqual(len(executed), len(POSTGRES_SCHEMA))
|
||||
self.assertTrue(fake.conn.committed)
|
||||
|
||||
|
||||
class RealPostgresSchemaTestCase(unittest.TestCase):
|
||||
"""真实 PG schema:无 DSN 或不可达时 skip。"""
|
||||
|
||||
def test_schema_creates_tables_and_roundtrips(self):
|
||||
params = _real_pg_params()
|
||||
if params is None:
|
||||
self.skipTest("未设置可达的 CLOUD_TEST_PG_DSN,跳过真实 PG 测试")
|
||||
if not _has_real_driver():
|
||||
self.skipTest("未装 psycopg 驱动,跳过真实 PG 测试")
|
||||
host, port, dbname, user, password = params
|
||||
db = PostgresDatabase(host=host, port=port, dbname=dbname,
|
||||
user=user, password=password)
|
||||
conn = db.connect()
|
||||
try:
|
||||
init_postgres_schema(conn)
|
||||
rows = db.read_all(
|
||||
conn,
|
||||
"SELECT table_name FROM information_schema.tables"
|
||||
" WHERE table_schema = 'public'",
|
||||
())
|
||||
names = {r["table_name"] for r in rows}
|
||||
for table in SchemaStatementTestCase.EXPECTED_TABLES:
|
||||
self.assertIn(table, names)
|
||||
# 最小往返:插一条项目再按名读回,随后删掉测试行。
|
||||
with db.transaction(conn):
|
||||
conn.execute(
|
||||
"INSERT INTO projects(id, name, owner_id, created_at)"
|
||||
" VALUES(?, ?, ?, ?)",
|
||||
("pg-test-proj", "pg-test", "pg-test-user", 0.0))
|
||||
rows = db.read_all(
|
||||
conn, "SELECT * FROM projects WHERE id = ?",
|
||||
("pg-test-proj",))
|
||||
self.assertEqual(len(rows), 1)
|
||||
self.assertEqual(rows[0]["name"], "pg-test")
|
||||
finally:
|
||||
try:
|
||||
with db.transaction(conn):
|
||||
conn.execute("DELETE FROM projects WHERE id = ?",
|
||||
("pg-test-proj",))
|
||||
except Exception:
|
||||
pass
|
||||
conn.close()
|
||||
@@ -0,0 +1,172 @@
|
||||
"""Postgres 集成测试:真实服务器上的端到端路径。
|
||||
|
||||
全部测试在无真实 Postgres 时 skip,绝不 fail。判定条件与
|
||||
test_postgres_backend.py 的 _real_pg_params() 一致:须显式
|
||||
设置 CLOUD_TEST_PG_DSN(postgres://user:pass@host:port/dbname)
|
||||
且主机端口 2 秒内可达,且装有 psycopg 驱动。DSN 必须指向
|
||||
一次性测试库;本文件只写业务表与测试行,跑完删掉测试行,
|
||||
不删库、不动其它库。DSN 只从环境变量读,测试不打印它。
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
from cloud.tests.test_postgres_backend import (
|
||||
_has_real_driver,
|
||||
_real_pg_params,
|
||||
)
|
||||
from cloud.db import PostgresDatabase, init_postgres_schema
|
||||
from cloud.migrations import migrate_sqlite_to_postgres
|
||||
from cloud.server import init_db
|
||||
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
from cloud.tests.test_postgres_backend import (
|
||||
_has_real_driver,
|
||||
_real_pg_params,
|
||||
)
|
||||
from cloud.db import PostgresDatabase, init_postgres_schema
|
||||
from cloud.migrations import migrate_sqlite_to_postgres
|
||||
from cloud.server import init_db
|
||||
|
||||
|
||||
def _pg_db_or_skip(testcase):
|
||||
params = _real_pg_params()
|
||||
if params is None:
|
||||
testcase.skipTest("未设置可达的 CLOUD_TEST_PG_DSN,跳过真实 PG 测试")
|
||||
if not _has_real_driver():
|
||||
testcase.skipTest("未装 psycopg 驱动,跳过真实 PG 测试")
|
||||
host, port, dbname, user, password = params
|
||||
return PostgresDatabase(host=host, port=port, dbname=dbname,
|
||||
user=user, password=password)
|
||||
|
||||
|
||||
class PostgresIntegrationTestCase(unittest.TestCase):
|
||||
def test_server_init_db_and_adapter_crud(self):
|
||||
db = _pg_db_or_skip(self)
|
||||
conn = db.connect()
|
||||
try:
|
||||
init_postgres_schema(conn)
|
||||
# 适配层 CRUD:projects 插、读、改、删均走 ? 占位符。
|
||||
with db.transaction(conn):
|
||||
conn.execute(
|
||||
"INSERT INTO projects(id, name, owner_id, created_at)"
|
||||
" VALUES(?, ?, ?, ?)",
|
||||
("pg-int-proj", "pg-int", "pg-int-user", 0.0))
|
||||
rows = db.read_all(
|
||||
conn, "SELECT * FROM projects WHERE id = ?",
|
||||
("pg-int-proj",))
|
||||
self.assertEqual(len(rows), 1)
|
||||
self.assertEqual(rows[0]["name"], "pg-int")
|
||||
with db.transaction(conn):
|
||||
conn.execute(
|
||||
"UPDATE projects SET name = ? WHERE id = ?",
|
||||
("pg-int-2", "pg-int-proj"))
|
||||
rows = db.read_all(
|
||||
conn, "SELECT name FROM projects WHERE id = ?",
|
||||
("pg-int-proj",))
|
||||
self.assertEqual(rows[0]["name"], "pg-int-2")
|
||||
finally:
|
||||
try:
|
||||
with db.transaction(conn):
|
||||
conn.execute("DELETE FROM projects WHERE id = ?",
|
||||
("pg-int-proj",))
|
||||
except Exception:
|
||||
pass
|
||||
conn.close()
|
||||
|
||||
def test_sqlite_to_postgres_migration_is_idempotent(self):
|
||||
import sqlite3
|
||||
db = _pg_db_or_skip(self)
|
||||
# 先建一个最小 SQLite 库:含 projects/tasks 两行。
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
src = os.path.join(tmp, "src.db")
|
||||
sconn = sqlite3.connect(src)
|
||||
sconn.execute(
|
||||
"CREATE TABLE projects(id TEXT PRIMARY KEY, name TEXT,"
|
||||
" owner_id TEXT, created_at REAL)")
|
||||
sconn.execute(
|
||||
"CREATE TABLE tasks(id TEXT PRIMARY KEY, project_id TEXT,"
|
||||
" prompt TEXT, status TEXT)")
|
||||
sconn.execute(
|
||||
"INSERT INTO projects VALUES(?, ?, ?, ?)",
|
||||
("mig-proj", "mig", "mig-user", 1.5))
|
||||
sconn.execute(
|
||||
"INSERT INTO tasks VALUES(?, ?, ?, ?)",
|
||||
("mig-task", "mig-proj", "hello", "queued"))
|
||||
sconn.commit()
|
||||
sconn.close()
|
||||
conn = db.connect()
|
||||
try:
|
||||
first = migrate_sqlite_to_postgres(src, pg_conn=conn)
|
||||
self.assertGreaterEqual(first.get("projects", 0), 1)
|
||||
self.assertGreaterEqual(first.get("tasks", 0), 1)
|
||||
# 第二次跑:ON CONFLICT DO NOTHING,不重复插。
|
||||
second = migrate_sqlite_to_postgres(src, pg_conn=conn)
|
||||
self.assertEqual(second.get("projects", -1), 0)
|
||||
self.assertEqual(second.get("tasks", -1), 0)
|
||||
rows = db.read_all(
|
||||
conn, "SELECT * FROM projects WHERE id = ?",
|
||||
("mig-proj",))
|
||||
self.assertEqual(len(rows), 1)
|
||||
self.assertEqual(rows[0]["name"], "mig")
|
||||
finally:
|
||||
try:
|
||||
with db.transaction(conn):
|
||||
conn.execute("DELETE FROM tasks WHERE id = ?",
|
||||
("mig-task",))
|
||||
conn.execute("DELETE FROM projects WHERE id = ?",
|
||||
("mig-proj",))
|
||||
except Exception:
|
||||
pass
|
||||
conn.close()
|
||||
|
||||
def test_server_init_db_on_postgres(self):
|
||||
# 真走 server.init_db 的 PG 分支:环境变量切到 postgres
|
||||
# 后,init_db 不应建任何 SQLite 文件,而应在 PG 里建表
|
||||
# (含 tasks 的 6 个结果列);本测试不启动 HTTP server。
|
||||
import os
|
||||
from unittest import mock
|
||||
params = _real_pg_params()
|
||||
if params is None:
|
||||
self.skipTest("未设置可达的 CLOUD_TEST_PG_DSN,跳过真实 PG 测试")
|
||||
if not _has_real_driver():
|
||||
self.skipTest("未装 psycopg 驱动,跳过真实 PG 测试")
|
||||
host, port, dbname, user, password = params
|
||||
env = {
|
||||
"CLOUD_DB_BACKEND": "postgres",
|
||||
"CLOUD_DB_HOST": host,
|
||||
"CLOUD_DB_PORT": str(port),
|
||||
"CLOUD_DB_NAME": dbname,
|
||||
"CLOUD_DB_USER": user or "",
|
||||
"CLOUD_DB_PASSWORD": password or "",
|
||||
}
|
||||
with mock.patch.dict(os.environ, env):
|
||||
init_db("ignored-in-postgres-mode.db")
|
||||
db = PostgresDatabase(host=host, port=port, dbname=dbname,
|
||||
user=user, password=password)
|
||||
conn = db.connect()
|
||||
try:
|
||||
rows = db.read_all(
|
||||
conn,
|
||||
"SELECT table_name FROM information_schema.tables"
|
||||
" WHERE table_schema = 'public'",
|
||||
())
|
||||
names = {r["table_name"] for r in rows}
|
||||
self.assertIn("projects", names)
|
||||
self.assertIn("tasks", names)
|
||||
self.assertIn("audit_log", names)
|
||||
cols = db.read_all(
|
||||
conn,
|
||||
"SELECT column_name FROM information_schema.columns"
|
||||
" WHERE table_name = 'tasks'",
|
||||
())
|
||||
colnames = {r["column_name"] for r in cols}
|
||||
for col in ("result_text", "error_text", "actual_model",
|
||||
"started_at", "finished_at", "child_pid"):
|
||||
self.assertIn(col, colnames)
|
||||
finally:
|
||||
conn.close()
|
||||
Reference in new issue
Block a user