cloud: fill gaps (memory/KB/cost/models endpoints, WAL concurrency, OAuth scaffold, push subscriptions)

cloud/tests/test_db.py

@@ -0,0 +1,249 @@
+"""cloud/db.py 单测：SQLite 并发 helper。
+
+只测 db.py，不改 server.py / host.py / 其它模块。
+重试逻辑的覆盖方式（本文件写清楚，供协调员核对）：
+- 确定性 fake 方式：FakeConn 前两次抛 database is locked、第三次
+  成功，断言重试后成功且退避序列为 0.05/0.1/0.2s；另测重试耗尽
+  抛原始错误、非锁错误一次都不重试。
+- 真实并发方式：8 线程各写 20 条（server.py 同款 tasks 表），
+  断言无错误、最终 160 条、id 无重复。
+两者都测了，前者证重试逻辑本身正确，后者证真实写冲突不再直接炸掉。
+"""
+
+import os
+import sqlite3
+import sys
+import tempfile
+import threading
+import time
+import unittest
+from pathlib import Path
+from unittest import mock
+
+try:
+    from cloud import db as db_module
+    from cloud.db import (
+        connect,
+        execute_with_retry,
+        is_locked_error,
+        read_all,
+        transaction,
+    )
+except ImportError:  # discover -s cloud/tests 时把项目根目录补进 sys.path
+    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
+    from cloud import db as db_module
+    from cloud.db import (
+        connect,
+        execute_with_retry,
+        is_locked_error,
+        read_all,
+        transaction,
+    )
+
+# 与 server.py init_db 的 tasks 表定义逐字段一致（db.py 本身不建表，
+# 这里复刻定义只是为了让并发测试写进同款表，证明 schema 兼容）。
+TASKS_DDL = (
+    "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 REAL, "
+    "updated_at REAL, "
+    "last_heartbeat REAL, "
+    "UNIQUE(project_id, idempotency_key))"
+)
+
+
+class FakeConn:
+    """受控 fake：前 fail_times 次 execute 抛锁错误，之后返回哨兵。"""
+
+    def __init__(self, fail_times, error_text="database is locked"):
+        self.fail_times = fail_times
+        self.error_text = error_text
+        self.calls = 0
+
+    def execute(self, sql, params=()):
+        self.calls += 1
+        if self.calls <= self.fail_times:
+            raise sqlite3.OperationalError(self.error_text)
+        return "ok-sentinel"
+
+
+class DbHelperTestCase(unittest.TestCase):
+    def setUp(self):
+        self.tmp = tempfile.TemporaryDirectory()
+        self.db_path = os.path.join(self.tmp.name, "concurrency.db")
+        conn = connect(self.db_path)
+        try:
+            conn.execute(TASKS_DDL)
+            conn.commit()
+        finally:
+            conn.close()
+
+    def tearDown(self):
+        self.tmp.cleanup()
+
+    # -- 连接属性：WAL + busy_timeout + Row ------------------------------
+    def test_connect_sets_wal_busy_timeout_row_factory(self):
+        conn = connect(self.db_path)
+        try:
+            self.assertEqual(conn.row_factory, sqlite3.Row)
+            journal = conn.execute(
+                "PRAGMA journal_mode").fetchone()[0]
+            self.assertEqual(journal, "wal")
+            busy = conn.execute("PRAGMA busy_timeout").fetchone()[0]
+            self.assertEqual(busy, 5000)
+            # Row 行为：按列名访问，与 server.py 的 _task_to_dict 兼容。
+            conn.execute(
+                "INSERT INTO tasks(id, project_id, prompt, status)"
+                " VALUES(?, ?, ?, 'queued')",
+                ("t-row", "p1", "hello"),
+            )
+            conn.commit()
+            row = read_all(
+                conn, "SELECT * FROM tasks WHERE id = ?", ("t-row",))[0]
+            self.assertEqual(row["id"], "t-row")
+            self.assertEqual(row["prompt"], "hello")
+        finally:
+            conn.close()
+
+    def test_connect_custom_busy_timeout(self):
+        conn = connect(self.db_path, busy_timeout_ms=1234)
+        try:
+            busy = conn.execute("PRAGMA busy_timeout").fetchone()[0]
+            self.assertEqual(busy, 1234)
+        finally:
+            conn.close()
+
+    # -- 锁判定 ---------------------------------------------------------
+    def test_is_locked_error_only_for_locked_operational_error(self):
+        self.assertTrue(is_locked_error(
+            sqlite3.OperationalError("database is locked")))
+        self.assertTrue(is_locked_error(
+            sqlite3.OperationalError("Database Is Locked (retry)")))
+        self.assertFalse(is_locked_error(
+            sqlite3.OperationalError("no such table: tasks")))
+        self.assertFalse(is_locked_error(ValueError("database is locked")))
+
+    # -- 重试逻辑：fake 受控方式（确定性）--------------------------------
+    def test_retry_succeeds_after_two_locked_failures(self):
+        fake = FakeConn(fail_times=2)
+        with mock.patch.object(db_module.time, "sleep") as sleep_mock:
+            result = execute_with_retry(fake, "INSERT ...", ())
+        self.assertEqual(result, "ok-sentinel")
+        self.assertEqual(fake.calls, 3)  # 1 次初始 + 2 次重试
+        self.assertEqual(
+            [c.args[0] for c in sleep_mock.call_args_list],
+            [0.05, 0.1, 0.2][:2])
+
+    def test_retry_backoff_sequence_is_exponential(self):
+        fake = FakeConn(fail_times=3)
+        with mock.patch.object(db_module.time, "sleep") as sleep_mock:
+            execute_with_retry(fake, "INSERT ...", ())
+        self.assertEqual(
+            [c.args[0] for c in sleep_mock.call_args_list],
+            [0.05, 0.1, 0.2])
+
+    def test_retry_exhausted_reraises_original_error(self):
+        fake = FakeConn(fail_times=99)
+        with mock.patch.object(db_module.time, "sleep"):
+            with self.assertRaises(sqlite3.OperationalError) as ctx:
+                execute_with_retry(fake, "INSERT ...", ())
+        self.assertIn("database is locked", str(ctx.exception))
+        # 默认 max_retries=3：1 次初始 + 3 次重试 = 4 次尝试，不多不少。
+        self.assertEqual(fake.calls, 4)
+
+    def test_non_locked_error_is_not_retried(self):
+        fake = FakeConn(fail_times=99, error_text="no such table: tasks")
+        with mock.patch.object(db_module.time, "sleep") as sleep_mock:
+            with self.assertRaises(sqlite3.OperationalError):
+                execute_with_retry(fake, "SELECT ...", ())
+        self.assertEqual(fake.calls, 1)
+        sleep_mock.assert_not_called()
+
+    # -- 事务：成功提交 / 失败回滚 ---------------------------------------
+    def test_transaction_commits_and_rolls_back(self):
+        conn = connect(self.db_path)
+        try:
+            with transaction(conn):
+                execute_with_retry(
+                    conn,
+                    "INSERT INTO tasks(id, project_id, prompt, status)"
+                    " VALUES(?, ?, ?, 'queued')",
+                    ("t-ok", "p1", "ok"),
+                )
+            self.assertEqual(
+                len(read_all(conn, "SELECT * FROM tasks")), 1)
+            with self.assertRaises(sqlite3.IntegrityError):
+                with transaction(conn):
+                    execute_with_retry(
+                        conn,
+                        "INSERT INTO tasks(id, project_id, prompt, status)"
+                        " VALUES(?, ?, ?, 'queued')",
+                        ("t-ok", "p1", "dup"),
+                    )
+            # 冲突行回滚后仍只有原来 1 条，不留半写。
+            self.assertEqual(
+                len(read_all(conn, "SELECT * FROM tasks")), 1)
+        finally:
+            conn.close()
+
+    # -- 真实并发：8 线程 x 20 条 ----------------------------------------
+    def test_concurrent_writes_8_threads_x_20_no_loss(self):
+        errors = []
+        barrier = threading.Barrier(8)
+
+        def worker(thread_idx):
+            try:
+                barrier.wait(timeout=10)
+                conn = connect(self.db_path)
+                try:
+                    for i in range(20):
+                        with transaction(conn):
+                            execute_with_retry(
+                                conn,
+                                "INSERT INTO tasks(id, project_id, prompt,"
+                                " status, idempotency_key, depends_on,"
+                                " priority, created_by, created_at,"
+                                " updated_at, last_heartbeat)"
+                                " VALUES(?, ?, ?, 'queued', ?, '[]',"
+                                " 0, 'worker', ?, ?, NULL)",
+                                ("t-%d-%d" % (thread_idx, i),
+                                 "p-concurrent",
+                                 "prompt %d-%d" % (thread_idx, i),
+                                 "key-%d-%d" % (thread_idx, i),
+                                 time.time(), time.time()),
+                            )
+                finally:
+                    conn.close()
+            except Exception as exc:  # noqa: BLE001 - 汇总后统一断言
+                errors.append(exc)
+
+        threads = [threading.Thread(target=worker, args=(n,))
+                   for n in range(8)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join(timeout=60)
+        self.assertFalse([t for t in threads if t.is_alive()],
+                         "worker 线程未在 60s 内结束")
+        self.assertEqual(errors, [], "并发写不应有任何异常: %r" % (errors,))
+        conn = connect(self.db_path)
+        try:
+            rows = read_all(
+                conn,
+                "SELECT id FROM tasks WHERE project_id = ?",
+                ("p-concurrent",))
+        finally:
+            conn.close()
+        self.assertEqual(len(rows), 160)
+        self.assertEqual(len({r["id"] for r in rows}), 160)
+
+
+if __name__ == "__main__":
+    unittest.main()