deer-flow/backend/tests/test_persistence_engine_postgres_config.py

"""Tests for hardened PostgreSQL async engine configuration."""

from __future__ import annotations

import asyncio
import sys
from time import monotonic
from types import ModuleType
from unittest.mock import AsyncMock, MagicMock, call, patch

import pytest
import yaml
from pydantic import ValidationError

from deerflow.config.database_config import DatabaseConfig
from deerflow.persistence import engine as engine_mod


class _PostgresError(Exception):
    """Driver error carrying a SQLSTATE, as asyncpg's exceptions do."""

    def __init__(self, sqlstate: str, message: str) -> None:
        super().__init__(message)
        self.sqlstate = sqlstate


def test_postgres_engine_kwargs_include_connection_hardening() -> None:
    kwargs = engine_mod._postgres_engine_kwargs(echo=False, pool_size=5)

    assert kwargs["echo"] is False
    assert kwargs["pool_size"] == 5
    assert kwargs["pool_pre_ping"] is True
    assert kwargs["pool_recycle"] == engine_mod.POSTGRES_POOL_RECYCLE_SECONDS
    assert kwargs["connect_args"]["command_timeout"] == engine_mod.POSTGRES_COMMAND_TIMEOUT_SECONDS
    assert kwargs["json_serializer"] is engine_mod._json_serializer


def test_database_command_timeout_defaults_to_30_seconds() -> None:
    config = DatabaseConfig()

    assert config.command_timeout == 30


def test_database_pool_recycle_defaults_to_300_seconds() -> None:
    config = DatabaseConfig()

    assert config.pool_recycle == 300


def test_postgres_engine_kwargs_preserve_caller_values() -> None:
    kwargs = engine_mod._postgres_engine_kwargs(echo=True, pool_size=20, pool_recycle=120, command_timeout=90)

    assert kwargs["echo"] is True
    assert kwargs["pool_size"] == 20
    assert kwargs["pool_recycle"] == 120
    assert kwargs["connect_args"] == {"command_timeout": 90}


def test_postgres_engine_kwargs_allow_command_timeout_opt_out() -> None:
    config = DatabaseConfig(command_timeout=None)
    kwargs = engine_mod._postgres_engine_kwargs(echo=False, pool_size=5, command_timeout=config.command_timeout)

    assert config.command_timeout is None
    assert kwargs["connect_args"] == {}


@pytest.mark.parametrize(
    ("raw", "expected"),
    [
        (None, None),
        (30, 30.0),
        (0.5, 0.5),
        ("60", 60.0),
    ],
)
def test_database_command_timeout_accepts_finite_positive_seconds(raw, expected: float | None) -> None:
    assert DatabaseConfig(command_timeout=raw).command_timeout == expected


@pytest.mark.parametrize(
    "raw",
    [
        True,
        False,
        yaml.safe_load("on"),
        yaml.safe_load("off"),
        yaml.safe_load("yes"),
        yaml.safe_load("no"),
        0,
        -1,
        -0.5,
        float("inf"),
        float("-inf"),
        float("nan"),
        yaml.safe_load(".inf"),
        yaml.safe_load(".nan"),
        "1e999",
    ],
)
def test_database_command_timeout_rejects_boolean_and_non_finite_values(raw) -> None:
    with pytest.raises(ValidationError, match="command_timeout"):
        DatabaseConfig(command_timeout=raw)


@pytest.mark.parametrize(
    ("field", "invalid_value"),
    [
        ("pool_size", True),
        ("pool_size", False),
        ("pool_size", 0),
        ("pool_size", -1),
        ("pool_recycle", True),
        ("pool_recycle", False),
        ("pool_recycle", 0),
        ("pool_recycle", -1),
    ],
)
def test_database_pool_settings_reject_booleans_and_non_positive_integers(field: str, invalid_value) -> None:
    with pytest.raises(ValidationError):
        DatabaseConfig(**{field: invalid_value})


