390 lines
16 KiB
Python
390 lines
16 KiB
Python
"""数据库接口层（云端控制面单机阶段）。
|
||
|
||
本模块分两层，调用方按需选用：
|
||
|
||
1. 模块级函数（历史接口，原样保留，见下文）：``connect`` /
|
||
``is_locked_error`` / ``execute_with_retry`` / ``read_all`` /
|
||
``transaction``。server.py 与旧测试仍走这条路，行为不变。
|
||
2. ``Database`` 接口 + 两个实现：``SQLiteDatabase``（现有 SQLite
|
||
行为的类封装）与 ``PostgresDatabase``（骨架，见下文）。
|
||
新代码应走接口 + ``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 骨架的诚实边界：标准库没有 Postgres 驱动，
|
||
``PostgresDatabase`` 初始化时找不到 psycopg/psycopg2 会立即
|
||
明确报错（``PostgresDriverMissingError``，文案含 psycopg），
|
||
绝不静默回落到 SQLite；即使驱动已安装，其连接/执行/查询/
|
||
事务方法目前仍是 ``NotImplementedError`` 骨架，尚未接通真实
|
||
Postgres——换生产数据库还需要真实部署、驱动安装与迁移脚本，
|
||
不是改一个环境变量就自动完成。
|
||
"""
|
||
|
||
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):
|
||
"""返回一个已设好 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_NOT_IMPLEMENTED = (
|
||
"PostgresDatabase 是骨架，尚未实现：需要真实 Postgres 部署、"
|
||
"psycopg 驱动安装与迁移脚本后才能接通，当前不假装已接通 Postgres。"
|
||
)
|
||
|
||
|
||
class PostgresDatabase(Database):
|
||
"""Database 的 Postgres 骨架（未接通真实 Postgres）。
|
||
|
||
连接参数全部走环境变量（见 POSTGRES_ENV_VARS），显式参数优先
|
||
于环境变量；未给出的参数保持 None，不硬编码任何地址或密码。
|
||
初始化时即探测驱动：缺驱动立即抛 PostgresDriverMissingError，
|
||
因此 server 选 postgres 后端时在启动阶段就失败，不会在运行时
|
||
中途才炸。即使驱动已安装，connect/execute/read/transaction
|
||
目前仍抛 NotImplementedError——骨架不是可用后端。
|
||
本类的 __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(self):
|
||
raise NotImplementedError(_POSTGRES_NOT_IMPLEMENTED)
|
||
|
||
def execute_with_retry(self, conn, sql, params=(),
|
||
max_retries=DEFAULT_MAX_RETRIES,
|
||
backoff_seconds=DEFAULT_BACKOFF_SECONDS):
|
||
raise NotImplementedError(_POSTGRES_NOT_IMPLEMENTED)
|
||
|
||
def read_all(self, conn, sql, params=()):
|
||
raise NotImplementedError(_POSTGRES_NOT_IMPLEMENTED)
|
||
|
||
def transaction(self, conn):
|
||
raise NotImplementedError(_POSTGRES_NOT_IMPLEMENTED)
|
||
|
||
def is_locked_error(self, exc):
|
||
raise NotImplementedError(_POSTGRES_NOT_IMPLEMENTED)
|
||
|
||
|
||
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",
|
||
"create_database",
|
||
"database_from_env",
|
||
]
|