codex-app-server/cloud/tests/test_postgres_backend.py

"""Postgres 后端单测：适配层（fake 驱动，常跑）+ 真实 PG（skip）。

分两类，泾渭分明：
- FakeDriver 类：不连任何真实 Postgres，验证占位符转换、
  行按名访问、唯一冲突映射、事务回滚、锁重试、schema 语句
  覆盖——这些在无 Postgres 的机器（含 CI runner）上也必须通过。
- RealPostgres 类：需要真实 Postgres 时才跑，判定条件见
  _real_pg_params()：须显式设置 CLOUD_TEST_PG_DSN（形如
  postgres://user:pass@host:port/dbname）且主机端口 2 秒内
  可达，否则 skipTest，绝不 fail。DSN 必须指向一次性测试库，
  本文件不在其中建生产表之外的任何东西（schema 即业务表）。

测试里出现的密码一律是 fake/test 占位（"test-placeholder"），
绝不写真实凭证；DSN 只从环境变量读，测试不打印它。
"""

import os
import socket
import sys
import traceback
import unittest
from multiprocessing import Barrier, Process, Queue
from pathlib import Path
from unittest import mock
from urllib.parse import unquote, urlparse

try:
    from cloud.db import (
        POSTGRES_SCHEMA,
        PgConnection,
        PostgresDatabase,
        PostgresDriverMissingError,
        PostgresIntegrityError,
        PostgresOperationalError,
        _convert_placeholders,
        _translate_sql,
        init_postgres_schema,
    )
except ImportError:  # discover -s cloud/tests 时把项目根目录补进 sys.path
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
    from cloud.db import (
        POSTGRES_SCHEMA,
        PgConnection,
        PostgresDatabase,
        PostgresDriverMissingError,
        PostgresIntegrityError,
        PostgresOperationalError,
        _convert_placeholders,
        _translate_sql,
        init_postgres_schema,
    )


def _has_real_driver():
    for name in ("psycopg", "psycopg2"):
        try:
            __import__(name)
            return True
        except ImportError:
            continue
    return False


def _real_pg_params():
    """返回 (host, port, dbname, user, password) 或 None。

    仅当 CLOUD_TEST_PG_DSN 已设且主机端口 2 秒内可达时返回；
    其余情况返回 None，调用方 skipTest。绝不打印 DSN。
    """
    dsn = os.environ.get("CLOUD_TEST_PG_DSN")
    if not dsn:
        return None
    parsed = urlparse(dsn)
    if parsed.scheme not in ("postgres", "postgresql") or not parsed.hostname:
        return None
    host = parsed.hostname
    port = parsed.port or 5432
    try:
        with socket.create_connection((host, port), timeout=2):
            pass
    except OSError:
        return None
    return (
        host,
        port,
        parsed.path.lstrip("/") or "postgres",
        unquote(parsed.username) if parsed.username else None,
        unquote(parsed.password) if parsed.password else None,
    )


class _FakeUnique(Exception):
    pass


class _FakeDeadlock(Exception):
    pass


class _FakeCursor:
    """记录 SQL、按配置失败或返回预置行；description 为列名元组。"""

    def __init__(self, rows=None, fail_sequence=None):
        self.executed = []
        self.rows = list(rows or [])
        self.fail_sequence = list(fail_sequence or [])
        self.description = [("id",), ("name",)]

    def execute(self, sql, params=()):
        self.executed.append((sql, params))
        if self.fail_sequence:
            exc = self.fail_sequence.pop(0)
            if exc is not None:
                raise exc
        return self

    def fetchall(self):
        return list(self.rows)

    def fetchone(self):
        return self.rows[0] if self.rows else None

    @property
    def rowcount(self):
        return len(self.rows) if self.rows else 1

    def close(self):
        pass


class _FakeConn:
    def __init__(self, cursor=None):
        self.cursor_obj = cursor or _FakeCursor()
        self.committed = False
        self.rolled_back = False
        self.closed = False

    def cursor(self):
        return self.cursor_obj

    def commit(self):
        self.committed = True

    def rollback(self):
        self.rolled_back = True

    def close(self):
        self.closed = True


class _FakeDriver:
    UniqueViolation = _FakeUnique
    DeadlockDetected = _FakeDeadlock
    OperationalError = Exception
    Error = Exception

    def __init__(self, cursor=None):
        self.conn = _FakeConn(cursor)

    def connect(self, **kwargs):
        self.connect_kwargs = kwargs
        return self.conn