@pytest.mark.asyncio
async def test_configured_command_timeout_ends_stalled_command(monkeypatch) -> None:
    config = DatabaseConfig(
        backend="postgres",
        postgres_url="postgresql://user:password@localhost/deerflow",
        command_timeout=0.01,
    )

    class _StalledAsyncpgEngine:
        def __init__(self, command_timeout: float) -> None:
            self.command_timeout = command_timeout

        async def checkout(self) -> None:
            async with asyncio.timeout(self.command_timeout):
                await asyncio.Event().wait()

        async def dispose(self) -> None:
            return None

    def _create_engine(_url: str, **kwargs) -> _StalledAsyncpgEngine:
        assert kwargs["pool_pre_ping"] is True
        return _StalledAsyncpgEngine(kwargs["connect_args"]["command_timeout"])

    bootstrap_schema = AsyncMock()

    # Patch one key: patch.dict restores all of sys.modules and races with background imports.
    monkeypatch.setitem(sys.modules, "asyncpg", ModuleType("asyncpg"))
    with (
        patch.object(engine_mod, "create_async_engine", side_effect=_create_engine),
        patch.object(engine_mod, "async_sessionmaker", return_value=MagicMock()),
        patch("deerflow.persistence.bootstrap.bootstrap_schema", new=bootstrap_schema),
    ):
        try:
            await engine_mod.init_engine_from_config(config)
            engine = engine_mod.get_engine()
            assert isinstance(engine, _StalledAsyncpgEngine)
            started_at = monotonic()
            with pytest.raises(TimeoutError):
                await engine.checkout()
            elapsed = monotonic() - started_at

            assert engine.command_timeout == config.command_timeout
            assert elapsed < 1
        finally:
            await engine_mod.close_engine()


@pytest.mark.asyncio
async def test_init_engine_from_config_preserves_longer_command_timeout_override(monkeypatch) -> None:
    config = DatabaseConfig(
        backend="postgres",
        postgres_url="postgresql://user:password@localhost/deerflow",
        pool_recycle=120,
        command_timeout=90,
    )
    mock_engine = MagicMock()
    mock_engine.dispose = AsyncMock()
    bootstrap_schema = AsyncMock()

    monkeypatch.setitem(sys.modules, "asyncpg", ModuleType("asyncpg"))
    with (
        patch.object(engine_mod, "create_async_engine", return_value=mock_engine) as create_engine,
        patch.object(engine_mod, "async_sessionmaker", return_value=MagicMock()),
        patch("deerflow.persistence.bootstrap.bootstrap_schema", new=bootstrap_schema),
    ):
        try:
            await engine_mod.init_engine_from_config(config)

            kwargs = create_engine.call_args.kwargs
            assert kwargs["connect_args"]["command_timeout"] == 90
            assert kwargs["pool_recycle"] == 120
        finally:
            await engine_mod.close_engine()


@pytest.mark.asyncio
async def test_init_engine_postgres_uses_hardened_kwargs(monkeypatch) -> None:
    url = "postgresql+asyncpg://user:password@localhost/deerflow"
    mock_engine = MagicMock()
    mock_engine.dispose = AsyncMock()
    bootstrap_schema = AsyncMock()

    monkeypatch.setitem(sys.modules, "asyncpg", ModuleType("asyncpg"))
    with (
        patch.object(engine_mod, "create_async_engine", return_value=mock_engine) as create_engine,
        patch.object(engine_mod, "async_sessionmaker", return_value=MagicMock()),
        patch("deerflow.persistence.bootstrap.bootstrap_schema", new=bootstrap_schema),
    ):
        try:
            await engine_mod.init_engine(backend="postgres", url=url, echo=True, pool_size=12)

            create_engine.assert_called_once_with(url, **engine_mod._postgres_engine_kwargs(echo=True, pool_size=12))
            bootstrap_schema.assert_awaited_once_with(mock_engine, backend="postgres", postgres_schema="")
        finally:
            await engine_mod.close_engine()


