diff --git a/.gitignore b/.gitignore index d816355..4ef1c6e 100644 --- a/.gitignore +++ b/.gitignore @@ -25,3 +25,5 @@ config.*.yaml venv .venv dist + +.claude diff --git a/tests/api/conftest.py b/tests/api/conftest.py index eef0b32..e76b0ac 100644 --- a/tests/api/conftest.py +++ b/tests/api/conftest.py @@ -31,9 +31,15 @@ async def app(db_backend_config: DatabaseConfig) -> AsyncGenerator[FastAPI, None # create tables in the test database async with app.state.db_engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) - yield app - async with app.state.db_engine.begin() as conn: - await conn.run_sync(Base.metadata.drop_all) + try: + yield app + + async with app.state.db_engine.begin() as conn: + await conn.run_sync(Base.metadata.drop_all) + finally: + # See tests/scheduler/conftest.py: the database is shared, so the + # engine has to be disposed even when the teardown above fails. + await app.state.db_engine.dispose() @pytest_asyncio.fixture diff --git a/tests/scheduler/conftest.py b/tests/scheduler/conftest.py index 12b83f2..f681d80 100644 --- a/tests/scheduler/conftest.py +++ b/tests/scheduler/conftest.py @@ -1,24 +1,29 @@ """Pytest fixture and configurations""" import pytest_asyncio -from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine +from sqlalchemy.ext.asyncio import async_sessionmaker -from warden.lib.db.database import Base, build_db_url +from warden.lib.db.database import Base, build_engine @pytest_asyncio.fixture(scope="function") async def db_engine(config_db): - engine = create_async_engine(build_db_url(config_db.database)) + engine = build_engine(config_db.database) async with engine.begin() as conn: # Create all tables once await conn.run_sync(Base.metadata.create_all) - yield engine - - async with engine.begin() as conn: - # Delete tables - await conn.run_sync(Base.metadata.drop_all) - await engine.dispose() + try: + yield engine + + async with engine.begin() as conn: + # Delete tables + await conn.run_sync(Base.metadata.drop_all) + finally: + # Always dispose: every backend shares one database across the test + # session, so an engine left open on a failing teardown leaks its + # connections - and their transactions - into every later test. + await engine.dispose() @pytest_asyncio.fixture(scope="function") diff --git a/warden/api/routes/dependencies/db.py b/warden/api/routes/dependencies/db.py index 6af1224..4ccf6b8 100644 --- a/warden/api/routes/dependencies/db.py +++ b/warden/api/routes/dependencies/db.py @@ -1,15 +1,15 @@ from typing import Annotated, AsyncGenerator from fastapi import Depends, FastAPI, Request -from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from warden.lib.config import DatabaseConfig -from warden.lib.db.database import build_db_url +from warden.lib.db.database import build_engine def init_db(app: FastAPI, db_config: DatabaseConfig): """Initialize the async engine and session factory with the given DB URL.""" - engine = create_async_engine(build_db_url(db_config), echo=db_config.echo) + engine = build_engine(db_config) # TODO: ensure isolation between concurrent requests session_factory = async_sessionmaker(bind=engine, expire_on_commit=False) diff --git a/warden/lib/db/database.py b/warden/lib/db/database.py index 496d60c..f33cd87 100644 --- a/warden/lib/db/database.py +++ b/warden/lib/db/database.py @@ -1,6 +1,7 @@ """Warden db utils""" from sqlalchemy.engine.url import URL +from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine from sqlalchemy.orm import declarative_base from warden.lib.config import DatabaseConfig @@ -35,3 +36,11 @@ def build_db_url(cfg: DatabaseConfig) -> str: ).render_as_string(hide_password=False) raise ValueError(f"Unsupported backend: {cfg.backend}") + + +def build_engine(cfg: DatabaseConfig) -> AsyncEngine: + """Build the async engine for `cfg`.""" + + engine = create_async_engine(build_db_url(cfg), echo=cfg.echo) + + return engine diff --git a/warden/scheduler/main.py b/warden/scheduler/main.py index c48a989..7b2f1f4 100644 --- a/warden/scheduler/main.py +++ b/warden/scheduler/main.py @@ -9,11 +9,10 @@ from sqlalchemy.ext.asyncio import ( AsyncEngine, async_sessionmaker, - create_async_engine, ) from warden.lib.config import Config -from warden.lib.db.database import build_db_url +from warden.lib.db.database import build_engine from warden.lib.models import Job from warden.scheduler.cancellation_worker import cancellation_worker from warden.scheduler.db import job_update_commiter @@ -144,7 +143,7 @@ async def main_async(conf: Config | None = None): conf = Config() logging.config.dictConfig(config=conf.logging) - engine = create_async_engine(build_db_url(conf.database), echo=conf.database.echo) + engine = build_engine(conf.database) loop = asyncio.get_running_loop() stop_event = asyncio.Event()