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.
422 lines
17 KiB
Python
422 lines
17 KiB
Python
"""cloud/db.py 接口层单测:Database + 双后端 + 工厂 + server 启动校验。
|
||
|
||
不删、不改 cloud/tests/test_db.py 的原有测试;本文件只覆盖新增的
|
||
接口层。测试里出现的 Postgres 密码一律是 fake/test 占位
|
||
("test-placeholder"),不是真实格式的密钥,绝不写真实凭证。
|
||
|
||
覆盖点(与任务一一对应):
|
||
- SQLite 后端原有行为不变:WAL / busy_timeout / Row / 有界重试 /
|
||
事务提交回滚,全部经 SQLiteDatabase 实例走一遍。
|
||
- Postgres 后端无驱动时明确报错:实例化即抛 RuntimeError 且文案含
|
||
psycopg,且不是静默回落(返回值绝不是 SQLiteDatabase)。
|
||
- 后端切换不丢接口契约:Database 的抽象方法集合 == 两个后端类
|
||
都提供的方法集合,且签名(除 self 外)一致;Postgres 方法体在
|
||
骨架阶段抛 NotImplementedError/驱动缺失错误。
|
||
- server 启动校验:CLOUD_DB_BACKEND=postgres 时 create_server 在
|
||
建表前就报错;默认/sqlite 时仍能正常建 server。
|
||
"""
|
||
|
||
import inspect
|
||
import os
|
||
import sqlite3
|
||
import sys
|
||
import tempfile
|
||
import unittest
|
||
from pathlib import Path
|
||
from unittest import mock
|
||
|
||
try:
|
||
from cloud.db import (
|
||
Database,
|
||
PostgresDatabase,
|
||
PostgresDriverMissingError,
|
||
SQLiteDatabase,
|
||
create_database,
|
||
database_from_env,
|
||
)
|
||
from cloud.server import create_server
|
||
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||
from cloud.db import (
|
||
Database,
|
||
PostgresDatabase,
|
||
PostgresDriverMissingError,
|
||
SQLiteDatabase,
|
||
create_database,
|
||
database_from_env,
|
||
)
|
||
from cloud.server import create_server
|
||
|
||
INTERFACE_METHODS = (
|
||
"connect",
|
||
"execute_with_retry",
|
||
"read_all",
|
||
"transaction",
|
||
"is_locked_error",
|
||
)
|
||
|
||
TASKS_DDL = (
|
||
"CREATE TABLE IF NOT EXISTS tasks("
|
||
"id TEXT PRIMARY KEY, project_id TEXT, prompt TEXT, status TEXT)"
|
||
)
|
||
|
||
|
||
class _EnvCleanMixin:
|
||
def setUp(self):
|
||
self._env_backup = {
|
||
k: os.environ.get(k) for k in
|
||
("CLOUD_DB_BACKEND", "CLOUD_DB_HOST", "CLOUD_DB_PORT",
|
||
"CLOUD_DB_NAME", "CLOUD_DB_USER", "CLOUD_DB_PASSWORD")
|
||
}
|
||
for key in self._env_backup:
|
||
os.environ.pop(key, None)
|
||
|
||
def tearDown(self):
|
||
for key, value in self._env_backup.items():
|
||
if value is None:
|
||
os.environ.pop(key, None)
|
||
else:
|
||
os.environ[key] = value
|
||
|
||
|
||
class SQLiteBackendPreservedTestCase(_EnvCleanMixin, unittest.TestCase):
|
||
def test_sqlite_database_keeps_wal_busy_timeout_row_retry(self):
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
path = os.path.join(tmp, "backend.db")
|
||
db = SQLiteDatabase(path)
|
||
conn = db.connect()
|
||
try:
|
||
self.assertEqual(conn.row_factory, sqlite3.Row)
|
||
self.assertEqual(
|
||
conn.execute("PRAGMA journal_mode").fetchone()[0], "wal")
|
||
self.assertEqual(
|
||
conn.execute("PRAGMA busy_timeout").fetchone()[0], 5000)
|
||
conn.execute(TASKS_DDL)
|
||
conn.commit()
|
||
with db.transaction(conn):
|
||
db.execute_with_retry(
|
||
conn,
|
||
"INSERT INTO tasks(id, project_id, prompt, status)"
|
||
" VALUES(?, ?, ?, 'queued')",
|
||
("t1", "p1", "hello"))
|
||
rows = db.read_all(
|
||
conn, "SELECT * FROM tasks WHERE id = ?", ("t1",))
|
||
self.assertEqual(len(rows), 1)
|
||
self.assertEqual(rows[0]["prompt"], "hello")
|
||
# 事务回滚:重复主键失败后仍只有 1 条。
|
||
with self.assertRaises(sqlite3.IntegrityError):
|
||
with db.transaction(conn):
|
||
db.execute_with_retry(
|
||
conn,
|
||
"INSERT INTO tasks(id, project_id, prompt,"
|
||
" status) VALUES(?, ?, ?, 'queued')",
|
||
("t1", "p1", "dup"))
|
||
self.assertEqual(
|
||
len(db.read_all(conn, "SELECT * FROM tasks")), 1)
|
||
finally:
|
||
conn.close()
|
||
# 锁判定仍走原规则:锁错误 True,其它 False。
|
||
self.assertTrue(db.is_locked_error(
|
||
sqlite3.OperationalError("database is locked")))
|
||
self.assertFalse(db.is_locked_error(
|
||
sqlite3.OperationalError("no such table: tasks")))
|
||
|
||
def test_sqlite_database_custom_busy_timeout(self):
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
path = os.path.join(tmp, "backend.db")
|
||
db = SQLiteDatabase(path, busy_timeout_ms=1234)
|
||
conn = db.connect()
|
||
try:
|
||
self.assertEqual(
|
||
conn.execute("PRAGMA busy_timeout").fetchone()[0], 1234)
|
||
finally:
|
||
conn.close()
|
||
|
||
def test_sqlite_retry_still_bounded_via_instance(self):
|
||
# 用 fake 连接走实例方法:重试耗尽抛原始错误,不无限循环。
|
||
class FakeConn:
|
||
def __init__(self):
|
||
self.calls = 0
|
||
|
||
def execute(self, sql, params=()):
|
||
self.calls += 1
|
||
raise sqlite3.OperationalError("database is locked")
|
||
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
db = SQLiteDatabase(os.path.join(tmp, "x.db"))
|
||
fake = FakeConn()
|
||
with mock.patch("cloud.db.time.sleep"):
|
||
with self.assertRaises(sqlite3.OperationalError):
|
||
db.execute_with_retry(fake, "INSERT ...", ())
|
||
self.assertEqual(fake.calls, 4) # 1 次初始 + 3 次重试
|
||
|
||
|
||
class PostgresSkeletonTestCase(_EnvCleanMixin, unittest.TestCase):
|
||
def test_postgres_init_without_driver_raises_psycopg_error(self):
|
||
# 本环境无 psycopg/psycopg2(已实测);若将来装了驱动,
|
||
# 本用例应改为跳过——当前断言的是“无驱动时明确报错”。
|
||
try:
|
||
import psycopg # noqa: F401
|
||
has_driver = True
|
||
except ImportError:
|
||
try:
|
||
import psycopg2 # noqa: F401
|
||
has_driver = True
|
||
except ImportError:
|
||
has_driver = False
|
||
if has_driver:
|
||
self.skipTest("环境已装 psycopg 驱动,无驱动报错路径不适用")
|
||
with self.assertRaises(PostgresDriverMissingError) as ctx:
|
||
PostgresDatabase(host="test-host", dbname="test-db",
|
||
user="test-user",
|
||
password="test-placeholder")
|
||
self.assertIn("psycopg", str(ctx.exception))
|
||
# 绝不静默回落:异常不是 SQLite 实例,工厂同样抛错而非返回。
|
||
with self.assertRaises(PostgresDriverMissingError) as ctx2:
|
||
create_database(backend="postgres")
|
||
self.assertIn("psycopg", str(ctx2.exception))
|
||
self.assertIsInstance(ctx2.exception, RuntimeError)
|
||
|
||
def test_postgres_params_come_from_env_not_hardcoded(self):
|
||
# 驱动探测 mock 成已装(不连真实 Postgres),验证参数来源:
|
||
# 环境变量进属性、显式参数优先于环境变量、repr 不泄密码。
|
||
from cloud.db import POSTGRES_ENV_VARS
|
||
self.assertEqual(POSTGRES_ENV_VARS["host"], "CLOUD_DB_HOST")
|
||
self.assertEqual(POSTGRES_ENV_VARS["port"], "CLOUD_DB_PORT")
|
||
self.assertEqual(POSTGRES_ENV_VARS["dbname"], "CLOUD_DB_NAME")
|
||
self.assertEqual(POSTGRES_ENV_VARS["user"], "CLOUD_DB_USER")
|
||
self.assertEqual(POSTGRES_ENV_VARS["password"], "CLOUD_DB_PASSWORD")
|
||
env = {
|
||
"CLOUD_DB_HOST": "env-host",
|
||
"CLOUD_DB_PORT": "5433",
|
||
"CLOUD_DB_NAME": "env-db",
|
||
"CLOUD_DB_USER": "env-user",
|
||
"CLOUD_DB_PASSWORD": "test-placeholder",
|
||
}
|
||
with mock.patch.dict(os.environ, env):
|
||
with mock.patch("cloud.db._load_postgres_driver",
|
||
return_value=("psycopg", object())):
|
||
db = PostgresDatabase()
|
||
self.assertEqual(db.host, "env-host")
|
||
self.assertEqual(db.port, 5433)
|
||
self.assertEqual(db.dbname, "env-db")
|
||
self.assertEqual(db.user, "env-user")
|
||
self.assertNotIn("test-placeholder", repr(db))
|
||
# 显式参数优先于环境变量;未给出的参数不硬编码成默认值。
|
||
db2 = PostgresDatabase(host="explicit-host")
|
||
self.assertEqual(db2.host, "explicit-host")
|
||
self.assertEqual(db2.dbname, "env-db")
|
||
|
||
def test_postgres_methods_are_implemented_via_fake_driver(self):
|
||
# 骨架已由真实实现取代:用 fake 驱动(不连真实 Postgres)
|
||
# 验证 connect / read_all / transaction / is_locked_error
|
||
# 均可用,且驱动唯一冲突被映射成模块内完整性错误。
|
||
from cloud.db import (
|
||
PgConnection,
|
||
PostgresIntegrityError,
|
||
PostgresOperationalError,
|
||
)
|
||
|
||
class _FakeUnique(Exception):
|
||
pass
|
||
|
||
class _FakeCursor:
|
||
def __init__(self):
|
||
self.executed = []
|
||
self.description = [("id",), ("name",)]
|
||
|
||
def execute(self, sql, params=()):
|
||
self.executed.append((sql, params))
|
||
if "dup" in sql:
|
||
raise _FakeUnique("duplicate key")
|
||
return self
|
||
|
||
def fetchall(self):
|
||
return [("t1", "hello")]
|
||
|
||
def fetchone(self):
|
||
return ("t1", "hello")
|
||
|
||
@property
|
||
def rowcount(self):
|
||
return 1
|
||
|
||
def close(self):
|
||
pass
|
||
|
||
class _FakeConn:
|
||
def __init__(self):
|
||
self.cursor_obj = _FakeCursor()
|
||
self.committed = False
|
||
self.rolled_back = False
|
||
|
||
def cursor(self):
|
||
return self.cursor_obj
|
||
|
||
def commit(self):
|
||
self.committed = True
|
||
|
||
def rollback(self):
|
||
self.rolled_back = True
|
||
|
||
def close(self):
|
||
pass
|
||
|
||
class _FakeDriver:
|
||
UniqueViolation = _FakeUnique
|
||
OperationalError = Exception
|
||
Error = Exception
|
||
|
||
def __init__(self):
|
||
self.conn = _FakeConn()
|
||
|
||
def connect(self, **kwargs):
|
||
return self.conn
|
||
|
||
fake = _FakeDriver()
|
||
with mock.patch("cloud.db._load_postgres_driver",
|
||
return_value=("psycopg", fake)):
|
||
db = PostgresDatabase(host="test-host", dbname="test-db")
|
||
conn = db.connect()
|
||
self.assertIsInstance(conn, PgConnection)
|
||
rows = db.read_all(conn, "SELECT * FROM t WHERE id = ?",
|
||
("t1",))
|
||
self.assertEqual(rows[0]["name"], "hello")
|
||
self.assertEqual(rows[0][0], "t1")
|
||
# 唯一冲突 -> 模块内错误,不是驱动原始异常、不是 sqlite 异常。
|
||
with self.assertRaises(PostgresIntegrityError):
|
||
db.execute_with_retry(conn, "SELECT dup", ())
|
||
# 事务回滚:第二条失败后 rollback 被调用且异常原样抛出。
|
||
with self.assertRaises(PostgresIntegrityError):
|
||
with db.transaction(conn):
|
||
conn.execute("SELECT dup", ())
|
||
self.assertTrue(fake.conn.rolled_back)
|
||
# 锁判定:本模块锁错误 True,其它错误 False。
|
||
self.assertTrue(db.is_locked_error(
|
||
PostgresOperationalError("deadlock", locked=True)))
|
||
self.assertFalse(db.is_locked_error(ValueError("x")))
|
||
|
||
def test_postgres_repr_never_contains_password(self):
|
||
skeleton = object.__new__(PostgresDatabase)
|
||
skeleton.host = "test-host"
|
||
skeleton.port = None
|
||
skeleton.dbname = "test-db"
|
||
skeleton.user = "test-user"
|
||
skeleton.password = "test-placeholder"
|
||
self.assertNotIn("test-placeholder", repr(skeleton))
|
||
|
||
|
||
class InterfaceContractTestCase(_EnvCleanMixin, unittest.TestCase):
|
||
def test_database_is_abstract(self):
|
||
with self.assertRaises(TypeError):
|
||
Database()
|
||
|
||
def test_both_backends_expose_same_interface_methods(self):
|
||
abstract = set(Database.__abstractmethods__)
|
||
self.assertEqual(set(INTERFACE_METHODS), abstract)
|
||
for cls in (SQLiteDatabase, PostgresDatabase):
|
||
for name in INTERFACE_METHODS:
|
||
self.assertTrue(
|
||
callable(getattr(cls, name, None)),
|
||
"%s 缺少接口方法 %s" % (cls.__name__, name))
|
||
|
||
def test_both_backends_signatures_match_interface(self):
|
||
for cls in (SQLiteDatabase, PostgresDatabase):
|
||
for name in INTERFACE_METHODS:
|
||
expected = inspect.signature(getattr(Database, name))
|
||
actual = inspect.signature(getattr(cls, name))
|
||
expected_params = [p for p in expected.parameters
|
||
if p != "self"]
|
||
actual_params = [p for p in actual.parameters
|
||
if p != "self"]
|
||
self.assertEqual(
|
||
expected_params, actual_params,
|
||
"%s.%s 签名与 Database 不一致" % (cls.__name__, name))
|
||
|
||
def test_factory_backend_switching(self):
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
path = os.path.join(tmp, "factory.db")
|
||
# 默认(环境变量未设)仍是 SQLite。
|
||
db = create_database(db_path=path)
|
||
self.assertIsInstance(db, SQLiteDatabase)
|
||
self.assertIsInstance(db, Database)
|
||
# 显式 sqlite。
|
||
self.assertIsInstance(
|
||
create_database(backend="sqlite", db_path=path),
|
||
SQLiteDatabase)
|
||
# 环境变量 sqlite。
|
||
with mock.patch.dict(os.environ,
|
||
{"CLOUD_DB_BACKEND": "sqlite"}):
|
||
self.assertIsInstance(
|
||
database_from_env(db_path=path), SQLiteDatabase)
|
||
# 未知后端:明确 ValueError,不静默当 sqlite。
|
||
with self.assertRaises(ValueError):
|
||
create_database(backend="mysql", db_path=path)
|
||
with mock.patch.dict(os.environ,
|
||
{"CLOUD_DB_BACKEND": "mysql"}):
|
||
with self.assertRaises(ValueError):
|
||
database_from_env(db_path=path)
|
||
# 环境变量 postgres:无驱动时抛含 psycopg 的错误,
|
||
# 绝不返回一个 SQLite 实例假装切换成功。
|
||
with mock.patch.dict(os.environ,
|
||
{"CLOUD_DB_BACKEND": "postgres"}):
|
||
try:
|
||
import psycopg # noqa: F401
|
||
has_driver = True
|
||
except ImportError:
|
||
try:
|
||
import psycopg2 # noqa: F401
|
||
has_driver = True
|
||
except ImportError:
|
||
has_driver = False
|
||
if has_driver:
|
||
self.skipTest("环境已装 psycopg 驱动")
|
||
with self.assertRaises(RuntimeError) as ctx:
|
||
database_from_env(db_path=path)
|
||
self.assertIn("psycopg", str(ctx.exception))
|
||
|
||
|
||
class ServerBackendSelectionTestCase(_EnvCleanMixin, unittest.TestCase):
|
||
def test_server_postgres_backend_fails_at_startup(self):
|
||
try:
|
||
import psycopg # noqa: F401
|
||
has_driver = True
|
||
except ImportError:
|
||
try:
|
||
import psycopg2 # noqa: F401
|
||
has_driver = True
|
||
except ImportError:
|
||
has_driver = False
|
||
if has_driver:
|
||
self.skipTest("环境已装 psycopg 驱动")
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
path = os.path.join(tmp, "should-not-be-created.db")
|
||
with mock.patch.dict(os.environ,
|
||
{"CLOUD_DB_BACKEND": "postgres"}):
|
||
with self.assertRaises(RuntimeError) as ctx:
|
||
create_server(path, "127.0.0.1", 0)
|
||
self.assertIn("psycopg", str(ctx.exception))
|
||
# 启动前失败:不应已建出 SQLite 库文件假装可用。
|
||
self.assertFalse(os.path.exists(path))
|
||
|
||
def test_server_default_and_sqlite_still_work(self):
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
path = os.path.join(tmp, "ok.db")
|
||
server = create_server(path, "127.0.0.1", 0)
|
||
try:
|
||
self.assertTrue(os.path.exists(path))
|
||
finally:
|
||
server.server_close()
|
||
with mock.patch.dict(os.environ,
|
||
{"CLOUD_DB_BACKEND": "sqlite"}):
|
||
path2 = os.path.join(tmp, "ok2.db")
|
||
server2 = create_server(path2, "127.0.0.1", 0)
|
||
try:
|
||
self.assertTrue(os.path.exists(path2))
|
||
finally:
|
||
server2.server_close()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|