mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
6dc28f5868
* fix(persistence): auto-create the Postgres database only when it is missing init_engine retried bootstrap through CREATE DATABASE whenever the error text contained "does not exist". Postgres words missing relations, columns, functions and roles the same way, so a bootstrap that failed inside an existing database (for example a migration reading a dropped column) went on to CREATE DATABASE. Startup then failed with DuplicateDatabaseError, or with "permission denied to create database" for a role without CREATEDB, and the original error appeared nowhere in the traceback. Recognize a missing database by SQLSTATE 3D000 (invalid_catalog_name) instead, following __cause__ for a wrapped copy, and re-raise every other bootstrap error unchanged. A live PostgreSQL test covers both the missing-database path and an existing database whose migration fails. * docs(changelog): reference #7410 for the Postgres auto-create fix * fix(persistence): walk orig, __cause__ and __context__ for the missing-database SQLSTATE Match _is_unique_violation's chain walk, with its cycle guard, so a wrapper that keeps the driver error only on `orig` or `__context__` still triggers auto-create. Drop the live test's English message assertion; the SQLSTATE check already proves the original error surfaced. --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
75 lines
2.4 KiB
Python
75 lines
2.4 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import sys
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from deerflow.persistence import bootstrap as bootstrap_mod
|
|
from deerflow.persistence import engine as engine_mod
|
|
|
|
|
|
class _BlockingDisposeEngine:
|
|
def __init__(self) -> None:
|
|
self.dispose_started = asyncio.Event()
|
|
self.allow_dispose = asyncio.Event()
|
|
self.dispose_finished = asyncio.Event()
|
|
|
|
async def dispose(self) -> None:
|
|
self.dispose_started.set()
|
|
await self.allow_dispose.wait()
|
|
self.dispose_finished.set()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_postgres_auto_create_retry_drains_provisional_engine_dispose_across_cancellation(monkeypatch) -> None:
|
|
engine = _BlockingDisposeEngine()
|
|
previous_engine = engine_mod._engine
|
|
previous_factory = engine_mod._session_factory
|
|
monkeypatch.setitem(sys.modules, "asyncpg", SimpleNamespace())
|
|
monkeypatch.setattr(engine_mod, "create_async_engine", lambda *args, **kwargs: engine)
|
|
monkeypatch.setattr(engine_mod, "async_sessionmaker", lambda *args, **kwargs: object())
|
|
|
|
class MissingDatabaseError(Exception):
|
|
sqlstate = "3D000"
|
|
|
|
async def fail_bootstrap(*args, **kwargs) -> None:
|
|
raise MissingDatabaseError('database "deerflow" does not exist')
|
|
|
|
async def fake_auto_create(_url: str) -> None:
|
|
return None
|
|
|
|
monkeypatch.setattr(bootstrap_mod, "bootstrap_schema", fail_bootstrap)
|
|
monkeypatch.setattr(engine_mod, "_auto_create_postgres_db", fake_auto_create)
|
|
|
|
task = asyncio.create_task(
|
|
engine_mod.init_engine(
|
|
"postgres",
|
|
url="postgresql+asyncpg://user:password@localhost/deerflow",
|
|
)
|
|
)
|
|
try:
|
|
await asyncio.wait_for(engine.dispose_started.wait(), timeout=2)
|
|
|
|
task.cancel()
|
|
for _ in range(5):
|
|
await asyncio.sleep(0)
|
|
task.cancel()
|
|
for _ in range(5):
|
|
await asyncio.sleep(0)
|
|
|
|
assert not task.done()
|
|
engine.allow_dispose.set()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
finally:
|
|
engine.allow_dispose.set()
|
|
if not task.done():
|
|
task.cancel()
|
|
await asyncio.gather(task, return_exceptions=True)
|
|
engine_mod._engine = previous_engine
|
|
engine_mod._session_factory = previous_factory
|
|
|
|
assert engine.dispose_finished.is_set()
|