153 lines
5.3 KiB
Python
153 lines
5.3 KiB
Python
"""Pool checkout/command deadlines survive ORM bootstrap and DB-create retry."""
|
|
|
|
import asyncio
|
|
from contextlib import asynccontextmanager
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
from pydantic import ValidationError
|
|
|
|
from deerflow.config.database_config import DatabaseConfig
|
|
from deerflow.persistence import engine as persistence
|
|
|
|
|
|
def test_backend_defaults_and_explicit_overrides_are_preserved():
|
|
assert DatabaseConfig(backend="postgres").pool_recycle == 300
|
|
assert DatabaseConfig(backend="postgres").command_timeout == 30
|
|
assert DatabaseConfig(backend="mysql").pool_recycle == 1800
|
|
assert DatabaseConfig(backend="gauss").pool_recycle == 1800
|
|
config = DatabaseConfig(backend="postgres", pool_recycle=120, command_timeout=None)
|
|
assert config.pool_recycle == 120
|
|
assert config.command_timeout is None
|
|
|
|
|
|
@pytest.mark.parametrize("value", [True, False, 0, -1, float("nan"), float("inf"), float("-inf")])
|
|
def test_invalid_command_timeouts_fail_config_validation(value):
|
|
with pytest.raises(ValidationError):
|
|
DatabaseConfig(command_timeout=value)
|
|
|
|
|
|
@pytest.mark.parametrize("retry", [False, True])
|
|
@pytest.mark.parametrize("timeout", [7.5, None])
|
|
def test_postgres_uses_same_deadline_on_initial_and_retry_engines(monkeypatch, retry, timeout):
|
|
async def run():
|
|
created = []
|
|
|
|
class Engine:
|
|
dialect = SimpleNamespace(name="postgresql")
|
|
|
|
def __init__(self, fails):
|
|
self.fails = fails
|
|
self.dispose = AsyncMock()
|
|
|
|
@asynccontextmanager
|
|
async def begin(self):
|
|
async def run_sync(fn):
|
|
if self.fails:
|
|
self.fails = False
|
|
raise RuntimeError('database "review" does not exist')
|
|
|
|
yield SimpleNamespace(run_sync=run_sync)
|
|
|
|
def create(url, **kwargs):
|
|
instance = Engine(retry and not created)
|
|
created.append((url, kwargs, instance))
|
|
return instance
|
|
|
|
for name in ("_engine", "_session_factory", "_engine_loop"):
|
|
monkeypatch.setattr(persistence, name, None)
|
|
monkeypatch.setattr(persistence, "create_async_engine", create)
|
|
create_database = AsyncMock()
|
|
monkeypatch.setattr(persistence, "_auto_create_postgres_db", create_database)
|
|
await persistence.init_engine_from_config(
|
|
DatabaseConfig(
|
|
backend="postgres",
|
|
postgres_url="postgresql://user:pass@localhost/review",
|
|
pool_size=11,
|
|
max_overflow=3,
|
|
pool_timeout=9,
|
|
pool_recycle=120,
|
|
command_timeout=timeout,
|
|
)
|
|
)
|
|
assert len(created) == (2 if retry else 1)
|
|
assert create_database.await_count == int(retry)
|
|
for url, kwargs, _ in created:
|
|
assert url.startswith("postgresql+asyncpg://")
|
|
assert kwargs["connect_args"] == {"command_timeout": timeout}
|
|
assert kwargs["pool_pre_ping"] is True
|
|
assert kwargs["pool_recycle"] == 120
|
|
assert kwargs["pool_size"] == 11
|
|
assert kwargs["max_overflow"] == 3
|
|
assert kwargs["pool_timeout"] == 9
|
|
if retry:
|
|
created[0][2].dispose.assert_awaited_once()
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"backend,url",
|
|
[
|
|
("mysql", "mysql+asyncmy://user:pass@localhost/review"),
|
|
("postgres", "postgresql+psycopg://user:pass@localhost/review"),
|
|
],
|
|
)
|
|
def test_asyncpg_arguments_are_not_sent_to_other_drivers(backend, url):
|
|
kwargs = persistence._pooled_engine_kwargs(
|
|
backend=backend,
|
|
url=url,
|
|
echo=False,
|
|
pool_size=5,
|
|
max_overflow=10,
|
|
pool_timeout=30,
|
|
pool_recycle=1800,
|
|
command_timeout=7.5,
|
|
)
|
|
assert "connect_args" not in kwargs
|
|
assert kwargs["pool_recycle"] == 1800
|
|
|
|
|
|
def test_gauss_asyncpg_receives_command_timeout():
|
|
kwargs = persistence._pooled_engine_kwargs(
|
|
backend="gauss",
|
|
url="opengauss+asyncpg://user:pass@localhost/review",
|
|
echo=False,
|
|
pool_size=5,
|
|
max_overflow=10,
|
|
pool_timeout=30,
|
|
pool_recycle=1800,
|
|
command_timeout=10,
|
|
)
|
|
assert kwargs["connect_args"] == {"command_timeout": 10}
|
|
|
|
|
|
def test_sqlite_initialization_never_receives_network_pool_arguments(monkeypatch, tmp_path):
|
|
async def run():
|
|
actual_create = persistence.create_async_engine
|
|
calls = []
|
|
|
|
def capture(url, **kwargs):
|
|
calls.append(kwargs)
|
|
return actual_create(url, **kwargs)
|
|
|
|
for name in ("_engine", "_session_factory", "_engine_loop"):
|
|
monkeypatch.setattr(persistence, name, None)
|
|
monkeypatch.setattr(persistence, "create_async_engine", capture)
|
|
try:
|
|
await persistence.init_engine(
|
|
"sqlite",
|
|
url=f"sqlite+aiosqlite:///{tmp_path / 'review.db'}",
|
|
sqlite_dir=str(tmp_path),
|
|
command_timeout=0.01,
|
|
)
|
|
assert len(calls) == 1
|
|
assert "connect_args" not in calls[0]
|
|
assert "pool_recycle" not in calls[0]
|
|
finally:
|
|
if persistence.get_engine() is not None:
|
|
await persistence.get_engine().dispose()
|
|
|
|
asyncio.run(run())
|