def _make_db(fake):
    with mock.patch("cloud.db._load_postgres_driver",
                    return_value=("psycopg", fake)):
        return PostgresDatabase(host="test-host", dbname="test-db",
                                user="test-user",
                                password="test-placeholder")


class DriverMissingTestCase(unittest.TestCase):
    def test_missing_driver_raises_psycopg_error(self):
        if _has_real_driver():
            self.skipTest("环境已装 psycopg 驱动，无驱动报错路径不适用")
        with self.assertRaises(PostgresDriverMissingError) as ctx:
            PostgresDatabase(host="test-host", dbname="test-db")
        self.assertIn("psycopg", str(ctx.exception))


class PlaceholderTranslationTestCase(unittest.TestCase):
    def test_question_marks_outside_strings_become_percent_s(self):
        self.assertEqual(
            _convert_placeholders(
                "SELECT * FROM t WHERE a = ? AND b = 'x?y' AND c = ?"),
            "SELECT * FROM t WHERE a = %s AND b = 'x?y' AND c = %s")
        self.assertEqual(
            _convert_placeholders("SELECT 'it''s ?' , ?"),
            "SELECT 'it''s ?' , %s")

    def test_insert_or_replace_becomes_on_conflict_update(self):
        translated, kind = _translate_sql(
            "INSERT OR REPLACE INTO project_members"
            "(project_id, user_id, role) VALUES(?, ?, ?)")
        self.assertEqual(kind, "normal")
        self.assertIn("ON CONFLICT (project_id, user_id) DO UPDATE SET",
                      translated)
        self.assertIn("role = EXCLUDED.role", translated)
        self.assertIn("%s", translated)
        self.assertNotIn("?", translated)

    def test_auto_memory_tables_replace_updates_not_ignores(self):
        # 2026-10-28 生产教训：这三张表未登记主键时翻译退化为
        # ON CONFLICT DO NOTHING，首写之后的更新静默丢弃；SQLite
        # 原生 OR REPLACE 测不出来，必须在翻译层断言。
        cases = [
            ("INSERT OR REPLACE INTO memory_settings"
             "(user_id, generate_memories, updated_at) VALUES(?, ?, ?)",
             "ON CONFLICT (user_id) DO UPDATE SET",
             "generate_memories = EXCLUDED.generate_memories"),
            ("INSERT OR REPLACE INTO task_memory_flags"
             "(task_id, use_memories, generate_memories, updated_at)"
             " VALUES(?, ?, ?, ?)",
             "ON CONFLICT (task_id) DO UPDATE SET",
             "use_memories = EXCLUDED.use_memories"),
            ("INSERT OR REPLACE INTO memory_extraction_log"
             "(task_id, status, reason, processed_at) VALUES(?, ?, ?, ?)",
             "ON CONFLICT (task_id) DO UPDATE SET",
             "status = EXCLUDED.status"),
        ]
        for sql, conflict, setter in cases:
            translated, kind = _translate_sql(sql)
            self.assertEqual(kind, "normal")
            self.assertIn(conflict, translated)
            self.assertIn(setter, translated)
            self.assertNotIn("DO NOTHING", translated)
            self.assertNotIn("OR REPLACE", translated)

    def test_autoincrement_becomes_identity(self):
        translated, _ = _translate_sql(
            "CREATE TABLE t(id INTEGER PRIMARY KEY AUTOINCREMENT, x TEXT)")
        self.assertIn("GENERATED BY DEFAULT AS IDENTITY", translated)
        self.assertNotIn("AUTOINCREMENT", translated)

    def test_fts5_is_rejected_not_silently_created(self):
        with self.assertRaises(PostgresOperationalError):
            _translate_sql(
                "CREATE VIRTUAL TABLE kb_fts USING fts5(content)")


