diff --git a/backend/packages/harness/deerflow/persistence/engine.py b/backend/packages/harness/deerflow/persistence/engine.py index dd78051505..4c99ab392a 100644 --- a/backend/packages/harness/deerflow/persistence/engine.py +++ b/backend/packages/harness/deerflow/persistence/engine.py @@ -58,14 +58,19 @@ _session_factory: async_sessionmaker[AsyncSession] | None = None # SQLSTATE invalid_catalog_name: the database named in the URL does not exist. _MISSING_DATABASE_SQLSTATE = "3D000" +# SQLSTATE duplicate_database: the CREATE DATABASE issued by +# ``_auto_create_postgres_db`` lost a race against another Gateway booting the same +# target. The winner's catalog row is committed by the time the loser is refused, so +# the database exists and is usable. +_DUPLICATE_DATABASE_SQLSTATE = "42P04" -def _is_missing_database_error(exc: BaseException) -> bool: - """Return whether *exc* means the target PostgreSQL database is absent. +def _has_sqlstate(exc: BaseException, code: str) -> bool: + """Return whether *exc*, or anything it wraps, carries SQLSTATE *code*. Postgres words missing relations, columns and roles as "does not exist" too, so the message cannot tell them apart; only the SQLSTATE can. - asyncpg raises this one unwrapped from connect(); follow ``orig``, + asyncpg raises these unwrapped from connect(); follow ``orig``, ``__cause__`` and ``__context__`` (as ``_is_unique_violation`` in ``runtime/runs/manager.py`` does) so a wrapped copy is recognized as well. """ @@ -76,7 +81,7 @@ def _is_missing_database_error(exc: BaseException) -> bool: if id(current) in seen: continue seen.add(id(current)) - if getattr(current, "sqlstate", None) == _MISSING_DATABASE_SQLSTATE: + if getattr(current, "sqlstate", None) == code: return True for attr in ("orig", "__cause__", "__context__"): inner = getattr(current, attr, None) @@ -85,6 +90,16 @@ def _is_missing_database_error(exc: BaseException) -> bool: return False +def _is_missing_database_error(exc: BaseException) -> bool: + """Return whether *exc* means the target PostgreSQL database is absent.""" + return _has_sqlstate(exc, _MISSING_DATABASE_SQLSTATE) + + +def _is_duplicate_database_error(exc: BaseException) -> bool: + """Return whether *exc* means the target PostgreSQL database already exists.""" + return _has_sqlstate(exc, _DUPLICATE_DATABASE_SQLSTATE) + + async def _auto_create_postgres_db(url: str) -> None: """Connect to the ``postgres`` maintenance DB and CREATE DATABASE. @@ -105,9 +120,20 @@ async def _auto_create_postgres_db(url: str) -> None: maint_url = parsed.set(database="postgres") maint_engine = create_async_engine(maint_url, isolation_level="AUTOCOMMIT") try: - async with maint_engine.connect() as conn: - await conn.execute(text(f'CREATE DATABASE "{db_name}"')) - logger.info("Auto-created PostgreSQL database: %s", db_name) + try: + async with maint_engine.connect() as conn: + await conn.execute(text(f'CREATE DATABASE "{db_name}"')) + logger.info("Auto-created PostgreSQL database: %s", db_name) + except Exception as exc: + # Two Gateways booting the same missing database both reach this point and + # only one can win; for the loser the database exists and is usable, so the + # caller's retry should proceed instead of the boot failing. Everything else + # -- including "permission denied to create database" -- propagates + # untouched, because a CREATE DATABASE failure is not a reason to mask the + # bootstrap error that led here. + if not _is_duplicate_database_error(exc): + raise + logger.info("PostgreSQL database %s was created by a peer process; continuing", db_name) finally: # Drain: the maintenance engine is local to this helper, so an # abandoned dispose leaves its pool to the ``postgres`` database with diff --git a/backend/tests/test_persistence_engine_auto_create_race.py b/backend/tests/test_persistence_engine_auto_create_race.py new file mode 100644 index 0000000000..34ef1836df --- /dev/null +++ b/backend/tests/test_persistence_engine_auto_create_race.py @@ -0,0 +1,117 @@ +"""Two Gateways booting the same missing PostgreSQL database: only one can create it. + +``_auto_create_postgres_db`` runs when bootstrap fails with SQLSTATE 3D000. With two +processes starting against the same server, the loser's ``CREATE DATABASE`` is refused with +SQLSTATE 42P04 *after* the winner has committed the catalog row, so the target exists and is +usable — the boot must continue to the retry rather than fail on a race it effectively won. +""" + +from __future__ import annotations + +import logging + +import pytest + +from deerflow.persistence import engine as engine_mod + +_TARGET_URL = "postgresql+asyncpg://deerflow:secret@127.0.0.1:5432/deerflow_target" + + +class _SqlstateError(Exception): + """Stands in for the asyncpg exception, which the driver attaches ``sqlstate`` to.""" + + def __init__(self, sqlstate: str | None, message: str = "server said no") -> None: + super().__init__(message) + if sqlstate is not None: + self.sqlstate = sqlstate + + +class _FailingMaintenanceEngine: + """Fake throwaway ``postgres``-database engine whose CREATE DATABASE raises *error*.""" + + def __init__(self, error: BaseException) -> None: + self.error = error + self.statements: list[str] = [] + self.disposed = False + + def connect(self) -> _FailingMaintenanceEngine: + return self + + async def __aenter__(self) -> _FailingMaintenanceEngine: + return self + + async def __aexit__(self, *exc_info: object) -> None: + return None + + async def execute(self, statement: object) -> None: + self.statements.append(str(statement)) + raise self.error + + async def dispose(self) -> None: + self.disposed = True + + +@pytest.fixture +def maintenance(monkeypatch: pytest.MonkeyPatch): + installed: list[_FailingMaintenanceEngine] = [] + + def _install(error: BaseException) -> _FailingMaintenanceEngine: + engine = _FailingMaintenanceEngine(error) + installed.append(engine) + monkeypatch.setattr(engine_mod, "create_async_engine", lambda *args, **kwargs: engine) + return engine + + _install.installed = installed # type: ignore[attr-defined] + return _install + + +@pytest.mark.asyncio +async def test_duplicate_database_from_a_peer_continues_the_boot(maintenance, caplog) -> None: + engine = maintenance(_SqlstateError("42P04", 'database "deerflow_target" already exists')) + + with caplog.at_level(logging.INFO): + await engine_mod._auto_create_postgres_db(_TARGET_URL) + + assert 'CREATE DATABASE "deerflow_target"' in engine.statements[0] + assert "was created by a peer process" in caplog.text + assert engine.disposed, "the throwaway maintenance engine must still be disposed" + + +@pytest.mark.asyncio +async def test_wrapped_duplicate_database_is_recognised(maintenance) -> None: + """SQLSTATE lives on the driver exception, which reaches us wrapped in ``orig``.""" + wrapper = RuntimeError("bootstrap failed") + wrapper.__cause__ = _SqlstateError("42P04", 'database "deerflow_target" already exists') + engine = maintenance(wrapper) + + await engine_mod._auto_create_postgres_db(_TARGET_URL) + + assert engine.disposed + + +@pytest.mark.asyncio +async def test_permission_denied_still_propagates(maintenance) -> None: + """The other CREATE DATABASE failure mode is a real error and must not be swallowed.""" + engine = maintenance(_SqlstateError("42501", "permission denied to create database")) + + with pytest.raises(_SqlstateError): + await engine_mod._auto_create_postgres_db(_TARGET_URL) + + assert engine.disposed + + +@pytest.mark.asyncio +async def test_unrelated_failure_still_propagates(maintenance) -> None: + engine = maintenance(OSError("connection refused")) + + with pytest.raises(OSError): + await engine_mod._auto_create_postgres_db(_TARGET_URL) + + assert engine.disposed + + +def test_missing_and_duplicate_codes_are_distinct() -> None: + """Guards against collapsing the two SQLSTATEs into one constant.""" + assert engine_mod._MISSING_DATABASE_SQLSTATE != engine_mod._DUPLICATE_DATABASE_SQLSTATE + assert not engine_mod._is_missing_database_error(_SqlstateError("42P04")) + assert not engine_mod._is_duplicate_database_error(_SqlstateError("3D000"))