fix: other middlewares don't run when there's an exception in fastapi-sqla context manager (#16)

Author: vicrep

fastapi_sqla/__init__.py

@@ -8,6 +8,7 @@
 import structlog
 from fastapi import Depends, FastAPI, Query, Request
 from fastapi.concurrency import contextmanager_in_threadpool
+from fastapi.responses import PlainTextResponse
 from pydantic import BaseModel, Field
 from pydantic.generics import GenericModel
 from sqlalchemy import engine_from_config
@@ -106,13 +107,20 @@ def open_session() -> Session:
 
     try:
         yield session
-        session.commit()
 
     except Exception:
-        logger.exception("commit failed, rolling back")
+        logger.warning("context failed, rolling back", exc_info=True)
         session.rollback()
         raise
 
+    else:
+        try:
+            session.commit()
+        except Exception:
+            logger.exception("commit failed, rolling back")
+            session.rollback()
+            raise
+
     finally:
         session.close()
 
@@ -137,11 +145,27 @@ def get_users(session: fastapi_sqla.Session = Depends()):
     """
     async with contextmanager_in_threadpool(open_session()) as session:
         request.scope[_SESSION_KEY] = session
+
         response = await call_next(request)
+
+        loop = asyncio.get_running_loop()
+
+        # 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 loop.run_in_executor(None, 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.
-            loop = asyncio.get_running_loop()
+            logger.warning("http error, rolling back", status_code=response.status_code)
             await loop.run_in_executor(None, session.rollback)
 
     return response

tests/test_middleware.py

@@ -1,8 +1,8 @@
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
 import httpx
 from asgi_lifespan import LifespanManager
-from fastapi import Depends, FastAPI
+from fastapi import Depends, FastAPI, HTTPException
 from pydantic import BaseModel
 from pytest import fixture, mark
 from sqlalchemy import text
@@ -68,9 +68,26 @@ def create_user_with_async_session(
         ):
             session.add(User(**user.dict()))
 
+    @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):
 
@@ -120,24 +137,36 @@ def user_1(sqla_connection):
     yield
 
 
-async def test_commit_error_returns_500(client, user_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"}
+            "/users",
+            json={"id": 1, "first_name": "Bob", "last_name": "Morane"},
+            headers={"origin": "localhost"},
         )
 
     assert res.status_code == 500
+
     assert {
-        "event": "commit failed, rolling back",
-        "log_level": "error",
+        "event": "commit failed, returning http error",
         "exc_info": True,
+        "log_level": "error",
+    } in caplog
+
+    assert {
+        "event": "http error, rolling back",
+        "log_level": "warning",
+        "status_code": 500,
     } in caplog
 
+    mock_middleware.assert_called_once()
+
 
-async def test_rollback_on_http_exception(client):
+async def test_rollback_on_http_exception(client, mock_middleware):
     with patch("fastapi_sqla.open_session") as open_session:
         session = open_session.return_value.__enter__.return_value
 
         await client.get("/404")
 
         session.rollback.assert_called_once_with()
+        mock_middleware.assert_called_once()

tests/test_open_session.py

@@ -0,0 +1,81 @@
+from pytest import fixture, mark, raises
+from sqlalchemy import insert, select, text
+from sqlalchemy.exc import IntegrityError
+
+
+@fixture(autouse=True, scope="module")
+def module_setup_tear_down(engine, sqla_connection):
+    engine.execute(
+        "CREATE TABLE IF NOT EXISTS test_table (id integer primary key, value varchar)"
+    )
+    yield
+    engine.execute("DROP TABLE test_table")
+
+
+@fixture(autouse=True)
+def setup(sqla_connection):
+    from fastapi_sqla import _Session
+
+    _Session.configure(bind=sqla_connection)
+
+
+@fixture(scope="module")
+def TestTable(module_setup_tear_down):
+    from fastapi_sqla import Base, startup
+
+    class TestTable(Base):
+        __tablename__ = "test_table"
+
+    startup()
+
+    return TestTable
+
+
+@mark.sqlalchemy("1.4")
+def test_open_session():
+    from fastapi_sqla import open_session
+
+    with open_session() as session:
+        res = session.execute(select(text("'OK'"))).scalar()
+
+    assert res == "OK"
+
+
+@mark.sqlalchemy("1.4")
+def test_open_session_rollback_when_error_occurs_in_context(TestTable, session):
+    from fastapi_sqla import open_session
+
+    error = Exception("Error in context")
+
+    class Custom(Exception):
+        pass
+
+    with raises(Exception) as raise_info:
+        with open_session() as session:
+            session.execute(insert(TestTable).values(id=1, value="bobby drop tables"))
+            raise error
+
+    assert raise_info.value == error
+
+    res = session.execute(select(TestTable)).fetchall()
+    assert res == [], "insert has not been rolled back"
+
+
+@fixture
+def existing_record(TestTable, session):
+    id = 1
+    session.execute(insert(TestTable).values(id=id, value="bob morane was there."))
+    session.flush()
+    yield (1, "bob morane was there.")
+
+
+def test_open_session_re_raise_exception_when_commit_fails(
+    TestTable, existing_record, session
+):
+    from fastapi_sqla import open_session
+
+    with raises(Exception) as raise_info:
+        with open_session() as session:
+            session.add(TestTable(id=1, value="bobby already exists"))
+
+    assert isinstance(raise_info.value, IntegrityError)