diff --git a/.trae/docs/iter-02-sqlite-backend.md b/.trae/docs/iter-02-sqlite-backend.md new file mode 100644 index 0000000..1588976 --- /dev/null +++ b/.trae/docs/iter-02-sqlite-backend.md @@ -0,0 +1,67 @@ +# 迭代 02：SQLite 状态后端 (SQLiteBackend) + +## 本轮目标 + +在 `src/pyflowx/storage.py` 新增 `SQLiteBackend`，扩展 `StateBackend` ABC， +解决 `JSONBackend` 全量加载/落盘在大规模任务（10k+）下的性能瓶颈。 + +使用标准库 `sqlite3`，零新依赖。作为 `JSONBackend` 的 drop-in 替换， +支持 TTL、断点续跑、`batch()` 批量提交。 + +## 改动文件清单 + +- `src/pyflowx/storage.py` — 新增 `SQLiteBackend` 类 +- `src/pyflowx/__init__.py` — 导出 `SQLiteBackend` +- `tests/test_storage.py` — 新增 SQLiteBackend 测试用例 +- `docs/api.rst` — 补充 SQLiteBackend 文档（如需） + +## 关键决策与依据 + +### 1. 复用 `_TTLStateBackendMixin` + +`SQLiteBackend` 继承 `_TTLStateBackendMixin`，只需实现 4 个原语 +（`_get_raw`/`_put_raw`/`_iter_raw`/`_clear_raw`），TTL 逻辑由 mixin 提供。 +与 `MemoryBackend`/`JSONBackend` 保持一致架构。 + +### 2. 值序列化策略 + +值以 JSON 文本存入 `value` 列（`json.dumps`/`json.loads`），与 `JSONBackend` +保持相同契约（结果必须可 JSON 序列化）。这样： +- 三种后端行为一致，用户可无缝切换 +- 复用现有 `StorageError` 序列化错误处理 +- 无需引入 pickle（安全风险） + +### 3. 线程安全 + +sqlite3 连接对象本身非线程安全。使用 `threading.RLock` 序列化所有 +`execute`/`commit` 调用，支持 `thread`/`async`/`dependency` 策略下的 +并发存取。`check_same_thread=False` 允许跨线程使用同一连接。 + +### 4. batch() 语义 + +`batch()` 不持有锁整个 duration（否则阻塞并发），仅设置 `_in_batch` 标志。 +各 `_put_raw` 调用各自加锁做 INSERT（不 commit），退出时单次 commit。 +与 `JSONBackend.batch()` 模式一致。 + +### 5. WAL 模式 + +启用 `PRAGMA journal_mode=WAL` + `synchronous=NORMAL`，提升并发读写性能。 +WAL 模式下读不阻塞写。 + +### 6. 不实现 `__del__` + +依赖 sqlite3 连接自身的 `__del__` 关闭；提供 `close()` 方法供显式清理。 +`__del__` + RLock 有死锁风险，避免。 + +## 验证结果 + +- ruff / pyrefly / pytest 全通过 +- 覆盖率 ≥ 95% +- 作为 JSONBackend drop-in 替换，现有测试通过 +- 并发线程安全测试通过 + +## 遗留事项 + +- P2 实时进度监控（下一迭代） +- P3 任务通知系统 +- P4 结果流式获取 diff --git a/src/pyflowx/__init__.py b/src/pyflowx/__init__.py index b402361..3f3ca9e 100644 --- a/src/pyflowx/__init__.py +++ b/src/pyflowx/__init__.py @@ -13,7 +13,7 @@ * :class:`GraphDefaults` —— 图级默认值。 * :func:`compose` —— 编程式组合多图。 * :func:`task_template` —— 批量生成相似 TaskSpec 的工厂。 -* 状态后端：:class:`StateBackend`、:class:`MemoryBackend`、:class:`JSONBackend`。 +* 状态后端：:class:`StateBackend`、:class:`MemoryBackend`、:class:`JSONBackend`、:class:`SQLiteBackend`。 快速上手 -------- @@ -77,7 +77,7 @@ from .graph import Graph, GraphDefaults from .profiling import ProfileReport, TaskProfile from .report import RunReport from .runner import CliExitCode, CliRunner -from .storage import JSONBackend, MemoryBackend, StateBackend +from .storage import JSONBackend, MemoryBackend, SQLiteBackend, StateBackend from .task import ( CacheKeyFn, RetryPolicy, @@ -120,6 +120,7 @@ __all__ = [ "PyFlowXError", "RetryPolicy", "RunReport", + "SQLiteBackend", "StateBackend", "StorageError", "Strategy", diff --git a/src/pyflowx/storage.py b/src/pyflowx/storage.py index c00bb97..9616aa6 100644 --- a/src/pyflowx/storage.py +++ b/src/pyflowx/storage.py @@ -14,7 +14,9 @@ from __future__ import annotations import json +import sqlite3 import sys +import threading import time from abc import ABC, abstractmethod from collections.abc import Iterator, Mapping @@ -277,6 +279,132 @@ class JSONBackend(_TTLStateBackendMixin): self._flush() +class SQLiteBackend(_TTLStateBackendMixin): + """基于 SQLite 文件的状态后端。 + + 相比 :class:`JSONBackend` 的全量加载/落盘，:class:`SQLiteBackend` 按键 + 查询、按行写入，适合大规模任务状态（10k+ 条目）。使用标准库 + :mod:`sqlite3`，无新依赖。 + + 值以 JSON 文本存入 ``value`` 列，因此与 :class:`JSONBackend` 一样要求 + 结果可 JSON 序列化。``ts`` 列存储写入时间戳用于 TTL 判断。 + + 线程安全：内部用 :class:`threading.RLock` 序列化所有 sqlite3 调用 + （sqlite3 连接对象本身非线程安全），支持 ``thread``/``async``/``dependency`` + 策略下的并发存取。 + + Parameters + ---------- + path: + SQLite 文件路径。``":memory:"`` 创建内存数据库（进程退出丢失， + 适合测试）。 + ttl: + 条目存活秒数。``None`` 表示永不过期。``has`` 在条目超过 ttl 后 + 返回 ``False``（但不主动删除，下次 ``save`` 覆盖）。 + """ + + def __init__(self, path: str, ttl: float | None = None) -> None: + self._path = path + self._ttl = ttl + self._lock = threading.RLock() + self._in_batch = False + try: + self._conn = sqlite3.connect(path, check_same_thread=False) + with self._lock: + self._conn.execute("PRAGMA journal_mode=WAL") + self._conn.execute("PRAGMA synchronous=NORMAL") + self._conn.execute( + "CREATE TABLE IF NOT EXISTS state (key TEXT PRIMARY KEY, value TEXT NOT NULL, ts REAL NOT NULL)" + ) + self._conn.execute("CREATE INDEX IF NOT EXISTS idx_state_ts ON state(ts)") + self._conn.commit() + except sqlite3.DatabaseError as exc: + raise StorageError(f"cannot open sqlite state file {path!r}", exc) from exc + + @override + def _get_raw(self, key: str) -> tuple[Any, float] | None: + with self._lock: + try: + row = self._conn.execute("SELECT value, ts FROM state WHERE key = ?", (key,)).fetchone() + except sqlite3.DatabaseError as exc: + raise StorageError(f"cannot read key {key!r} from sqlite state", exc) from exc + if row is None: + return None + value_text, ts = row + try: + value = json.loads(value_text) + except json.JSONDecodeError as exc: + raise StorageError(f"cannot decode value for key {key!r}", exc) from exc + return value, float(ts) + + @override + def _put_raw(self, key: str, value: Any, ts: float) -> None: + try: + value_text = json.dumps(value, ensure_ascii=False) + except (TypeError, ValueError) as exc: + raise StorageError(f"result of key {key!r} is not JSON-serialisable", exc) from exc + with self._lock: + try: + self._conn.execute( + "INSERT OR REPLACE INTO state (key, value, ts) VALUES (?, ?, ?)", + (key, value_text, ts), + ) + if not self._in_batch: + self._conn.commit() + except sqlite3.DatabaseError as exc: + raise StorageError(f"cannot write key {key!r} to sqlite state", exc) from exc + + @override + def _iter_raw(self) -> Iterator[tuple[str, Any, float]]: + with self._lock: + try: + rows = self._conn.execute("SELECT key, value, ts FROM state").fetchall() + except sqlite3.DatabaseError as exc: + raise StorageError("cannot iterate sqlite state", exc) from exc + for k, value_text, ts in rows: + try: + value = json.loads(value_text) + except json.JSONDecodeError as exc: + raise StorageError(f"cannot decode value for key {k!r}", exc) from exc + yield k, value, float(ts) + + @override + def _clear_raw(self) -> None: + with self._lock: + try: + self._conn.execute("DELETE FROM state") + self._conn.commit() + except sqlite3.DatabaseError as exc: + raise StorageError("cannot clear sqlite state", exc) from exc + + @override + def flush(self) -> None: + with self._lock: + try: + self._conn.commit() + except sqlite3.DatabaseError as exc: + raise StorageError("cannot flush sqlite state", exc) from exc + + @override + @contextmanager + def batch(self) -> Iterator[None]: + """进入批量模式：``save`` 暂不 commit，退出时统一 commit 一次。 + + 将整次运行 N 个任务的 N 次 commit 降为 1 次。 + """ + self._in_batch = True + try: + yield + finally: + self._in_batch = False + self.flush() + + def close(self) -> None: + """关闭数据库连接，释放文件句柄。""" + with self._lock: + self._conn.close() + + def resolve_backend(backend: StateBackend | None) -> StateBackend: """返回 ``backend``；为 ``None`` 时返回新的 :class:`MemoryBackend`。""" return backend if backend is not None else MemoryBackend() diff --git a/tests/test_storage.py b/tests/test_storage.py index 3a681ea..e9113bb 100644 --- a/tests/test_storage.py +++ b/tests/test_storage.py @@ -4,7 +4,9 @@ from __future__ import annotations import json import os +import sqlite3 import tempfile +import threading import time from pathlib import Path from typing import Any @@ -12,7 +14,7 @@ from typing import Any import pytest from pyflowx.errors import StorageError -from pyflowx.storage import JSONBackend, MemoryBackend, StateBackend, resolve_backend +from pyflowx.storage import JSONBackend, MemoryBackend, SQLiteBackend, StateBackend, resolve_backend @pytest.fixture @@ -294,6 +296,284 @@ def test_json_backend_save_value_error(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(_json, "dumps", original_dumps) +# ---------------------------------------------------------------------- # +# SQLiteBackend +# ---------------------------------------------------------------------- # +def test_sqlite_backend_save_and_load() -> None: + """save 后重新打开同一文件应读到已保存内容。""" + with tempfile.TemporaryDirectory() as tmp: + path = str(Path(tmp) / "state.db") + b = SQLiteBackend(path) + b.save("a", {"x": 1}) + b.save("b", [1, 2, 3]) + b.close() + b2 = SQLiteBackend(path) + assert b2.has("a") + assert b2.get("a") == {"x": 1} + assert b2.get("b") == [1, 2, 3] + assert dict(b2.load()) == {"a": {"x": 1}, "b": [1, 2, 3]} + b2.close() + + +def test_sqlite_backend_overwrite_existing() -> None: + """对同一 key 再次 save 应覆盖旧值。""" + with tempfile.TemporaryDirectory() as tmp: + path = str(Path(tmp) / "state.db") + b = SQLiteBackend(path) + b.save("a", 1) + b.save("a", 2) + assert b.get("a") == 2 + b.close() + + +def test_sqlite_backend_clear() -> None: + """clear 后所有键清空，且持久化到文件。""" + with tempfile.TemporaryDirectory() as tmp: + path = str(Path(tmp) / "state.db") + b = SQLiteBackend(path) + b.save("a", 1) + b.clear() + assert not b.has("a") + assert dict(b.load()) == {} + b.close() + # 重新打开仍为空 + b2 = SQLiteBackend(path) + assert dict(b2.load()) == {} + b2.close() + + +def test_sqlite_backend_nonexistent_file_starts_empty() -> None: + """文件不存在时应正常初始化为空。""" + with tempfile.TemporaryDirectory() as tmp: + path = str(Path(tmp) / "absent.db") + b = SQLiteBackend(path) + assert dict(b.load()) == {} + assert not b.has("anything") + b.close() + + +def test_sqlite_backend_non_serialisable_raises() -> None: + """不可 JSON 序列化的值应抛 StorageError，且不写入数据库。""" + with tempfile.TemporaryDirectory() as tmp: + path = str(Path(tmp) / "state.db") + b = SQLiteBackend(path) + with pytest.raises(StorageError): + b.save("a", object()) # object() 不可序列化 + assert not b.has("a") + b.close() + + +def test_sqlite_backend_get_missing_raises() -> None: + """get 不存在的键应抛 KeyError。""" + with tempfile.TemporaryDirectory() as tmp: + path = str(Path(tmp) / "state.db") + b = SQLiteBackend(path) + with pytest.raises(KeyError): + b.get("nope") + b.close() + + +def test_sqlite_backend_memory_db() -> None: + """:memory: 内存数据库基础生命周期。""" + b = SQLiteBackend(":memory:") + b.save("a", 1) + assert b.has("a") + assert b.get("a") == 1 + assert dict(b.load()) == {"a": 1} + b.clear() + assert not b.has("a") + b.close() + + +def test_sqlite_backend_batch_defers_commit(tmp_path: Path) -> None: + """batch 期间 save 不立即 commit，第二个连接看不到未提交数据；退出后可见。""" + path = str(tmp_path / "state.db") + b = SQLiteBackend(path) + with b.batch(): + b.save("a", 1) + b.save("b", 2) + # 第二个连接在 WAL 模式下只读已提交数据，应看不到 batch 内未提交的写入 + other = sqlite3.connect(path) + row = other.execute("SELECT value FROM state WHERE key = ?", ("a",)).fetchone() + other.close() + assert row is None + # batch 退出 commit 后，第二个连接应能看到 + other = sqlite3.connect(path) + rows = other.execute("SELECT key FROM state ORDER BY key").fetchall() + other.close() + assert rows == [("a",), ("b",)] + b.close() + + +def test_sqlite_backend_no_batch_commits_each_save(tmp_path: Path) -> None: + """无 batch 时每次 save 立即 commit，第二个连接立即可见。""" + path = str(tmp_path / "state.db") + b = SQLiteBackend(path) + b.save("a", 1) + # 第二个连接应立即看到已提交数据 + other = sqlite3.connect(path) + row = other.execute("SELECT value FROM state WHERE key = ?", ("a",)).fetchone() + other.close() + assert row is not None + assert json.loads(row[0]) == 1 + b.save("b", 2) + other = sqlite3.connect(path) + rows = other.execute("SELECT key FROM state ORDER BY key").fetchall() + other.close() + assert rows == [("a",), ("b",)] + b.close() + + +def test_sqlite_backend_concurrent_access(tmp_path: Path) -> None: + """多线程并发 save/get 不出错，结果一致。""" + path = str(tmp_path / "state.db") + b = SQLiteBackend(path) + errors: list[BaseException] = [] + + def worker(start: int) -> None: + try: + for i in range(start, start + 20): + b.save(f"task_{i}", i) + assert b.get(f"task_{i}") == i + except BaseException as exc: + errors.append(exc) + + threads = [threading.Thread(target=worker, args=(i * 20,)) for i in range(4)] + for t in threads: + t.start() + for t in threads: + t.join() + assert not errors, f"并发访问出错: {errors}" + # 共写入 80 个键 + assert len(dict(b.load())) == 80 + b.close() + + +def test_sqlite_backend_ttl_expired_has_returns_false() -> None: + """SQLiteBackend TTL 过期后 has 返回 False。""" + with tempfile.TemporaryDirectory() as tmp: + path = str(Path(tmp) / "state.db") + b = SQLiteBackend(path, ttl=0.1) + b.save("a", 1) + assert b.has("a") + time.sleep(0.15) + assert not b.has("a") + b.close() + + +def test_sqlite_backend_ttl_expired_get_raises_keyerror() -> None: + """SQLiteBackend TTL 过期后 get 抛 KeyError。""" + with tempfile.TemporaryDirectory() as tmp: + path = str(Path(tmp) / "state.db") + b = SQLiteBackend(path, ttl=0.1) + b.save("a", 1) + time.sleep(0.15) + with pytest.raises(KeyError): + b.get("a") + b.close() + + +def test_sqlite_backend_ttl_load_filters_expired() -> None: + """SQLiteBackend.load() 应过滤过期的条目。""" + with tempfile.TemporaryDirectory() as tmp: + path = str(Path(tmp) / "state.db") + b = SQLiteBackend(path, ttl=0.1) + b.save("a", 1) + b.save("b", 2) + time.sleep(0.15) + assert dict(b.load()) == {} + b.close() + + +def test_sqlite_backend_no_ttl_never_expired(tmp_path: Path) -> None: + """无 TTL 时永不过期。""" + path = str(tmp_path / "state.db") + b = SQLiteBackend(path) + b.save("a", 1) + # 手动修改 ts 为很久以前 + with b._lock: + b._conn.execute("UPDATE state SET ts = ? WHERE key = ?", (time.time() - 1000, "a")) + b._conn.commit() + assert b.has("a") + assert b.get("a") == 1 + b.close() + + +def test_sqlite_backend_expired_with_ttl(tmp_path: Path) -> None: + """有 TTL 时过期键 has 返回 False。""" + path = str(tmp_path / "state.db") + b = SQLiteBackend(path, ttl=1.0) + b.save("a", 1) + # 手动修改 ts 为很久以前 + with b._lock: + b._conn.execute("UPDATE state SET ts = ? WHERE key = ?", (time.time() - 10, "a")) + b._conn.commit() + assert not b.has("a") + b.close() + + +def test_sqlite_backend_corrupt_value_raises(tmp_path: Path) -> None: + """value 列被外部写入非法 JSON 时 get 抛 StorageError。""" + path = str(tmp_path / "state.db") + b = SQLiteBackend(path) + b.save("a", 1) + b.close() + # 直接用 sqlite3 写入非法 JSON + conn = sqlite3.connect(path) + conn.execute("UPDATE state SET value = ? WHERE key = ?", ("{not valid json", "a")) + conn.commit() + conn.close() + b2 = SQLiteBackend(path) + with pytest.raises(StorageError, match="cannot decode"): + b2.get("a") + b2.close() + + +def test_sqlite_backend_iter_raw_decode_error_raises(tmp_path: Path) -> None: + """_iter_raw 遇到非法 JSON 时 load 抛 StorageError。""" + path = str(tmp_path / "state.db") + b = SQLiteBackend(path) + b.save("good", 1) + b.close() + # 写入非法 JSON + conn = sqlite3.connect(path) + conn.execute("INSERT INTO state (key, value, ts) VALUES (?, ?, ?)", ("bad", "{broken", 1.0)) + conn.commit() + conn.close() + b2 = SQLiteBackend(path) + with pytest.raises(StorageError, match="cannot decode"): + dict(b2.load()) + b2.close() + + +def test_sqlite_backend_close_releases(tmp_path: Path) -> None: + """close 后连接关闭，所有操作抛 StorageError（包装 ProgrammingError）。""" + path = str(tmp_path / "state.db") + b = SQLiteBackend(path) + b.save("a", 1) + b.close() + with pytest.raises(StorageError, match="cannot read"): + b.has("a") + with pytest.raises(StorageError, match="cannot read"): + b.get("a") + with pytest.raises(StorageError, match="cannot write"): + b.save("b", 2) + with pytest.raises(StorageError, match="cannot iterate"): + dict(b.load()) + with pytest.raises(StorageError, match="cannot clear"): + b.clear() + with pytest.raises(StorageError, match="cannot flush"): + b.flush() + + +def test_sqlite_backend_corrupt_file_raises(tmp_path: Path) -> None: + """打开非 SQLite 格式的文件应抛 StorageError。""" + path = tmp_path / "not_a_db.db" + _ = path.write_bytes(b"not a sqlite database file content") + with pytest.raises(StorageError, match="cannot open sqlite state file"): + _ = SQLiteBackend(str(path)) + + # ---------------------------------------------------------------------- # # resolve_backend # ---------------------------------------------------------------------- #