class AdapterFakeDriverTestCase(unittest.TestCase):
    def test_row_name_and_index_access(self):
        fake = _FakeDriver(_FakeCursor(rows=[("t1", "hello")]))
        db = _make_db(fake)
        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")
        # 占位符已转换：fake 游标收到的 SQL 应含 %s 不含 ?。
        self.assertIn("%s", fake.conn.cursor_obj.executed[0][0])
        self.assertNotIn("?", fake.conn.cursor_obj.executed[0][0])

    def test_pragma_info_notnull_flag_matches_sqlite(self):
        # PG 实证（2026-11-06）抓到的真 bug：适配层把
        # information_schema.is_nullable 翻反了（NO -> 0），
        # 与 sqlite PRAGMA table_info 的 notnull 语义相反。
        # 此前无调用方读 row[3]，故只潜伏未炸；ensure 的
        # outcome 放宽依赖该标志为真，此处钉住。
        rows = [(1, "task_id", "text", "NO", None, 1),
                (2, "outcome", "text", "YES", None, 0)]
        fake = _FakeDriver(_FakeCursor(rows=rows))
        db = _make_db(fake)
        conn = db.connect()
        info = conn.execute(
            "PRAGMA table_info(task_attempts)").fetchall()
        by_name = {row[1]: row for row in info}
        self.assertEqual(by_name["task_id"][3], 1)
        self.assertEqual(by_name["outcome"][3], 0)
        self.assertEqual([row[1] for row in info],
                         ["task_id", "outcome"])

    def test_unique_violation_maps_to_integrity_error(self):
        import sqlite3
        fake = _FakeDriver(_FakeCursor(fail_sequence=[_FakeUnique("dup")]))
        db = _make_db(fake)
        conn = db.connect()
        with self.assertRaises(PostgresIntegrityError) as ctx:
            conn.execute("INSERT INTO t VALUES(?)", ("x",))
        # 不是 sqlite3 异常：调用方按 sqlite 异常捕获时会漏掉，
        # 因此 server.py 已同步捕获本模块错误（见该文件）。
        self.assertNotIsInstance(ctx.exception, sqlite3.IntegrityError)

    def test_transaction_rolls_back_on_error(self):
        fake = _FakeDriver(_FakeCursor(fail_sequence=[None, _FakeUnique("dup")]))
        db = _make_db(fake)
        conn = db.connect()
        with self.assertRaises(PostgresIntegrityError):
            with db.transaction(conn):
                conn.execute("INSERT INTO t VALUES(?)", ("a",))
                conn.execute("INSERT INTO t VALUES(?)", ("b",))
        self.assertTrue(fake.conn.rolled_back)

    def test_execute_with_retry_retries_only_locked_errors(self):
        fake = _FakeDriver(_FakeCursor(
            fail_sequence=[_FakeDeadlock("deadlock"), None]))
        db = _make_db(fake)
        conn = db.connect()
        # DeadlockDetected 在适配层映射为 locked=True，应重试一次后成功。
        cur = db.execute_with_retry(conn, "UPDATE t SET x = ?", ("y",),
                                    backoff_seconds=(0,))
        self.assertEqual(len(fake.conn.cursor_obj.executed), 2)
        self.assertIsNotNone(cur)

    def test_repr_never_contains_password(self):
        fake = _FakeDriver()
        db = _make_db(fake)
        self.assertNotIn("test-placeholder", repr(db))
        self.assertNotIn("test-placeholder", repr(db.connect()))


class SchemaStatementTestCase(unittest.TestCase):
    EXPECTED_TABLES = (
        "users", "project_members", "audit_log", "projects", "tasks",
        "summaries", "kb_docs", "cost_records", "oauth_states",
        "slack_users", "push_subscriptions",
    )

    def test_schema_covers_all_business_tables(self):
        joined = "\n".join(POSTGRES_SCHEMA)
        for table in self.EXPECTED_TABLES:
            self.assertIn("CREATE TABLE IF NOT EXISTS %s(" % table, joined)
        # tasks 含迁移补的 6 列；绝无 FTS5 / AUTOINCREMENT 残留。
        tasks_stmt = [s for s in POSTGRES_SCHEMA
                      if "CREATE TABLE IF NOT EXISTS tasks(" in s][0]
        for col in ("result_text", "error_text", "actual_model",
                    "started_at", "finished_at", "child_pid"):
            self.assertIn(col, tasks_stmt)
        self.assertNotIn("fts5", joined.lower())
        self.assertNotIn("AUTOINCREMENT", joined)

    def test_schema_builds_via_fake_driver(self):
        fake = _FakeDriver()
        db = _make_db(fake)
        conn = db.connect()
        init_postgres_schema(conn)
        executed = [sql for sql, _ in fake.conn.cursor_obj.executed]
        # 改动理由：init_postgres_schema 现在还要给老库 projects 表补
        # repo_url 列（PRAGMA 查列 + ALTER TABLE），所以执行语句数
        # 不再精确等于 schema 语句数；断言改为“schema 语句全部执行
        # 且补列语句也执行”，不放宽 schema 本身的覆盖要求。
        for stmt in POSTGRES_SCHEMA:
            self.assertIn(stmt, executed)
        self.assertTrue(any("repo_url" in sql for sql in executed))
        self.assertTrue(fake.conn.committed)


