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
+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",
|
||||
]
|
||||
Reference in new issue
Block a user