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:
host committed 2026-10-26 03:23:06 +08:00
1 parent a19e77ee44
commit dc3f572b2b
10 files changed
+1542 -66

No files matched your search

+670 -22
View File
@@ -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",
]