class RealPostgresSchemaTestCase(unittest.TestCase):
    """真实 PG schema：无 DSN 或不可达时 skip。"""

    def test_schema_creates_tables_and_roundtrips(self):
        params = _real_pg_params()
        if params is None:
            self.skipTest("未设置可达的 CLOUD_TEST_PG_DSN，跳过真实 PG 测试")
        if not _has_real_driver():
            self.skipTest("未装 psycopg 驱动，跳过真实 PG 测试")
        host, port, dbname, user, password = params
        db = PostgresDatabase(host=host, port=port, dbname=dbname,
                              user=user, password=password)
        conn = db.connect()
        try:
            init_postgres_schema(conn)
            rows = db.read_all(
                conn,
                "SELECT table_name FROM information_schema.tables"
                " WHERE table_schema = 'public'",
                ())
            names = {r["table_name"] for r in rows}
            for table in SchemaStatementTestCase.EXPECTED_TABLES:
                self.assertIn(table, names)
            # 最小往返：插一条项目再按名读回，随后删掉测试行。
            with db.transaction(conn):
                conn.execute(
                    "INSERT INTO projects(id, name, owner_id, created_at)"
                    " VALUES(?, ?, ?, ?)",
                    ("pg-test-proj", "pg-test", "pg-test-user", 0.0))
            rows = db.read_all(
                conn, "SELECT * FROM projects WHERE id = ?",
                ("pg-test-proj",))
            self.assertEqual(len(rows), 1)
            self.assertEqual(rows[0]["name"], "pg-test")
        finally:
            try:
                with db.transaction(conn):
                    conn.execute("DELETE FROM projects WHERE id = ?",
                                 ("pg-test-proj",))
            except Exception:
                pass
            conn.close()


def _concurrent_schema_worker(barrier, queue, wid, params):
    """并发初始化子进程：与其它进程同时跑 init_postgres_schema。"""
    try:
        host, port, dbname, user, password = params
        barrier.wait(timeout=60)
        db = PostgresDatabase(host=host, port=port, dbname=dbname,
                              user=user, password=password)
        conn = db.connect()
        try:
            init_postgres_schema(conn)
        finally:
            conn.close()
        queue.put((wid, "ok", ""))
    except Exception:
        queue.put((wid, "FAIL", traceback.format_exc()))


class RealPostgresConcurrentInitTestCase(unittest.TestCase):
    """真实 PG 并发初始化：无 DSN 或不可达时 skip。

    回归：启动竞态（37ba687 与 7983567 两次部署均撞上）——
    cloud-server 与 worker 并发 init 时，CREATE TABLE / CREATE
    INDEX 的 IF NOT EXISTS 检查与建对象不是原子的，并发建同名
    对象会撞 pg_type_typname_nsp_index（建表）或
    pg_class_relname_nsp_index（建索引）的 UniqueViolation。
    修法是 schema 初始化持 advisory 锁串行化；本用例在空库上
    4 进程同时初始化，必须全部成功、表齐全、无一例异常外泄。
    """

    def test_concurrent_init_postgres_schema_no_unique_violation(self):
        params = _real_pg_params()
        if params is None:
            self.skipTest("未设置可达的 CLOUD_TEST_PG_DSN，跳过真实 PG 测试")
        if not _has_real_driver():
            self.skipTest("未装 psycopg 驱动，跳过真实 PG 测试")
        host, port, dbname, user, password = params
        db = PostgresDatabase(host=host, port=port, dbname=dbname,
                              user=user, password=password)
        conn = db.connect()
        try:
            conn.execute("DROP SCHEMA public CASCADE")
            conn.execute("CREATE SCHEMA public")
            conn.commit()
        finally:
            conn.close()
        nworkers = 4
        barrier = Barrier(nworkers)
        queue = Queue()
        procs = [Process(target=_concurrent_schema_worker,
                         args=(barrier, queue, i, params))
                 for i in range(nworkers)]
        try:
            for proc in procs:
                proc.start()
            for proc in procs:
                proc.join(timeout=120)
        finally:
            for proc in procs:
                if proc.is_alive():
                    proc.terminate()
        results = [queue.get() for _ in range(nworkers)]
        failures = [r for r in results if r[1] == "FAIL"]
        self.assertEqual([], failures,
                         "并发初始化有进程失败：%s" % (failures,))
        conn = db.connect()
        try:
            rows = db.read_all(
                conn,
                "SELECT table_name FROM information_schema.tables"
                " WHERE table_schema = 'public'",
                ())
            names = {r["table_name"] for r in rows}
            for table in SchemaStatementTestCase.EXPECTED_TABLES:
                self.assertIn(table, names)
        finally:
            conn.close()