fix: async middleware is equivalent to sync middleware - DIA-61984 (#107)

Author: arththebird

fastapi_sqla/_pytest_plugin.py

@@ -52,8 +52,12 @@ def db_url(db_host, db_user):
 
 
 @fixture(scope="session")
-def sqla_connection(db_url):
-    engine = create_engine(db_url)
+def engine(db_url):
+    return create_engine(db_url)
+
+
+@fixture(scope="session")
+def sqla_connection(engine):
     with engine.connect() as connection:
         yield connection
 
@@ -92,19 +96,19 @@ def sqla_modules():
 
 
 @fixture
-def sqla_reflection(sqla_modules, sqla_connection, db_url):
+def sqla_reflection(sqla_modules, sqla_connection):
     import fastapi_sqla
 
     fastapi_sqla.Base.metadata.bind = sqla_connection
     fastapi_sqla.Base.prepare(sqla_connection.engine)
 
 
 @fixture
-def patch_engine_from_config(request, db_url, sqla_connection, sqla_transaction):
+def patch_engine_from_config(request, sqla_connection, sqla_transaction):
     """So that all DB operations are never written to db for real."""
     from fastapi_sqla.sqla import _Session
 
-    if "dont_patch_engines" in request.keywords:
+    if "dont_patch_engines" in request.keywords:  # pragma: no cover
         yield
 
     else:
@@ -149,24 +153,29 @@ def async_sqlalchemy_url(db_url):
     return format_async_async_sqlalchemy_url(db_url)
 
 
-if asyncio_support:
+if asyncio_support:  # noqa: C901
 
     @fixture
-    async def async_engine(async_sqlalchemy_url):
+    def async_engine(async_sqlalchemy_url):
         return create_async_engine(async_sqlalchemy_url)
 
     @fixture
     async def async_sqla_connection(async_engine, event_loop):
-        async with async_engine.begin() as connection:
+        async with async_engine.connect() as connection:
             yield connection
-            await connection.rollback()
+
+    @fixture
+    async def async_sqla_transaction(async_sqla_connection):
+        async with async_sqla_connection.begin() as transaction:
+            yield transaction
+            await transaction.rollback()
 
     @fixture
     async def patch_new_engine(async_sqlalchemy_url, async_sqla_connection, request):
         """So that all async DB operations are never written to db for real."""
         from fastapi_sqla.async_sqla import _AsyncSession
 
-        if "dont_patch_engines" in request.keywords:
+        if "dont_patch_engines" in request.keywords:  # pragma: no cover
             yield
 
         else:
@@ -185,16 +194,13 @@ async def async_sqla_reflection(sqla_modules, async_sqla_connection):
 
     @fixture
     async def async_session(
-        async_sqla_connection, async_sqla_reflection, patch_new_engine
+        async_sqla_connection,
+        async_sqla_transaction,
+        async_sqla_reflection,
+        patch_new_engine,
     ):
         from fastapi_sqla.async_sqla import _AsyncSession
 
         session = _AsyncSession(bind=async_sqla_connection)
         yield session
         await session.close()
-
-else:
-
-    @fixture
-    async def patch_new_engine():
-        pass

fastapi_sqla/async_sqla.py

@@ -5,6 +5,7 @@
 
 import structlog
 from fastapi import Request
+from fastapi.responses import PlainTextResponse
 from sqlalchemy import text
 from sqlalchemy.ext.asyncio import AsyncEngine
 from sqlalchemy.ext.asyncio import AsyncSession as SqlaAsyncSession
@@ -87,13 +88,20 @@ async def open_session() -> AsyncGenerator[SqlaAsyncSession, None]:
 
     try:
         yield session
-        await session.commit()
 
     except Exception:
-        logger.exception("commit failed, rolling back")
+        logger.warning("context failed, rolling back", exc_info=True)
         await session.rollback()
         raise
 
+    else:
+        try:
+            await session.commit()
+        except Exception:
+            logger.exception("commit failed, rolling back")
+            await session.rollback()
+            raise
+
     finally:
         await session.close()
 
@@ -119,9 +127,30 @@ async def get_users(session: fastapi_sqla.AsyncSession = Depends()):
     async with open_session() as session:
         request.scope[_ASYNC_SESSION_KEY] = session
         response = await call_next(request)
+
+        is_dirty = bool(session.dirty or session.deleted or session.new)
+
+        # try to commit after response, so that we can return a proper 500 response
+        # and not raise a true internal server error
+        if response.status_code < 400:
+            try:
+                await session.commit()
+            except Exception:
+                logger.exception("commit failed, returning http error")
+                response = PlainTextResponse(
+                    content="Internal Server Error", status_code=500
+                )
+
         if response.status_code >= 400:
             # If ever a route handler returns an http exception, we do not want the
             # session opened by current context manager to commit anything in db.
+            if is_dirty:
+                # optimistically only log if there were uncommitted changes
+                logger.warning(
+                    "http error, rolling back possibly uncommitted changes",
+                    status_code=response.status_code,
+                )
+            # since this is no-op if session is not dirty, we can always call it
             await session.rollback()
 
     return response

tests/conftest.py

@@ -68,14 +68,6 @@ def environ(db_url, sqla_version_tuple, async_sqlalchemy_url):
         yield values
 
 
-@fixture(scope="session")
-def engine(environ):
-    from sqlalchemy import engine_from_config
-
-    engine = engine_from_config(environ, prefix="sqlalchemy_")
-    return engine
-
-
 @fixture(autouse=True)
 def tear_down(environ):
     from sqlalchemy.orm.session import close_all_sessions
@@ -88,6 +80,7 @@ def tear_down(environ):
     # reload fastapi_sqla to clear sqla deferred reflection mapping stored in Base
     importlib.reload(fastapi_sqla.models)
     importlib.reload(fastapi_sqla.sqla)
+    importlib.reload(fastapi_sqla.async_sqla)
     importlib.reload(fastapi_sqla)
 
 

tests/middleware/conftest.py

@@ -0,0 +1,83 @@
+from unittest.mock import Mock
+
+import httpx
+from asgi_lifespan import LifespanManager
+from fastapi import Depends, FastAPI, HTTPException
+from pydantic import BaseModel
+from pytest import fixture
+from sqlalchemy import text
+
+
+@fixture(scope="module", autouse=True)
+def setup_tear_down(sqla_connection):
+    with sqla_connection.begin():
+        sqla_connection.execute(
+            text(
+                """
+                CREATE TABLE IF NOT EXISTS public.user (
+                    id integer primary key,
+                    first_name varchar,
+                    last_name varchar
+                )
+                """
+            )
+        )
+    yield
+    with sqla_connection.begin():
+        sqla_connection.execute(text("DROP TABLE public.user"))
+
+
+@fixture
+def User():
+    from fastapi_sqla import Base
+
+    class User(Base):
+        __tablename__ = "user"
+
+    return User
+
+
+@fixture
+def app(User):
+    from fastapi_sqla import Session, setup
+
+    app = FastAPI()
+    setup(app)
+
+    class UserIn(BaseModel):
+        id: int
+        first_name: str
+        last_name: str
+
+    @app.post("/users")
+    def create_user(user: UserIn, session: Session = Depends()):
+        session.add(User(**dict(user)))
+
+    @app.get("/404")
+    def get_users(session: Session = Depends(Session)):
+        raise HTTPException(status_code=404, detail="YOLO")
+
+    return app
+
+
+@fixture
+def mock_middleware(app: FastAPI):
+    mock_middleware = Mock()
+
+    @app.middleware("http")
+    async def a_middleware(request, call_next):
+        res = await call_next(request)
+        mock_middleware()
+        return res
+
+    return mock_middleware
+
+
+@fixture
+async def client(app, mock_middleware):
+    async with LifespanManager(app):
+        transport = httpx.ASGITransport(app=app, raise_app_exceptions=False)
+        async with httpx.AsyncClient(
+            transport=transport, base_url="http://example.local"
+        ) as client:
+            yield client

tests/middleware/test_async_middleware.py

@@ -0,0 +1,109 @@
+from unittest.mock import patch
+
+from fastapi import Depends, FastAPI, HTTPException
+from pydantic import BaseModel
+from pytest import fixture, mark
+from sqlalchemy import text
+from structlog.testing import capture_logs
+
+pytestmark = [mark.sqlalchemy("1.4"), mark.require_asyncpg]
+
+
+@fixture
+def app(User):
+    from fastapi_sqla import AsyncSession, setup
+
+    app = FastAPI()
+    setup(app)
+
+    class UserIn(BaseModel):
+        id: int
+        first_name: str
+        last_name: str
+
+    @app.post("/users")
+    def create_user(user: UserIn, session: AsyncSession = Depends()):
+        session.add(User(**dict(user)))
+
+    @app.get("/404")
+    def get_users(session: AsyncSession = Depends(AsyncSession)):
+        raise HTTPException(status_code=404, detail="YOLO")
+
+    return app
+
+
+async def test_async_session_dependency(client, faker, async_session):
+    userid = faker.unique.random_int()
+    first_name = faker.first_name()
+    last_name = faker.last_name()
+    res = await client.post(
+        "/users", json={"id": userid, "first_name": first_name, "last_name": last_name}
+    )
+    assert res.status_code == 200, res.json()
+    row = (
+        await async_session.execute(
+            text(f"select * from public.user where id = {userid}")
+        )
+    ).fetchone()
+    assert row == (userid, first_name, last_name)
+
+
+@fixture
+async def user_1(async_sqla_connection):
+    async with async_sqla_connection.begin():
+        await async_sqla_connection.execute(
+            text("INSERT INTO public.user VALUES (1, 'bob', 'morane') ")
+        )
+    yield
+    async with async_sqla_connection.begin():
+        await async_sqla_connection.execute(
+            text("DELETE FROM public.user WHERE id = 1")
+        )
+
+
+async def test_commit_error_returns_500(client, user_1, mock_middleware):
+    with capture_logs() as caplog:
+        res = await client.post(
+            "/users",
+            json={"id": 1, "first_name": "Bob", "last_name": "Morane"},
+            headers={"origin": "localhost"},
+        )
+
+    assert res.status_code == 500
+
+    assert {
+        "event": "commit failed, returning http error",
+        "exc_info": True,
+        "log_level": "error",
+    } in caplog
+
+    assert {
+        "event": "http error, rolling back possibly uncommitted changes",
+        "log_level": "warning",
+        "status_code": 500,
+    } in caplog
+
+    mock_middleware.assert_called_once()
+
+
+async def test_rollback_on_http_exception(client, mock_middleware):
+    with patch("fastapi_sqla.async_sqla.open_session") as open_session:
+        session = open_session.return_value.__aenter__.return_value
+
+        await client.get("/404")
+
+        session.rollback.assert_awaited_once_with()
+        mock_middleware.assert_called_once()
+
+
+async def test_rollback_on_http_exception_silent(client, mock_middleware):
+    with capture_logs() as caplog:
+        await client.get("/404")
+
+    mock_middleware.assert_called_once()
+
+    assert {
+        "event": "http error, rolling back possibly uncommitted changes",
+        "log_level": "warning",
+        "status_code": 404,
+    } not in caplog