@pytest.mark.asyncio
async def test_init_engine_postgres_retry_uses_hardened_kwargs(monkeypatch) -> None:
    url = "postgresql+asyncpg://user:password@localhost/deerflow"
    initial_engine = MagicMock()
    initial_engine.dispose = AsyncMock()
    retry_engine = MagicMock()
    retry_engine.dispose = AsyncMock()
    bootstrap_schema = AsyncMock(side_effect=[_PostgresError("3D000", 'database "deerflow" does not exist'), None])
    auto_create = AsyncMock()

    monkeypatch.setitem(sys.modules, "asyncpg", ModuleType("asyncpg"))
    with (
        patch.object(engine_mod, "create_async_engine", side_effect=[initial_engine, retry_engine]) as create_engine,
        patch.object(engine_mod, "async_sessionmaker", return_value=MagicMock()),
        patch.object(engine_mod, "_auto_create_postgres_db", new=auto_create),
        patch("deerflow.persistence.bootstrap.bootstrap_schema", new=bootstrap_schema),
    ):
        try:
            await engine_mod.init_engine(backend="postgres", url=url, echo=False, pool_size=8)

            kwargs = engine_mod._postgres_engine_kwargs(echo=False, pool_size=8)
            assert create_engine.call_args_list == [call(url, **kwargs), call(url, **kwargs)]
            auto_create.assert_awaited_once_with(url)
            initial_engine.dispose.assert_awaited_once()
            assert bootstrap_schema.await_args_list == [
                call(initial_engine, backend="postgres", postgres_schema=""),
                call(retry_engine, backend="postgres", postgres_schema=""),
            ]
        finally:
            await engine_mod.close_engine()


def _wrapped_by_orig(missing: BaseException) -> BaseException:
    from sqlalchemy.exc import OperationalError

    return OperationalError("connect", {}, missing)


def _wrapped_by_cause(missing: BaseException) -> BaseException:
    wrapped = RuntimeError("bootstrap failed")
    wrapped.__cause__ = missing
    return wrapped


def _wrapped_by_context(missing: BaseException) -> BaseException:
    wrapped = RuntimeError("bootstrap failed")
    wrapped.__context__ = missing
    return wrapped


@pytest.mark.asyncio
@pytest.mark.parametrize("wrap", [_wrapped_by_orig, _wrapped_by_cause, _wrapped_by_context], ids=["orig", "cause", "context"])
async def test_init_engine_postgres_retries_a_wrapped_missing_database_error(monkeypatch, wrap) -> None:
    url = "postgresql+asyncpg://user:password@localhost/deerflow"
    engine = MagicMock()
    engine.dispose = AsyncMock()
    wrapped = wrap(_PostgresError("3D000", 'database "deerflow" does not exist'))
    bootstrap_schema = AsyncMock(side_effect=[wrapped, None])
    auto_create = AsyncMock()

    monkeypatch.setitem(sys.modules, "asyncpg", ModuleType("asyncpg"))
    with (
        patch.object(engine_mod, "create_async_engine", return_value=engine),
        patch.object(engine_mod, "async_sessionmaker", return_value=MagicMock()),
        patch.object(engine_mod, "_auto_create_postgres_db", new=auto_create),
        patch("deerflow.persistence.bootstrap.bootstrap_schema", new=bootstrap_schema),
    ):
        try:
            await engine_mod.init_engine(backend="postgres", url=url)

            auto_create.assert_awaited_once_with(url)
            assert bootstrap_schema.await_count == 2
        finally:
            await engine_mod.close_engine()


