feat: add middleware (#8)

Author: hadrien

README.md

@@ -2,29 +2,46 @@
 
 SqlAlchemy integration for FastAPI®
 
-
 ## Configuration
 
-* Configure environ variables:
+### Environment variables:
   The keys of interest in `os.environ` are prefixed with `sqlalchemy_`.
   Each matching key (after the prefix is stripped) is treated as though it were the
   corresponding keyword argument to [`sqlalchemy.create_engine`](https://docs.sqlalchemy.org/en/13/core/engines.html?highlight=create_engine#sqlalchemy.create_engine)  # noqa
   call.
 
   The only required key is `sqlalchemy_url`, which provides the database URL.
 
-* Setup the app:
-  ```python
-  import fastapi_sqla
-  from fastapi import FastAPI
-
-  app = FastAPI()
-  fastapi_sqla.setup(app)
-  ```
-* Adding a new entity class:
-  ```python
-  from fastapi_sqla import Base
-
-  class Entity(Base):
-      __tablename__ = "table-name-in-db"
-  ```
+### Setup the app:
+
+```python
+import fastapi_sqla
+from fastapi import FastAPI
+
+app = FastAPI()
+fastapi_sqla.setup(app)
+```
+
+### Adding a new entity class:
+
+```python
+from fastapi_sqla import Base
+
+
+class Entity(Base):
+    __tablename__ = "table-name-in-db"
+```
+
+### Getting an sqla orm session
+
+```python
+from fastapi import APIRouter, Depends
+from fastapi_sqla import Session, with_session
+
+router = APIRouter()
+
+
+@router.get("/example")
+def example(session: Session = Depends(with_session)):
+    return session.execute("SELECT now()").scalar()
+```

fastapi_sqla.py

@@ -1,12 +1,14 @@
 import os
+from contextlib import contextmanager
 
 import structlog
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
+from fastapi.concurrency import contextmanager_in_threadpool
 from sqlalchemy import engine_from_config
 from sqlalchemy.ext.declarative import DeferredReflection, declarative_base
-from sqlalchemy.orm.session import sessionmaker
+from sqlalchemy.orm.session import Session, sessionmaker
 
-__all__ = ["Base", "setup"]
+__all__ = ["Base", "setup", "with_session"]
 
 logger = structlog.get_logger(__name__)
 
@@ -15,6 +17,7 @@
 
 def setup(app: FastAPI):
     app.add_event_handler("startup", startup)
+    app.middleware("http")(add_session_to_request)
 
 
 def startup():
@@ -27,3 +30,76 @@ def startup():
 
 class Base(declarative_base(cls=DeferredReflection)):  # type: ignore
     __abstract__ = True
+
+
+@contextmanager
+def open_session() -> Session:
+    """Context manager that opens a session and properly closes session when exiting.
+
+    If no exception is raised before exiting context, session is committed when exiting
+    context. If an exception is raised, session is rollbacked.
+    """
+    session = _Session()
+    logger.bind(db_session=session)
+
+    try:
+        yield session
+        logger.debug("committing")
+        session.commit()
+    except Exception:
+        logger.exception("rolling back")
+        session.rollback()
+        raise
+    finally:
+        session.close()
+
+
+def with_session(request: Request) -> Session:
+    """Yield the sqlalchmey session for that request.
+
+    It is meant to be used as a FastAPI® dependency::
+
+        from er import sqla
+        from fastapi import APIRouter, Depends
+
+        router = APIRouter()
+
+        @router.get("/users")
+        def get_users(db: sqla.Session = Depends(sqla.with_session)):
+            pass
+    """
+    try:
+        yield request.scope["sqla_session"]
+    except KeyError:  # pragma: no cover
+        raise Exception(
+            "No session found in request, please ensure you've setup fastapi_sqla."
+        )
+
+
+async def add_session_to_request(request: Request, call_next):
+    """Middleware which injects a new sqla session into every request.
+
+    Handles creation of session, as well as commit, rollback, and closing of session.
+
+    Usage::
+
+        import fastapi_sqla
+        from fastapi import FastApi
+
+        app = FastApi()
+
+        fastapi_sqla.setup(app)  # includes middleware
+
+        @app.get("/users")
+        def get_users(session: sqla.Session = Depends(sqla.new_session)):
+            return session.query(...) # use your session here
+    """
+    async with contextmanager_in_threadpool(open_session()) as session:
+        request.scope["sqla_session"] = session
+        response = await call_next(request)
+        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.
+            session.rollback()
+
+    return response

tests/conftest.py

@@ -3,6 +3,7 @@
 from unittest.mock import patch
 
 from pytest import fixture
+from sqlalchemy.orm.session import close_all_sessions
 
 
 @fixture(scope="session")
@@ -22,5 +23,8 @@ def environ(db_uri):
 def tear_down():
     import fastapi_sqla
 
+    yield
+
+    close_all_sessions()
     # reload fastapi_sqla to clear sqla deferred reflection mapping stored in Base
     importlib.reload(fastapi_sqla)

tests/test_middleware.py

@@ -0,0 +1,104 @@
+from unittest.mock import patch
+
+import httpx
+from asgi_lifespan import LifespanManager
+from fastapi import Depends, FastAPI
+from pydantic import BaseModel
+from pytest import fixture, mark
+from sqlalchemy import engine_from_config
+from sqlalchemy.orm.session import close_all_sessions
+
+pytestmark = mark.asyncio
+
+
+@fixture
+def engine(environ):
+    engine = engine_from_config(environ, prefix="sqlalchemy_")
+    return engine
+
+
+@fixture(autouse=True)
+def setup_tear_down(engine):
+    engine.execute(
+        """
+        CREATE TABLE IF NOT EXISTS public.user (
+           id integer primary key,
+           first_name varchar,
+           last_name varchar
+        )
+    """
+    )
+    yield
+    close_all_sessions()
+    engine.execute("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 setup, with_session
+
+    app = FastAPI()
+    setup(app)
+
+    class UserIn(BaseModel):
+        id: int
+        first_name: str
+        last_name: str
+
+    @app.post("/users")
+    def create_user(user: UserIn, session=Depends(with_session)):
+        session.add(User(**user.dict()))
+        session.flush()
+        return {}
+
+    return app
+
+
+@fixture
+async def client(app):
+
+    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
+
+
+async def test_session_dependency(client):
+    res = await client.post(
+        "/users", json={"id": 1, "first_name": "Bob", "last_name": "Morane"}
+    )
+    assert res.status_code == 200
+
+
+@fixture
+def user_1(engine):
+    engine.execute("INSERT INTO public.user VALUES (1, 'bob', 'morane') ")
+    yield
+
+
+async def test_commit_error_returns_500(client, user_1):
+    res = await client.post(
+        "/users", json={"id": 1, "first_name": "Bob", "last_name": "Morane"}
+    )
+    assert res.status_code == 500
+
+
+async def test_rollback_on_http_exception(client):
+    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()

tests/test_startup.py

@@ -1,14 +1,7 @@
 import httpx
 from asgi_lifespan import LifespanManager
 from fastapi import FastAPI
-from pytest import fixture, mark
-from sqlalchemy.orm.session import close_all_sessions
-
-
-@fixture(autouse=True)
-def setup_tear_down():
-    yield
-    close_all_sessions()
+from pytest import mark
 
 
 def test_startup():