host
306cdf0d0c
Worker pool: workers table + register/heartbeat (heartbeat is
registration, no 404), GET /workers with alive (120s default),
POST /workers/heartbeat (X-API-Key, missing worker_id 400).
Tasks gain executor (local/docker/cloud, None = any worker);
worker claim filters by executor in pick_next and in the CAS
WHERE clause, CAS semantics unchanged; CloudExecutor _name
now correctly returns cloud, not local.
CloudExecutor: generic HTTP runner interface (endpoint from
explicit arg or CLOUD_RUNNER_ENDPOINT, auth header name only,
token value only from env into the request header, never into
logs/body/errors); no endpoint configured = explicit failure,
never silent downgrade. Vendor adapters (ECS/Cloud Run) are
out of scope, documented as such.
Auto PR: projects.auto_pr (0/1, create-time only, no PATCH
route exists); POST /tasks/{id}/pr (maintainer+ via
approve_merge) only for merged tasks (else 409); skip reasons
in order: auto_pr_disabled, no_worktree, no_remote, no_gh,
gh_not_authenticated — all return 200 with pr_skipped, task
stays merged, no push attempted; push/pr-create failure
returns 502 with step, task stays merged; success writes
pr_url/pr_number and audits pr_created. No default auto-push
from the worker. Remote URLs in errors are scrubbed of
userinfo; gh credentials are managed by gh itself, never
stored by the control plane.
Tests: host 27 pass, cloud 323 pass (8 skipped are
environment-dependent real-Postgres tests).
1062 lines
41 KiB
Python
1062 lines
41 KiB
Python
"""数据库接口层（云端控制面单机阶段）。
|
||
|
||
本模块分两层，调用方按需选用：
|
||
|
||
1. 模块级函数（历史接口，原样保留，见下文）：``connect`` /
|
||
``is_locked_error`` / ``execute_with_retry`` / ``read_all`` /
|
||
``transaction``。server.py 与旧测试仍走这条路，行为不变。
|
||
2. ``Database`` 接口 + 两个实现：``SQLiteDatabase``（现有 SQLite
|
||
行为的类封装）与 ``PostgresDatabase``（真实 Postgres 实现，
|
||
见下文的适配层）。
|
||
新代码应走接口 + ``create_database`` 工厂，不要再直接拼
|
||
sqlite3 连接。
|
||
|
||
只解决一件事：ThreadingHTTPServer 多线程 + 单个 SQLite 文件时的写冲突。
|
||
不碰任何业务逻辑、不定义新 schema、不改 server.py 的接线：
|
||
本模块只提供连接与读写 helper，协调员后续把 server 的 _connect
|
||
换成这里的 connect 即可（表名/字段与 server.py 的 init_db、
|
||
memory.py 的 summaries 完全一致，因为本模块根本不建表）。
|
||
|
||
使用约定（必须遵守，否则 helper 帮不了你）：
|
||
- 每个线程用自己的连接：绝不在多线程间共用一个 sqlite3 连接
|
||
（sqlite3 默认 check_same_thread=True，共用会直接报错）。
|
||
典型用法是每请求/每线程调用一次 connect()，用完 close()。
|
||
- 写操作走 execute_with_retry + transaction：一次语句失败重试是
|
||
安全的（单条 INSERT/UPDATE 被锁时 SQLite 已回滚该语句）；
|
||
不要把多条非幂等语句包在一个重试里——重试只针对单条语句。
|
||
|
||
单机 SQLite 的上限（诚实边界，见模块末尾注释）：
|
||
- WAL + busy_timeout + 有限重试只能缓解单机多线程写冲突，
|
||
不能变成多机/高并发数据库：写仍是单写者串行，读多写少尚可，
|
||
写密集或跨机器时必须换 Postgres（下一阶段，不在本轮假装解决）。
|
||
|
||
后端选择（环境变量，见 cloud/README.md 数据库一节）：
|
||
- ``CLOUD_DB_BACKEND``：``sqlite``（默认）或 ``postgres``。
|
||
- Postgres 连接参数全部走环境变量，不硬编码地址或密码：
|
||
``CLOUD_DB_HOST``、``CLOUD_DB_PORT``、``CLOUD_DB_NAME``、
|
||
``CLOUD_DB_USER``、``CLOUD_DB_PASSWORD``。
|
||
- Postgres 真后端：``PostgresDatabase`` 已接通真实 Postgres
|
||
（psycopg v3 优先，psycopg2 其次，延迟加载）。初始化时
|
||
找不到驱动仍立即明确报错（``PostgresDriverMissingError``，
|
||
文案含 psycopg），绝不静默回落到 SQLite。调用方仍写
|
||
sqlite3 风格 SQL（``?`` 占位符、``row["col"]``、``conn.execute``
|
||
/ ``fetchall`` / ``commit`` / ``rollback`` / ``close``），
|
||
由本模块的 psycopg 适配层（``PgConnection`` / ``PgCursor`` /
|
||
``PgRow``）翻译成 Postgres 方言；psycopg 异常被映射成
|
||
本模块的 ``PostgresIntegrityError`` /
|
||
``PostgresOperationalError`` / ``PostgresConnectionError``，
|
||
绝不把驱动异常原样漏成 sqlite3 异常。模块级 ``connect()``
|
||
按 ``CLOUD_DB_BACKEND`` 分发：``postgres`` 时返回适配连接
|
||
（``db_path`` 参数被忽略，连接参数只走环境变量），其余
|
||
情况仍是原 SQLite 连接——因此 server / AuthStore /
|
||
MemoryStore / KBStore / CostStore 等沿用 ``connect()`` 的
|
||
调用点无需重构即可走 Postgres；但 ``backup.py`` 的
|
||
SQLite 在线备份 API 仍只适用于 SQLite，Postgres 备份请用
|
||
pg_dump（见 cloud/README.md）。换生产数据库仍不是“改
|
||
一个环境变量就自动完成”：还需要真实 Postgres 部署、驱动
|
||
安装、建库与迁移脚本（见 cloud/migrations.py 与
|
||
cloud/migrate_sqlite_to_postgres.py）。
|
||
"""
|
||
|
||
import os
|
||
import sqlite3
|
||
import time
|
||
from abc import ABC, abstractmethod
|
||
from contextlib import contextmanager
|
||
|
||
# 与 server.py 旧 _connect 的 timeout=30s 保持一致的连接级等待；
|
||
# PRAGMA busy_timeout 是语句级等待，两者不冲突，后者以毫秒为单位
|
||
# 在每个连接上单独设置（不随库文件持久化，见 connect 的注释）。
|
||
DEFAULT_BUSY_TIMEOUT_MS = 5000
|
||
DEFAULT_CONNECT_TIMEOUT_S = 30
|
||
DEFAULT_MAX_RETRIES = 3
|
||
# 指数退避：0.05s、0.1s、0.2s；重试耗尽即抛原始错误，绝不无限循环。
|
||
DEFAULT_BACKOFF_SECONDS = (0.05, 0.1, 0.2)
|
||
_LOCKED_MARKER = "database is locked"
|
||
|
||
# Postgres 连接参数的环境变量名（名字即文档，见 cloud/README.md）。
|
||
POSTGRES_ENV_VARS = {
|
||
"host": "CLOUD_DB_HOST",
|
||
"port": "CLOUD_DB_PORT",
|
||
"dbname": "CLOUD_DB_NAME",
|
||
"user": "CLOUD_DB_USER",
|
||
"password": "CLOUD_DB_PASSWORD",
|
||
}
|
||
BACKEND_ENV_VAR = "CLOUD_DB_BACKEND"
|
||
DEFAULT_BACKEND = "sqlite"
|
||
DEFAULT_SQLITE_PATH = "cloud-control-plane.db"
|
||
|
||
|
||
def connect(db_path, busy_timeout_ms=DEFAULT_BUSY_TIMEOUT_MS,
|
||
timeout=DEFAULT_CONNECT_TIMEOUT_S):
|
||
"""返回一个新连接；后端按 CLOUD_DB_BACKEND 分发。
|
||
|
||
- ``postgres`` / ``postgresql``：返回 Postgres 适配连接
|
||
（``PgConnection``，见下文）；``db_path`` 被忽略，连接参数
|
||
只走 ``CLOUD_DB_*`` 环境变量；驱动缺失时抛
|
||
``PostgresDriverMissingError``，绝不回落到 SQLite。
|
||
- 其它（默认 ``sqlite``）：原 SQLite 行为不变。
|
||
"""
|
||
backend = os.environ.get(BACKEND_ENV_VAR) or DEFAULT_BACKEND
|
||
if str(backend).strip().lower() in ("postgres", "postgresql"):
|
||
return PostgresDatabase().connect()
|
||
return _connect_sqlite(db_path, busy_timeout_ms, timeout)
|
||
|
||
|
||
def _connect_sqlite(db_path, busy_timeout_ms=DEFAULT_BUSY_TIMEOUT_MS,
|
||
timeout=DEFAULT_CONNECT_TIMEOUT_S):
|
||
"""返回一个已设好 WAL + busy_timeout + Row 的新连接。
|
||
|
||
- PRAGMA journal_mode=WAL：读不阻塞写、写不阻塞读，是多线程
|
||
server 的前提；该模式随库文件持久化（:memory: 库除外，
|
||
内存库的 journal_mode 恒为 memory，属 SQLite 本身行为）。
|
||
- PRAGMA busy_timeout：写冲突时语句级等待窗口，单位毫秒；
|
||
它是每个连接的属性，不随库文件持久化，所以每次 connect 都设。
|
||
- row_factory=sqlite3.Row：与 server.py 各 handler 在 _db() 里
|
||
事后设置的 row_factory 一致，行可用 row["id"] 按列名访问。
|
||
- 不设 check_same_thread=False：跨线程共用连接是错误用法，
|
||
保持默认 True 让误用尽早报错，而不是静默出竞态。
|
||
"""
|
||
conn = sqlite3.connect(str(db_path), timeout=timeout)
|
||
try:
|
||
conn.execute("PRAGMA journal_mode=WAL")
|
||
conn.execute("PRAGMA busy_timeout=%d" % int(busy_timeout_ms))
|
||
conn.row_factory = sqlite3.Row
|
||
except Exception:
|
||
conn.close()
|
||
raise
|
||
return conn
|
||
|
||
|
||
def is_locked_error(exc):
|
||
"""只有 OperationalError 且文案含 database is locked 才算写锁冲突。
|
||
|
||
其它 OperationalError（如 no such table、语法错误）重试无用，
|
||
必须立即抛出，不能被重试掩盖成“锁问题”。
|
||
"""
|
||
return (
|
||
isinstance(exc, sqlite3.OperationalError)
|
||
and _LOCKED_MARKER in str(exc).lower()
|
||
)
|
||
|
||
|
||
def execute_with_retry(conn, sql, params=(),
|
||
max_retries=DEFAULT_MAX_RETRIES,
|
||
backoff_seconds=DEFAULT_BACKOFF_SECONDS):
|
||
"""执行单条写语句，遇写锁冲突时有限重试，返回 cursor。
|
||
|
||
重试次数 = max_retries（默认 3 次，即最多尝试 1 + 3 = 4 次），
|
||
每次重试前按 backoff_seconds 依次 sleep（默认 0.05/0.1/0.2s，
|
||
退避表短于重试次数时取最后一项，不越界）。
|
||
重试耗尽后把最后一次的原始 OperationalError 原样抛出，不吞掉、
|
||
不包装成别的异常类型，调用方看到的仍是 sqlite3.OperationalError。
|
||
非锁错误一次都不重试。
|
||
"""
|
||
attempts = 0
|
||
while True:
|
||
try:
|
||
return conn.execute(sql, params)
|
||
except sqlite3.OperationalError as exc:
|
||
if not is_locked_error(exc) or attempts >= max_retries:
|
||
raise
|
||
delay = backoff_seconds[
|
||
min(attempts, len(backoff_seconds) - 1)
|
||
]
|
||
time.sleep(delay)
|
||
attempts += 1
|
||
|
||
|
||
def read_all(conn, sql, params=()):
|
||
"""读操作 helper：不重试，只要求连接本身已走 WAL + busy_timeout。
|
||
|
||
读在 WAL 下不阻塞写，一般不需要重试；如果调用方传的是旧式
|
||
连接（非本模块 connect 建的），读结果仍正确，只是没有 WAL/
|
||
busy_timeout 保障——这属于调用方误用，本函数不静默替它建连接。
|
||
"""
|
||
return conn.execute(sql, params).fetchall()
|
||
|
||
|
||
@contextmanager
|
||
def transaction(conn):
|
||
"""单连接事务上下文：成功 commit，异常 rollback 后原样抛出。
|
||
|
||
与 execute_with_retry 配合的典型写法::
|
||
|
||
conn = connect(db_path)
|
||
try:
|
||
with transaction(conn):
|
||
execute_with_retry(conn, "INSERT INTO tasks(...) ...", (...))
|
||
finally:
|
||
conn.close()
|
||
|
||
注意：重试发生在语句级（transaction 内部），事务本身不重试；
|
||
若事务内某条语句最终抛错，整个事务回滚，不会留下半写状态。
|
||
"""
|
||
try:
|
||
yield conn
|
||
conn.commit()
|
||
except Exception:
|
||
try:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
raise
|
||
|
||
|
||
# ----------------------------------------------------------------------
|
||
# Database 接口层
|
||
# ----------------------------------------------------------------------
|
||
# 方法名按实际代码为准，只覆盖模块级函数已提供的能力，不 invent
|
||
# 不存在的方法名：connect（连接 + WAL 设置）、execute_with_retry
|
||
# （执行）、read_all（查询）、transaction（事务）、is_locked_error
|
||
# （锁冲突判定）。两个后端的方法签名必须一致，见 test_db_backend.py
|
||
# 的接口契约测试（按类检查方法名与签名，不依赖驱动是否安装）。
|
||
|
||
|
||
class Database(ABC):
|
||
"""数据库后端接口。实现者不得静默回落到另一个后端。"""
|
||
|
||
@abstractmethod
|
||
def connect(self):
|
||
"""返回一个新连接（SQLite：已设 WAL + busy_timeout + Row）。"""
|
||
|
||
@abstractmethod
|
||
def execute_with_retry(self, conn, sql, params=(),
|
||
max_retries=DEFAULT_MAX_RETRIES,
|
||
backoff_seconds=DEFAULT_BACKOFF_SECONDS):
|
||
"""执行单条写语句，遇写锁冲突时有限重试，返回 cursor。"""
|
||
|
||
@abstractmethod
|
||
def read_all(self, conn, sql, params=()):
|
||
"""执行查询并返回全部行。"""
|
||
|
||
@abstractmethod
|
||
def transaction(self, conn):
|
||
"""返回单连接事务上下文：成功 commit，异常 rollback 后原样抛出。"""
|
||
|
||
@abstractmethod
|
||
def is_locked_error(self, exc):
|
||
"""判定异常是否为可重试的写锁冲突；非锁错误必须返回 False。"""
|
||
|
||
|
||
class SQLiteDatabase(Database):
|
||
"""Database 的 SQLite 实现：原样保留 WAL/busy_timeout/有界重试。
|
||
|
||
模块级函数仍是同一份实现，本类只是类封装，逐方法委托，
|
||
因此旧调用与新接口看到的行为完全一致。
|
||
"""
|
||
|
||
def __init__(self, db_path, busy_timeout_ms=DEFAULT_BUSY_TIMEOUT_MS,
|
||
timeout=DEFAULT_CONNECT_TIMEOUT_S):
|
||
self.db_path = str(db_path)
|
||
self.busy_timeout_ms = int(busy_timeout_ms)
|
||
self.timeout = timeout
|
||
|
||
def __repr__(self):
|
||
return ("SQLiteDatabase(db_path=%r, busy_timeout_ms=%r, "
|
||
"timeout=%r)" % (self.db_path, self.busy_timeout_ms,
|
||
self.timeout))
|
||
|
||
def connect(self):
|
||
return connect(self.db_path,
|
||
busy_timeout_ms=self.busy_timeout_ms,
|
||
timeout=self.timeout)
|
||
|
||
def execute_with_retry(self, conn, sql, params=(),
|
||
max_retries=DEFAULT_MAX_RETRIES,
|
||
backoff_seconds=DEFAULT_BACKOFF_SECONDS):
|
||
return execute_with_retry(conn, sql, params,
|
||
max_retries=max_retries,
|
||
backoff_seconds=backoff_seconds)
|
||
|
||
def read_all(self, conn, sql, params=()):
|
||
return read_all(conn, sql, params)
|
||
|
||
@contextmanager
|
||
def transaction(self, conn):
|
||
with transaction(conn) as cur:
|
||
yield cur
|
||
|
||
def is_locked_error(self, exc):
|
||
return is_locked_error(exc)
|
||
|
||
|
||
class PostgresDriverMissingError(RuntimeError):
|
||
"""Postgres 驱动缺失：明确报错，绝不静默回落到 SQLite。"""
|
||
|
||
|
||
def _load_postgres_driver():
|
||
"""返回 (驱动名, 模块)；两者都缺时抛 PostgresDriverMissingError。
|
||
|
||
优先 psycopg（v3），其次 psycopg2；只做 import 探测，不连接、
|
||
不读任何真实数据库、不碰任何密码。
|
||
"""
|
||
for name in ("psycopg", "psycopg2"):
|
||
try:
|
||
module = __import__(name)
|
||
except ImportError:
|
||
continue
|
||
return name, module
|
||
raise PostgresDriverMissingError(
|
||
"Postgres 后端需安装 psycopg 驱动（例如 pip install psycopg），"
|
||
"当前环境未找到 psycopg/psycopg2；已拒绝静默回落到 SQLite，"
|
||
"请安装驱动后再选 CLOUD_DB_BACKEND=postgres。"
|
||
)
|
||
|
||
|
||
# ----------------------------------------------------------------------
|
||
# Postgres 适配层：把 sqlite3 风格调用翻译成 psycopg 调用
|
||
# ----------------------------------------------------------------------
|
||
# 不重构所有 store 的前提：现有代码到处用 ``?`` 占位符、
|
||
# ``row["col"]``、``conn.execute().fetchall()``、``commit`` /
|
||
# ``rollback`` / ``close``、``cursor.lastrowid``、``cursor.rowcount``，
|
||
# 以及少量 SQLite 方言（``INSERT OR REPLACE``、``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, "
|
||
"repo_url TEXT, "
|
||
"auto_pr INTEGER)",
|
||
# 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, "
|
||
"executor TEXT, "
|
||
"pr_url TEXT, "
|
||
"pr_number INTEGER, "
|
||
"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 连接混进来。
|
||
|
||
建表后还给老库 projects 表补 repo_url 列（新库建表时已含）：
|
||
PRAGMA table_info 经适配层翻译成 information_schema 查询，
|
||
调用方取 row[1] 即列名；缺列时 ALTER TABLE 补列，已迁移的
|
||
库是 no-op。server.init_db 的 Postgres 分支还会再跑一次
|
||
migrate_projects_table，两次补列互不冲突（第二次是 no-op）。
|
||
"""
|
||
for statement in POSTGRES_SCHEMA:
|
||
conn.execute(statement)
|
||
cols = {row[1] for row in conn.execute(
|
||
"PRAGMA table_info(projects)").fetchall()}
|
||
if "repo_url" not in cols:
|
||
conn.execute("ALTER TABLE projects ADD COLUMN repo_url TEXT")
|
||
if "auto_pr" not in cols:
|
||
conn.execute("ALTER TABLE projects ADD COLUMN auto_pr INTEGER")
|
||
tcols = {row[1] for row in conn.execute(
|
||
"PRAGMA table_info(tasks)").fetchall()}
|
||
for _name, _type in (("executor", "TEXT"), ("pr_url", "TEXT"),
|
||
("pr_number", "INTEGER")):
|
||
if _name not in tcols:
|
||
conn.execute(
|
||
"ALTER TABLE tasks ADD COLUMN %s %s" % (_name, _type))
|
||
conn.commit()
|
||
|
||
|
||
class PostgresDatabase(Database):
|
||
"""Database 的 Postgres 实现（已接通真实 Postgres）。
|
||
|
||
连接参数全部走环境变量（见 POSTGRES_ENV_VARS），显式参数优先
|
||
于环境变量；未给出的参数保持 None（传给驱动时省略，不硬编码
|
||
任何地址或密码）。初始化时即探测驱动：缺驱动立即抛
|
||
PostgresDriverMissingError。驱动分 psycopg（v3，优先）与
|
||
psycopg2，连接参数名两者通用（host/port/dbname/user/password）。
|
||
本类的 __repr__ 绝不输出 password。
|
||
"""
|
||
|
||
def __init__(self, host=None, port=None, dbname=None, user=None,
|
||
password=None):
|
||
env = POSTGRES_ENV_VARS
|
||
self.host = host if host is not None else os.environ.get(env["host"])
|
||
raw_port = port if port is not None else os.environ.get(env["port"])
|
||
if raw_port is None or raw_port == "":
|
||
self.port = None
|
||
else:
|
||
try:
|
||
self.port = int(raw_port)
|
||
except (TypeError, ValueError):
|
||
raise ValueError(
|
||
"环境变量 %s 须为整数端口，当前值无法解析" % env["port"])
|
||
self.dbname = (dbname if dbname is not None
|
||
else os.environ.get(env["dbname"]))
|
||
self.user = (user if user is not None
|
||
else os.environ.get(env["user"]))
|
||
# 只保存、不输出、不记录日志；调用方负责传入方式的安全。
|
||
self.password = (password if password is not None
|
||
else os.environ.get(env["password"]))
|
||
# 驱动探测在初始化时完成：缺驱动这里就报错，不等到首次查询。
|
||
self.driver_name, self._driver = _load_postgres_driver()
|
||
|
||
def __repr__(self):
|
||
return ("PostgresDatabase(host=%r, port=%r, dbname=%r, user=%r)"
|
||
% (self.host, self.port, self.dbname, self.user))
|
||
|
||
def _connect_kwargs(self):
|
||
kwargs = {}
|
||
if self.host is not None:
|
||
kwargs["host"] = self.host
|
||
if self.port is not None:
|
||
kwargs["port"] = self.port
|
||
if self.dbname is not None:
|
||
kwargs["dbname"] = self.dbname
|
||
if self.user is not None:
|
||
kwargs["user"] = self.user
|
||
if self.password is not None:
|
||
kwargs["password"] = self.password
|
||
return kwargs
|
||
|
||
def connect(self):
|
||
try:
|
||
real = self._driver.connect(**self._connect_kwargs())
|
||
except Exception as exc:
|
||
mapped = _map_driver_error(exc, self._driver)
|
||
if isinstance(mapped, PostgresDatabaseError):
|
||
raise mapped from exc
|
||
raise PostgresConnectionError(str(exc)) from exc
|
||
return PgConnection(real, self._driver)
|
||
|
||
def execute_with_retry(self, conn, sql, params=(),
|
||
max_retries=DEFAULT_MAX_RETRIES,
|
||
backoff_seconds=DEFAULT_BACKOFF_SECONDS):
|
||
attempts = 0
|
||
while True:
|
||
try:
|
||
return conn.execute(sql, params)
|
||
except PostgresOperationalError as exc:
|
||
if not exc.locked or attempts >= max_retries:
|
||
raise
|
||
# Postgres 里失败语句会让事务进入 aborted 状态，
|
||
# 不 rollback 就无法重试下一条语句。
|
||
try:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
delay = backoff_seconds[
|
||
min(attempts, len(backoff_seconds) - 1)]
|
||
time.sleep(delay)
|
||
attempts += 1
|
||
|
||
def read_all(self, conn, sql, params=()):
|
||
return conn.execute(sql, params).fetchall()
|
||
|
||
@contextmanager
|
||
def transaction(self, conn):
|
||
try:
|
||
yield conn
|
||
conn.commit()
|
||
except Exception:
|
||
try:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
raise
|
||
|
||
def is_locked_error(self, exc):
|
||
if isinstance(exc, PostgresOperationalError):
|
||
return bool(exc.locked)
|
||
# 驱动原始异常直接传入时也按名/SQLSTATE 判定（不依赖实例化）。
|
||
name = type(exc).__name__
|
||
code = getattr(exc, "sqlstate", None) or getattr(exc, "pgcode", None)
|
||
return name in ("DeadlockDetected", "LockNotAvailable") \
|
||
or code in ("40P01", "55P03")
|
||
|
||
|
||
def create_database(backend=None, db_path=None, **kwargs):
|
||
"""后端工厂：按 backend（或环境变量）返回 Database 实现。
|
||
|
||
- backend 为 None 时读 ``CLOUD_DB_BACKEND``，未设则默认 sqlite。
|
||
- ``sqlite``：返回 SQLiteDatabase（db_path 未给则用默认文件名，
|
||
与 server 旧默认值一致）；kwargs 可传 busy_timeout_ms/timeout。
|
||
- ``postgres``/``postgresql``：返回 PostgresDatabase；驱动缺失
|
||
时初始化即抛 PostgresDriverMissingError（文案含 psycopg），
|
||
绝不回落到 SQLite；kwargs 可传 host/port/dbname/user/password
|
||
覆盖环境变量（测试只能用 fake/test 占位）。
|
||
- 其它值抛 ValueError，不静默当成 sqlite。
|
||
"""
|
||
resolved = backend
|
||
if resolved is None:
|
||
resolved = os.environ.get(BACKEND_ENV_VAR) or DEFAULT_BACKEND
|
||
resolved = str(resolved).strip().lower()
|
||
if resolved == "sqlite":
|
||
path = db_path if db_path is not None else DEFAULT_SQLITE_PATH
|
||
return SQLiteDatabase(path, **kwargs)
|
||
if resolved in ("postgres", "postgresql"):
|
||
return PostgresDatabase(**kwargs)
|
||
raise ValueError(
|
||
"未知数据库后端 %r（%s），仅支持 sqlite/postgres"
|
||
% (resolved, BACKEND_ENV_VAR))
|
||
|
||
|
||
def database_from_env(db_path=None, **kwargs):
|
||
"""create_database 的环境变量版：backend 一律读 CLOUD_DB_BACKEND。"""
|
||
return create_database(backend=None, db_path=db_path, **kwargs)
|
||
|
||
|
||
__all__ = [
|
||
"DEFAULT_BUSY_TIMEOUT_MS",
|
||
"DEFAULT_CONNECT_TIMEOUT_S",
|
||
"DEFAULT_MAX_RETRIES",
|
||
"DEFAULT_BACKOFF_SECONDS",
|
||
"POSTGRES_ENV_VARS",
|
||
"BACKEND_ENV_VAR",
|
||
"DEFAULT_BACKEND",
|
||
"connect",
|
||
"is_locked_error",
|
||
"execute_with_retry",
|
||
"read_all",
|
||
"transaction",
|
||
"Database",
|
||
"SQLiteDatabase",
|
||
"PostgresDatabase",
|
||
"PostgresDriverMissingError",
|
||
"PostgresDatabaseError",
|
||
"PostgresIntegrityError",
|
||
"PostgresOperationalError",
|
||
"PostgresConnectionError",
|
||
"PgConnection",
|
||
"PgCursor",
|
||
"PgRow",
|
||
"POSTGRES_SCHEMA",
|
||
"init_postgres_schema",
|
||
"create_database",
|
||
"database_from_env",
|
||
]
|