def _cyclic_chain() -> BaseException:
    # The classifier must terminate on a chain that loops back on itself.
    error = _PostgresError("42703", 'column "updated_at" does not exist')
    wrapper = RuntimeError("bootstrap failed")
    wrapper.__cause__ = error
    error.__context__ = wrapper
    return error


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "error",
    [
        pytest.param(_PostgresError("42703", 'column "updated_at" does not exist'), id="undefined-column"),
        pytest.param(_PostgresError("42P01", 'relation "runs" does not exist'), id="undefined-table"),
        pytest.param(_PostgresError("42883", "function gen_random_uuid() does not exist"), id="undefined-function"),
        pytest.param(_PostgresError("28000", 'role "deerflow" does not exist'), id="undefined-role"),
        pytest.param(Exception('database "deerflow" does not exist'), id="no-sqlstate"),
        pytest.param(_cyclic_chain(), id="cyclic-chain"),
    ],
)
async def test_init_engine_postgres_reraises_other_does_not_exist_errors(monkeypatch, error) -> None:
    # Postgres words all of these "does not exist". Only a missing database
    # may trigger CREATE DATABASE; anything else must surface unchanged
    # instead of as "already exists" or "permission denied to create database".
    url = "postgresql+asyncpg://user:password@localhost/deerflow"
    engine = MagicMock()
    engine.dispose = AsyncMock()
    bootstrap_schema = AsyncMock(side_effect=error)
    auto_create = AsyncMock()

    monkeypatch.setitem(sys.modules, "asyncpg", ModuleType("asyncpg"))
    with (
        patch.object(engine_mod, "create_async_engine", return_value=engine) as create_engine,
        patch.object(engine_mod, "async_sessionmaker", return_value=MagicMock()),
        patch.object(engine_mod, "_auto_create_postgres_db", new=auto_create),
        patch("deerflow.persistence.bootstrap.bootstrap_schema", new=bootstrap_schema),
    ):
        try:
            with pytest.raises(type(error)) as raised:
                await engine_mod.init_engine(backend="postgres", url=url)

            assert raised.value is error
            auto_create.assert_not_awaited()
            create_engine.assert_called_once()
            bootstrap_schema.assert_awaited_once()
        finally:
            await engine_mod.close_engine()


@pytest.mark.asyncio
async def test_init_engine_sqlite_omits_postgres_kwargs_and_keeps_wal_listener(tmp_path) -> None:
    url = f"sqlite+aiosqlite:///{tmp_path / 'deerflow.db'}"
    mock_engine = MagicMock()
    mock_engine.sync_engine = object()
    mock_engine.dispose = AsyncMock()
    bootstrap_schema = AsyncMock()
    registered: dict[str, object] = {}

    def _capture_listener(target, event_name):
        assert target is mock_engine.sync_engine
        assert event_name == "connect"

        def _decorator(fn):
            registered["listener"] = fn
            return fn

        return _decorator

    with (
        patch.object(engine_mod, "create_async_engine", return_value=mock_engine) as create_engine,
        patch.object(engine_mod, "async_sessionmaker", return_value=MagicMock()),
        patch("sqlalchemy.event.listens_for", new=_capture_listener),
        patch("deerflow.persistence.bootstrap.bootstrap_schema", new=bootstrap_schema),
    ):
        try:
            await engine_mod.init_engine(backend="sqlite", url=url, echo=True, sqlite_dir=str(tmp_path))

            create_engine.assert_called_once_with(url, echo=True, json_serializer=engine_mod._json_serializer)

            cursor = MagicMock()
            dbapi_connection = MagicMock()
            dbapi_connection.cursor.return_value = cursor
            listener = registered["listener"]
            listener(dbapi_connection, None)
            assert [entry.args[0] for entry in cursor.execute.call_args_list] == [
                "PRAGMA journal_mode=WAL;",
                "PRAGMA synchronous=NORMAL;",
                "PRAGMA foreign_keys=ON;",
                "PRAGMA busy_timeout=30000;",
            ]
            cursor.close.assert_called_once_with()
        finally:
            await engine_mod.close_engine()