fastapi-sqla 2.10.0__py3-none-any.whl → 3.0.0__py3-none-any.whl

fastapi_sqla/__init__.py

@@ -1,7 +1,7 @@
 from fastapi_sqla.base import setup
 from fastapi_sqla.models import Base, Collection, Item, Page
 from fastapi_sqla.pagination import Paginate, PaginateSignature, Pagination
-from fastapi_sqla.sqla import Session, open_session
+from fastapi_sqla.sqla import Session, SessionDependency, SqlaSession, open_session
 
 __all__ = [
     "Base",
@@ -12,20 +12,33 @@ __all__ = [
     "PaginateSignature",
     "Pagination",
     "Session",
+    "SessionDependency",
+    "SqlaSession",
     "open_session",
     "setup",
 ]
 
 
 try:
-    from fastapi_sqla.async_pagination import AsyncPaginate, AsyncPagination
-    from fastapi_sqla.async_sqla import AsyncSession
+    from fastapi_sqla.async_pagination import (
+        AsyncPaginate,
+        AsyncPaginateSignature,
+        AsyncPagination,
+    )
+    from fastapi_sqla.async_sqla import (
+        AsyncSession,
+        AsyncSessionDependency,
+        SqlaAsyncSession,
+    )
     from fastapi_sqla.async_sqla import open_session as open_async_session
 
     __all__ += [
         "AsyncPaginate",
+        "AsyncPaginateSignature",
         "AsyncPagination",
         "AsyncSession",
+        "AsyncSessionDependency",
+        "SqlaAsyncSession",
         "open_async_session",
     ]
     has_asyncio_support = True

fastapi_sqla/_pytest_plugin.py

@@ -6,6 +6,7 @@ from alembic import command
 from alembic.config import Config
 from pytest import fixture
 from sqlalchemy import create_engine, text
+from sqlalchemy.orm.session import sessionmaker
 
 try:
     import asyncpg  # noqa
@@ -52,8 +53,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,7 +97,7 @@ 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
@@ -100,17 +105,15 @@ def sqla_reflection(sqla_modules, sqla_connection, db_url):
 
 
 @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:
         with patch("fastapi_sqla.sqla.engine_from_config") as engine_from_config:
             engine_from_config.return_value = sqla_connection
-            _Session.configure(bind=sqla_connection)
             yield engine_from_config
 
 
@@ -121,18 +124,27 @@ def sqla_transaction(sqla_connection):
     transaction.rollback()
 
 
+@fixture
+def session_factory():
+    return sessionmaker()
+
+
 @fixture
 def session(
-    sqla_transaction, sqla_connection, sqla_reflection, patch_engine_from_config
+    session_factory,
+    sqla_connection,
+    sqla_transaction,
+    sqla_reflection,
+    patch_engine_from_config,
 ):
     """Sqla session to use when creating db fixtures.
 
     While it does not write any record in DB, the application will still be able to
     access any record committed with that session.
     """
-    import fastapi_sqla.sqla
-
-    yield fastapi_sqla.sqla._Session(bind=sqla_connection)
+    session = session_factory(bind=sqla_connection)
+    yield session
+    session.close()
 
 
 def format_async_async_sqlalchemy_url(url):
@@ -149,32 +161,33 @@ 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 patch_new_engine(async_sqlalchemy_url, async_sqla_connection, request):
+    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(request, async_sqla_connection, async_sqla_transaction):
         """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:
             with patch("fastapi_sqla.async_sqla.new_engine") as new_engine:
                 new_engine.return_value = async_sqla_connection
-                _AsyncSession.configure(
-                    bind=async_sqla_connection, expire_on_commit=False
-                )
                 yield new_engine
 
     @fixture
@@ -183,18 +196,20 @@ if asyncio_support:
 
         await async_sqla_connection.run_sync(lambda conn: Base.prepare(conn.engine))
 
+    @fixture
+    def async_session_factory():
+        from fastapi_sqla.async_sqla import SqlaAsyncSession
+
+        return sessionmaker(class_=SqlaAsyncSession)
+
     @fixture
     async def async_session(
-        async_sqla_connection, async_sqla_reflection, patch_new_engine
+        async_session_factory,
+        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)
+        session = async_session_factory(bind=async_sqla_connection)
         yield session
         await session.close()
-
-else:
-
-    @fixture
-    async def patch_new_engine():
-        pass

fastapi_sqla/async_pagination.py

@@ -1,28 +1,31 @@
 import math
 from collections.abc import Awaitable, Callable
-from typing import Iterator, Optional, Union, cast
+from typing import Annotated, Iterator, Optional, Union, cast
 
 from fastapi import Depends, Query
 from sqlalchemy.sql import Select, func, select
 
-from fastapi_sqla.async_sqla import AsyncSession
+from fastapi_sqla.async_sqla import AsyncSessionDependency, SqlaAsyncSession
 from fastapi_sqla.models import Page
+from fastapi_sqla.sqla import _DEFAULT_SESSION_KEY
 
 QueryCountDependency = Callable[..., Awaitable[int]]
-PaginateSignature = Callable[[Select, Optional[bool]], Awaitable[Page]]
-DefaultDependency = Callable[[AsyncSession, int, int], PaginateSignature]
-WithQueryCountDependency = Callable[[AsyncSession, int, int, int], PaginateSignature]
+AsyncPaginateSignature = Callable[[Select, Optional[bool]], Awaitable[Page]]
+DefaultDependency = Callable[[SqlaAsyncSession, int, int], AsyncPaginateSignature]
+WithQueryCountDependency = Callable[
+    [SqlaAsyncSession, int, int, int], AsyncPaginateSignature
+]
 PaginateDependency = Union[DefaultDependency, WithQueryCountDependency]
 
 
-async def default_query_count(session: AsyncSession, query: Select) -> int:
+async def default_query_count(session: SqlaAsyncSession, query: Select) -> int:
     result = await session.execute(select(func.count()).select_from(query.subquery()))
     return cast(int, result.scalar())
 
 
 async def paginate_query(
     query: Select,
-    session: AsyncSession,
+    session: SqlaAsyncSession,
     total_items: int,
     offset: int,
     limit: int,
@@ -48,15 +51,16 @@ async def paginate_query(
 
 
 def AsyncPagination(
+    session_key: str = _DEFAULT_SESSION_KEY,
     min_page_size: int = 10,
     max_page_size: int = 100,
     query_count: Union[QueryCountDependency, None] = None,
 ) -> PaginateDependency:
     def default_dependency(
-        session: AsyncSession = Depends(),
+        session: SqlaAsyncSession = Depends(AsyncSessionDependency(key=session_key)),
         offset: int = Query(0, ge=0),
         limit: int = Query(min_page_size, ge=1, le=max_page_size),
-    ) -> PaginateSignature:
+    ) -> AsyncPaginateSignature:
         async def paginate(query: Select, scalars=True) -> Page:
             total_items = await default_query_count(session, query)
             return await paginate_query(
@@ -66,11 +70,11 @@ def AsyncPagination(
         return paginate
 
     def with_query_count_dependency(
-        session: AsyncSession = Depends(),
+        session: SqlaAsyncSession = Depends(AsyncSessionDependency(key=session_key)),
         offset: int = Query(0, ge=0),
         limit: int = Query(min_page_size, ge=1, le=max_page_size),
         total_items: int = Depends(query_count),
-    ):
+    ) -> AsyncPaginateSignature:
         async def paginate(query: Select, scalars=True) -> Page:
             return await paginate_query(
                 query, session, total_items, offset, limit, scalars=scalars
@@ -84,4 +88,4 @@ def AsyncPagination(
         return default_dependency
 
 
-AsyncPaginate: PaginateDependency = AsyncPagination()
+AsyncPaginate = Annotated[AsyncPaginateSignature, Depends(AsyncPagination())]

fastapi_sqla/async_sqla.py

@@ -1,10 +1,10 @@
-import os
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
-from typing import cast
+from typing import Annotated
 
 import structlog
-from fastapi import Request
+from fastapi import Depends, Request
+from fastapi.responses import PlainTextResponse
 from sqlalchemy import text
 from sqlalchemy.ext.asyncio import AsyncEngine
 from sqlalchemy.ext.asyncio import AsyncSession as SqlaAsyncSession
@@ -12,93 +12,86 @@ from sqlalchemy.orm.session import sessionmaker
 
 from fastapi_sqla import aws_aurora_support, aws_rds_iam_support
 from fastapi_sqla.models import Base
-from fastapi_sqla.sqla import new_engine
+from fastapi_sqla.sqla import _DEFAULT_SESSION_KEY, new_engine
 
 logger = structlog.get_logger(__name__)
-_ASYNC_SESSION_KEY = "fastapi_sqla_async_session"
-_AsyncSession = sessionmaker(class_=SqlaAsyncSession)
 
+_ASYNC_REQUEST_SESSION_KEY = "fastapi_sqla_async_session"
+_async_session_factories: dict[str, sessionmaker] = {}
 
-def new_async_engine():
-    envvar_prefix = None
-    if "async_sqlalchemy_url" in os.environ:
-        envvar_prefix = "async_sqlalchemy_"
 
-    engine = new_engine(envvar_prefix=envvar_prefix)
+def new_async_engine(key: str = _DEFAULT_SESSION_KEY):
+    engine = new_engine(key)
     return AsyncEngine(engine)
 
 
-async def startup():
-    engine = new_async_engine()
+async def startup(key: str = _DEFAULT_SESSION_KEY):
+    engine = new_async_engine(key)
     aws_rds_iam_support.setup(engine.sync_engine)
     aws_aurora_support.setup(engine.sync_engine)
 
-    # Fail early:
+    # Fail early
     try:
         async with engine.connect() as connection:
             await connection.execute(text("select 'ok'"))
     except Exception:
         logger.critical(
-            "Failed querying db: is sqlalchemy_url or async_sqlalchemy_url envvar "
-            "correctly configured?"
+            f"Failed querying db for key '{key}': "
+            "are the the environment variables correctly configured for this key?"
         )
         raise
 
     async with engine.connect() as connection:
         await connection.run_sync(lambda conn: Base.prepare(conn.engine))
 
-    _AsyncSession.configure(bind=engine, expire_on_commit=False)
-    logger.info("startup", async_engine=engine)
+    _async_session_factories[key] = sessionmaker(
+        class_=SqlaAsyncSession, bind=engine, expire_on_commit=False
+    )
 
-
-class AsyncSession(SqlaAsyncSession):
-    def __new__(cls, request: Request):
-        """Yield the sqlalchmey async session for that request.
-
-        It is meant to be used as a FastAPI dependency::
-
-            from fastapi import APIRouter, Depends
-            from fastapi_sqla import AsyncSession
-
-            router = APIRouter()
-
-            @router.get("/users")
-            async def get_users(session: AsyncSession = Depends()):
-                pass
-        """
-        try:
-            return request.scope[_ASYNC_SESSION_KEY]
-        except KeyError:  # pragma: no cover
-            raise Exception(
-                "No async session found in request, please ensure you've setup "
-                "fastapi_sqla."
-            )
+    logger.info("engine startup", engine_key=key, async_engine=engine)
 
 
 @asynccontextmanager
-async def open_session() -> AsyncGenerator[SqlaAsyncSession, None]:
+async def open_session(
+    key: str = _DEFAULT_SESSION_KEY,
+) -> AsyncGenerator[SqlaAsyncSession, None]:
     """Context manager to open an async session and properly close it 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 = cast(SqlaAsyncSession, _AsyncSession())
+    try:
+        session: SqlaAsyncSession = _async_session_factories[key]()
+    except KeyError as exc:
+        raise KeyError(
+            f"No async session with key '{key}' found, "
+            "please ensure you've configured the environment variables for this key."
+        ) from exc
+
     logger.bind(db_async_session=session)
 
     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()
 
 
-async def add_session_to_request(request: Request, call_next):
+async def add_session_to_request(
+    request: Request, call_next, key: str = _DEFAULT_SESSION_KEY
+):
     """Middleware which injects a new sqla async session into every request.
 
     Handles creation of session, as well as commit, rollback, and closing of session.
@@ -113,15 +106,71 @@ async def add_session_to_request(request: Request, call_next):
         fastapi_sqla.setup(app)  # includes middleware
 
         @app.get("/users")
-        async def get_users(session: fastapi_sqla.AsyncSession = Depends()):
+        async def get_users(session: fastapi_sqla.AsyncSession):
             return await session.execute(...) # use your session here
     """
-    async with open_session() as session:
-        request.scope[_ASYNC_SESSION_KEY] = session
+    async with open_session(key) as session:
+        setattr(request.state, f"{_ASYNC_REQUEST_SESSION_KEY}_{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
+
+
+class AsyncSessionDependency:
+    def __init__(self, key: str = _DEFAULT_SESSION_KEY) -> None:
+        self.key = key
+
+    def __call__(self, request: Request) -> SqlaAsyncSession:
+        """Yield the sqlalchemy async session for that request.
+
+        It is meant to be used as a FastAPI dependency::
+
+            from fastapi import APIRouter, Depends
+            from fastapi_sqla import SqlaAsyncSession, AsyncSessionDependency
+
+            router = APIRouter()
+
+            @router.get("/users")
+            async def get_users(
+                session: SqlaAsyncSession = Depends(AsyncSessionDependency())
+            ):
+                pass
+        """
+        try:
+            return getattr(request.state, f"{_ASYNC_REQUEST_SESSION_KEY}_{self.key}")
+        except AttributeError:
+            logger.exception(
+                f"No async session with key '{self.key}' found in request, "
+                "please ensure you've setup fastapi_sqla.",
+                session_key=self.key,
+            )
+            raise
+
+
+default_async_session_dep = AsyncSessionDependency()
+AsyncSession = Annotated[SqlaAsyncSession, Depends(default_async_session_dep)]

fastapi_sqla/aws_rds_iam_support.py

@@ -1,4 +1,5 @@
 from os import environ
+from typing import Any
 
 try:
     import boto3
@@ -11,9 +12,10 @@ except ImportError as err:
 from functools import lru_cache
 
 from sqlalchemy import event
+from sqlalchemy.engine import Engine
 
 
-def setup(engine):
+def setup(engine: Engine):
     lc_environ = {k.lower(): v for k, v in environ.items()}
     aws_rds_iam_enabled = lc_environ.get("fastapi_sqla_aws_rds_iam_enabled") == "true"
 
@@ -30,13 +32,13 @@ def get_rds_client():
     return session.client("rds")
 
 
-def get_authentication_token(host, port, user):
+def get_authentication_token(host: str, port: int, user: str):
     client = get_rds_client()
     token = client.generate_db_auth_token(DBHostname=host, Port=port, DBUsername=user)
     return token
 
 
-def set_connection_token(dialect, conn_rec, cargs, cparams):
+def set_connection_token(dialect, conn_rec, cargs, cparams: dict[str, Any]):
     cparams["password"] = get_authentication_token(
         host=cparams["host"], port=cparams.get("port", 5432), user=cparams["user"]
     )

fastapi_sqla/base.py

@@ -1,4 +1,6 @@
+import functools
 import os
+import re
 
 from fastapi import FastAPI
 from sqlalchemy.engine import Engine
@@ -15,19 +17,41 @@ except ImportError as err:  # pragma: no cover
     asyncio_support_err = str(err)
 
 
-def setup(app: FastAPI):
-    engine = sqla.new_engine()
-
-    if not is_async_dialect(engine):
-        app.add_event_handler("startup", sqla.startup)
-        app.middleware("http")(sqla.add_session_to_request)
+_ENGINE_KEYS_REGEX = re.compile(r"fastapi_sqla__(?!_)(.+)(?<!_)__(?!_).+")
 
-    has_async_config = "async_sqlalchemy_url" in os.environ or is_async_dialect(engine)
-    if has_async_config:
-        assert has_asyncio_support, asyncio_support_err
-        app.add_event_handler("startup", async_sqla.startup)
-        app.middleware("http")(async_sqla.add_session_to_request)
 
-
-def is_async_dialect(engine: Engine):
+def setup(app: FastAPI):
+    engine_keys = _get_engine_keys()
+    engines = {key: sqla.new_engine(key) for key in engine_keys}
+    for key, engine in engines.items():
+        if not _is_async_dialect(engine):
+            app.add_event_handler("startup", functools.partial(sqla.startup, key=key))
+            app.middleware("http")(
+                functools.partial(sqla.add_session_to_request, key=key)
+            )
+        else:
+            assert has_asyncio_support, asyncio_support_err
+            app.add_event_handler(
+                "startup", functools.partial(async_sqla.startup, key=key)
+            )
+            app.middleware("http")(
+                functools.partial(async_sqla.add_session_to_request, key=key)
+            )
+
+
+def _get_engine_keys() -> set[str]:
+    keys = {sqla._DEFAULT_SESSION_KEY}
+
+    lowercase_environ = {k.lower(): v for k, v in os.environ.items()}
+    for env_var in lowercase_environ:
+        match = _ENGINE_KEYS_REGEX.search(env_var)
+        if not match:
+            continue
+
+        keys.add(match.group(1))
+
+    return keys
+
+
+def _is_async_dialect(engine: Engine):
     return engine.dialect.is_async if hasattr(engine.dialect, "is_async") else False

fastapi_sqla/pagination.py

@@ -1,24 +1,24 @@
 import math
 from collections.abc import Callable
 from functools import singledispatch
-from typing import Iterator, Optional, Union, cast
+from typing import Annotated, Iterator, Optional, Union, cast
 
 from fastapi import Depends, Query
 from sqlalchemy.orm import Query as LegacyQuery
 from sqlalchemy.sql import Select, func, select
 
 from fastapi_sqla.models import Page
-from fastapi_sqla.sqla import Session
+from fastapi_sqla.sqla import _DEFAULT_SESSION_KEY, SessionDependency, SqlaSession
 
 DbQuery = Union[LegacyQuery, Select]
 QueryCountDependency = Callable[..., int]
 PaginateSignature = Callable[[DbQuery, Optional[bool]], Page]
-DefaultDependency = Callable[[Session, int, int], PaginateSignature]
-WithQueryCountDependency = Callable[[Session, int, int, int], PaginateSignature]
+DefaultDependency = Callable[[SqlaSession, int, int], PaginateSignature]
+WithQueryCountDependency = Callable[[SqlaSession, int, int, int], PaginateSignature]
 PaginateDependency = Union[DefaultDependency, WithQueryCountDependency]
 
 
-def default_query_count(session: Session, query: DbQuery) -> int:
+def default_query_count(session: SqlaSession, query: DbQuery) -> int:
     """Default function used to count items returned by a query.
 
     It is slower than a manually written query could be: It runs the query in a
@@ -46,7 +46,7 @@ def default_query_count(session: Session, query: DbQuery) -> int:
 @singledispatch
 def paginate_query(
     query: DbQuery,
-    session: Session,
+    session: SqlaSession,
     total_items: int,
     offset: int,
     limit: int,
@@ -59,7 +59,7 @@ def paginate_query(
 @paginate_query.register
 def _paginate_legacy(
     query: LegacyQuery,
-    session: Session,
+    session: SqlaSession,
     total_items: int,
     offset: int,
     limit: int,
@@ -81,7 +81,7 @@ def _paginate_legacy(
 @paginate_query.register
 def _paginate(
     query: Select,
-    session: Session,
+    session: SqlaSession,
     total_items: int,
     offset: int,
     limit: int,
@@ -107,12 +107,13 @@ def _paginate(
 
 
 def Pagination(
+    session_key: str = _DEFAULT_SESSION_KEY,
     min_page_size: int = 10,
     max_page_size: int = 100,
     query_count: Union[QueryCountDependency, None] = None,
 ) -> PaginateDependency:
     def default_dependency(
-        session: Session = Depends(),
+        session: SqlaSession = Depends(SessionDependency(key=session_key)),
         offset: int = Query(0, ge=0),
         limit: int = Query(min_page_size, ge=1, le=max_page_size),
     ) -> PaginateSignature:
@@ -125,7 +126,7 @@ def Pagination(
         return paginate
 
     def with_query_count_dependency(
-        session: Session = Depends(),
+        session: SqlaSession = Depends(SessionDependency(key=session_key)),
         offset: int = Query(0, ge=0),
         limit: int = Query(min_page_size, ge=1, le=max_page_size),
         total_items: int = Depends(query_count),
@@ -143,4 +144,4 @@ def Pagination(
         return default_dependency
 
 
-Paginate: PaginateDependency = Pagination()
+Paginate = Annotated[PaginateSignature, Depends(Pagination())]

fastapi_sqla/sqla.py

@@ -2,10 +2,10 @@ import asyncio
 import os
 from collections.abc import Generator
 from contextlib import contextmanager
-from typing import Union
+from typing import Annotated
 
 import structlog
-from fastapi import Request
+from fastapi import Depends, Request
 from fastapi.concurrency import contextmanager_in_threadpool
 from fastapi.responses import PlainTextResponse
 from sqlalchemy import engine_from_config, text
@@ -18,75 +18,62 @@ from fastapi_sqla.models import Base
 
 logger = structlog.get_logger(__name__)
 
-_SESSION_KEY = "fastapi_sqla_session"
+_DEFAULT_SESSION_KEY = "default"
+_REQUEST_SESSION_KEY = "fastapi_sqla_session"
+_session_factories: dict[str, sessionmaker] = {}
 
-_Session = sessionmaker()
 
+def new_engine(key: str = _DEFAULT_SESSION_KEY) -> Engine:
+    envvar_prefix = "sqlalchemy_"
+    if key != _DEFAULT_SESSION_KEY:
+        envvar_prefix = f"fastapi_sqla__{key}__{envvar_prefix}"
 
-def new_engine(*, envvar_prefix: Union[str, None] = None) -> Engine:
-    envvar_prefix = envvar_prefix if envvar_prefix else "sqlalchemy_"
-    lowercase_environ = {
-        k.lower(): v for k, v in os.environ.items() if k.lower() != "sqlalchemy_warn_20"
-    }
+    lowercase_environ = {k.lower(): v for k, v in os.environ.items()}
+    lowercase_environ.pop(f"{envvar_prefix}warn_20", None)
     return engine_from_config(lowercase_environ, prefix=envvar_prefix)
 
 
-def startup():
-    engine = new_engine()
+def startup(key: str = _DEFAULT_SESSION_KEY):
+    engine = new_engine(key)
     aws_rds_iam_support.setup(engine.engine)
     aws_aurora_support.setup(engine.engine)
 
-    # Fail early:
+    # Fail early
     try:
         with engine.connect() as connection:
             connection.execute(text("select 'OK'"))
     except Exception:
         logger.critical(
-            "Fail querying db: is sqlalchemy_url envvar correctly configured?"
+            f"Failed querying db for key '{key}': "
+            "are the the environment variables correctly configured for this key?"
         )
         raise
 
     Base.prepare(engine)
-    _Session.configure(bind=engine)
-    logger.info("startup", engine=engine)
+    _session_factories[key] = sessionmaker(bind=engine)
 
-
-class Session(SqlaSession):
-    def __new__(cls, request: Request):
-        """Yield the sqlalchmey session for that request.
-
-        It is meant to be used as a FastAPI dependency::
-
-            from fastapi import APIRouter, Depends
-            from fastapi_sqla import Session
-
-            router = APIRouter()
-
-            @router.get("/users")
-            def get_users(session: Session = Depends()):
-                pass
-        """
-        try:
-            return request.scope[_SESSION_KEY]
-        except KeyError:  # pragma: no cover
-            raise Exception(
-                "No session found in request, please ensure you've setup fastapi_sqla."
-            )
+    logger.info("engine startup", engine_key=key, engine=engine)
 
 
 @contextmanager
-def open_session() -> Generator[SqlaSession, None, None]:
+def open_session(key: str = _DEFAULT_SESSION_KEY) -> Generator[SqlaSession, None, None]:
     """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()
+    try:
+        session: SqlaSession = _session_factories[key]()
+    except KeyError as exc:
+        raise KeyError(
+            f"No session with key '{key}' found, "
+            "please ensure you've configured the environment variables for this key."
+        ) from exc
+
     logger.bind(db_session=session)
 
     try:
         yield session
-
     except Exception:
         logger.warning("context failed, rolling back", exc_info=True)
         session.rollback()
@@ -104,7 +91,9 @@ def open_session() -> Generator[SqlaSession, None, None]:
         session.close()
 
 
-async def add_session_to_request(request: Request, call_next):
+async def add_session_to_request(
+    request: Request, call_next, key: str = _DEFAULT_SESSION_KEY
+):
     """Middleware which injects a new sqla session into every request.
 
     Handles creation of session, as well as commit, rollback, and closing of session.
@@ -119,11 +108,11 @@ async def add_session_to_request(request: Request, call_next):
         fastapi_sqla.setup(app)  # includes middleware
 
         @app.get("/users")
-        def get_users(session: fastapi_sqla.Session = Depends()):
+        def get_users(session: fastapi_sqla.Session):
             return session.execute(...) # use your session here
     """
-    async with contextmanager_in_threadpool(open_session()) as session:
-        request.scope[_SESSION_KEY] = session
+    async with contextmanager_in_threadpool(open_session(key)) as session:
+        setattr(request.state, f"{_REQUEST_SESSION_KEY}_{key}", session)
 
         response = await call_next(request)
 
@@ -155,3 +144,36 @@ async def add_session_to_request(request: Request, call_next):
             await loop.run_in_executor(None, session.rollback)
 
     return response
+
+
+class SessionDependency:
+    def __init__(self, key: str = _DEFAULT_SESSION_KEY) -> None:
+        self.key = key
+
+    def __call__(self, request: Request) -> SqlaSession:
+        """Yield the sqlalchemy session for that request.
+
+        It is meant to be used as a FastAPI dependency::
+
+            from fastapi import APIRouter, Depends
+            from fastapi_sqla import SqlaSession, SessionDependency
+
+            router = APIRouter()
+
+            @router.get("/users")
+            def get_users(session: SqlaSession = Depends(SessionDependency())):
+                pass
+        """
+        try:
+            return getattr(request.state, f"{_REQUEST_SESSION_KEY}_{self.key}")
+        except AttributeError:
+            logger.exception(
+                f"No session with key '{self.key}' found in request, "
+                "please ensure you've setup fastapi_sqla.",
+                session_key=self.key,
+            )
+            raise
+
+
+default_session_dep = SessionDependency()
+Session = Annotated[SqlaSession, Depends(default_session_dep)]

{fastapi_sqla-2.10.0.dist-info → fastapi_sqla-3.0.0.dist-info}/METADATA

@@ -1,13 +1,13 @@
 Metadata-Version: 2.1
 Name: fastapi-sqla
-Version: 2.10.0
+Version: 3.0.0
 Summary: SQLAlchemy extension for FastAPI with support for pagination, asyncio, and pytest, ready for production.
 Home-page: https://github.com/dialoguemd/fastapi-sqla
 License: MIT
 Keywords: FastAPI,SQLAlchemy,asyncio,pytest,alembic
 Author: Hadrien David
 Author-email: hadrien.david@dialogue.co
-Requires-Python: >=3.7,<4.0
+Requires-Python: >=3.9,<4.0
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Environment :: Web Environment
 Classifier: Framework :: AsyncIO
@@ -19,8 +19,6 @@ Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.7
-Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
@@ -43,7 +41,7 @@ Requires-Dist: asgi_lifespan (>=1.0.1,<2.0.0) ; extra == "tests"
 Requires-Dist: asyncpg (>=0.28.0,<0.29.0) ; extra == "asyncpg"
 Requires-Dist: black (>=22.8.0,<23.0.0) ; extra == "tests"
 Requires-Dist: boto3 (>=1.24.74,<2.0.0) ; extra == "aws-rds-iam"
-Requires-Dist: fastapi (>=0.61)
+Requires-Dist: fastapi (>=0.95.1)
 Requires-Dist: greenlet (>=1.1.3,<2.0.0) ; extra == "tests"
 Requires-Dist: httpx (>=0.23.0,<0.24.0) ; extra == "tests"
 Requires-Dist: isort (>=5.5.3,<6.0.0) ; extra == "tests"
@@ -90,7 +88,7 @@ unique `email`:
 
 ```python
 # main.py
-from fastapi import Depends, FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException
 from fastapi_sqla import Base, Item, Page, Paginate, Session, setup
 from pydantic import BaseModel, EmailStr
 from sqlalchemy import select
@@ -118,12 +116,12 @@ class UserModel(UserIn):
 
 
 @app.get("/users", response_model=Page[UserModel])
-def list_users(paginate: Paginate = Depends()):
+def list_users(paginate: Paginate):
     return paginate(select(User))
 
 
 @app.get("/users/{user_id}", response_model=Item[UserModel])
-def get_user(user_id: int, session: Session = Depends()):
+def get_user(user_id: int, session: Session):
     user = session.get(User, user_id)
     if user is None:
         raise HTTPException(404)
@@ -131,7 +129,7 @@ def get_user(user_id: int, session: Session = Depends()):
 
 
 @app.post("/users", response_model=Item[UserModel])
-def create_user(new_user: UserIn, session: Session = Depends()):
+def create_user(new_user: UserIn, session: Session):
     user = User(**new_user.model_dump())
     session.add(user)
     try:
@@ -173,6 +171,21 @@ The only required key is `sqlalchemy_url`, which provides the database URL, exam
 export sqlalchemy_url=postgresql://postgres@localhost
 ```
 
+### Multi-session support
+
+In order to configure multiple sessions for the application, 
+set the environment variables with this prefix format: `fastapi_sqla__MY_KEY__`.
+
+As with the default session, each matching key (after the prefix is stripped) 
+is treated as though it were the corresponding keyword argument to [`sqlalchemy.create_engine`]
+call.
+
+For example, to configure a session with the `read_only` key:
+
+```bash
+export fastapi_sqla__read_only__sqlalchemy_url=postgresql://postgres@localhost
+```
+
 ### `asyncio` support using [`asyncpg`]
 
 SQLAlchemy `>= 1.4` supports `asyncio`.
@@ -182,7 +195,7 @@ To enable `asyncio` support against a Postgres DB, install `asyncpg`:
 pip install asyncpg
 ```
 
-And define environment variable `sqlalchemy_url` with `postgres+asyncpg` scheme:
+And define the environment variable `sqlalchemy_url` with `postgres+asyncpg` scheme:
 
 ```bash
 export sqlalchemy_url=postgresql+asyncpg://postgres@localhost
@@ -216,23 +229,71 @@ class Entity(Base):
 
 Use [FastAPI dependency injection] to get a session as a parameter of a path operation
 function.
-SQLAlchemy session is committed before response is returned or rollbacked if any
+
+The SQLAlchemy session is committed before the response is returned or rollbacked if any
 exception occurred:
 
 ```python
+from fastapi import APIRouter
+from fastapi_sqla import AsyncSession, Session
+
+router = APIRouter()
+
+
+@router.get("/example")
+def example(session: Session):
+    return session.execute("SELECT now()").scalar()
+
+
+@router.get("/async_example")
+async def async_example(session: AsyncSession):
+    return await session.scalar("SELECT now()")
+```
+
+In order to get a session configured with a custom key:
+
+```python
+from typing import Annotated
+
 from fastapi import APIRouter, Depends
-from fastapi_sqla import Session, AsyncSession
+from fastapi_sqla import (
+    AsyncSessionDependency,
+    SessionDependency,
+    SqlaAsyncSession,
+    SqlaSession,
+)
 
 router = APIRouter()
 
 
+# Preferred
+
+ReadOnlySession = Annotated[SqlaSession, Depends(SessionDependency(key="read_only"))]
+AsyncReadOnlySession = Annotated[
+    SqlaAsyncSession, Depends(AsyncSessionDependency(key="read_only"))
+]
+
 @router.get("/example")
-def example(session: Session = Depends()):
+def example(session: ReadOnlySession):
     return session.execute("SELECT now()").scalar()
 
 
 @router.get("/async_example")
-async def async_example(session: AsyncSession = Depends()):
+async def async_example(session: AsyncReadOnlySession):
+    return await session.scalar("SELECT now()")
+
+
+# Alternative
+
+@router.get("/example/alt")
+def example_alt(session: SqlaSession = Depends(SessionDependency(key="read_only"))):
+    return session.execute("SELECT now()").scalar()
+
+
+@router.get("/async_example/alt")
+async def async_example_alt(
+    session: SqlaAsyncSession = Depends(AsyncSessionDependency(key="read_only")),
+):
     return await session.scalar("SELECT now()")
 ```
 
@@ -240,12 +301,12 @@ async def async_example(session: AsyncSession = Depends()):
 
 When needing a session outside of a path operation, like when using
 [FastAPI background tasks], use `fastapi_sqla.open_session` context manager.
-SQLAlchemy session is committed when exiting context or rollbacked if any exception
+The SQLAlchemy session is committed when exiting context or rollbacked if any exception
 occurred:
 
 ```python
 from fastapi import APIRouter, BackgroundTasks
-from fastapi_sqla import open_session, open_async_session
+from fastapi_sqla import open_async_session, open_session
 
 router = APIRouter()
 
@@ -260,16 +321,23 @@ def run_bg():
     with open_session() as session:
         session.execute("SELECT now()").scalar()
 
+def run_bg_with_key():
+    with open_session(key="read_only") as session:
+        session.execute("SELECT now()").scalar()
 
 async def run_async_bg():
     async with open_async_session() as session:
         await session.scalar("SELECT now()")
+
+async def run_async_bg_with_key():
+    async with open_async_session(key="read_only") as session:
+        await session.scalar("SELECT now()")
 ```
 
 ## Pagination
 
 ```python
-from fastapi import APIRouter, Depends
+from fastapi import APIRouter
 from fastapi_sqla import Base, Page, Paginate
 from pydantic import BaseModel
 from sqlalchemy import select
@@ -290,7 +358,7 @@ class UserModel(BaseModel):
 
 
 @router.get("/users", response_model=Page[UserModel])
-def all_users(paginate: Paginate = Depends()):
+def all_users(paginate: Paginate):
     return paginate(select(User))
 ```
 
@@ -327,7 +395,7 @@ To paginate a query which doesn't return [scalars], specify `scalars=False` when
 `paginate`:
 
 ```python
-from fastapi import APIRouter, Depends
+from fastapi import APIRouter
 from fastapi_sqla import Base, Page, Paginate
 from pydantic import BaseModel
 from sqlalchemy import func, select
@@ -352,7 +420,7 @@ class UserModel(BaseModel):
 
 
 @router.get("/users", response_model=Page[UserModel])
-def all_users(paginate: Paginate = Depends()):
+def all_users(paginate: Paginate):
     query = (
         select(User.id, User.name, func.count(Note.id).label("notes_count"))
         .join(Note)
@@ -388,15 +456,15 @@ class UserModel(BaseModel):
     name: str
 
 
-def query_count(session: Session = Depends()) -> int:
+def query_count(session: Session) -> int:
     return session.execute(select(func.count()).select_from(User)).scalar()
 
 
-Paginate = Pagination(min_page_size=5, max_page_size=500, query_count=query_count)
+CustomPaginate = Pagination(min_page_size=5, max_page_size=500, query_count=query_count)
 
 
 @router.get("/users", response_model=Page[UserModel])
-def all_users(paginate: Paginate = Depends()):
+def all_users(paginate: CustomPaginate = Depends()):
     return paginate(select(User))
 ```
 
@@ -405,7 +473,7 @@ def all_users(paginate: Paginate = Depends()):
 When using the asyncio support, use the `AsyncPaginate` dependency:
 
 ```python
-from fastapi import APIRouter, Depends
+from fastapi import APIRouter
 from fastapi_sqla import Base, Page, AsyncPaginate
 from pydantic import BaseModel
 from sqlalchemy import select
@@ -426,7 +494,7 @@ class UserModel(BaseModel):
 
 
 @router.get("/users", response_model=Page[UserModel])
-async def all_users(paginate: AsyncPaginate = Depends()):
+async def all_users(paginate: AsyncPaginate):
     return await paginate(select(User))
 ```
 
@@ -450,16 +518,85 @@ class UserModel(BaseModel):
     name: str
 
 
-async def query_count(session: AsyncSession = Depends()) -> int:
+async def query_count(session: AsyncSession) -> int:
     result = await session.execute(select(func.count()).select_from(User))
     return result.scalar()
 
 
-Paginate = AsyncPagination(min_page_size=5, max_page_size=500, query_count=query_count)
+CustomPaginate = AsyncPagination(min_page_size=5, max_page_size=500, query_count=query_count)
 
 
 @router.get("/users", response_model=Page[UserModel])
-def all_users(paginate: CustomPaginate = Depends()):
+async def all_users(paginate: CustomPaginate = Depends()):
+    return await paginate(select(User))
+```
+
+### Multi-session support
+
+Pagination supports multiple sessions as well. To paginate using a session 
+configured with a custom key:
+
+```python
+from typing import Annotated
+
+from fastapi import APIRouter, Depends
+from fastapi_sqla import (
+    AsyncPaginateSignature,
+    AsyncPagination,
+    Base,
+    Page,
+    PaginateSignature,
+    Pagination,
+)
+from pydantic import BaseModel
+from sqlalchemy import func, select
+
+router = APIRouter()
+
+
+class User(Base):
+    __tablename__ = "user"
+
+
+class UserModel(BaseModel):
+    id: int
+    name: str
+
+
+# Preferred
+
+ReadOnlyPaginate = Annotated[
+    PaginateSignature, Depends(Pagination(session_key="read_only"))
+]
+AsyncReadOnlyPaginate = Annotated[
+    AsyncPaginateSignature, Depends(AsyncPagination(session_key="read_only"))
+]
+
+@router.get("/users", response_model=Page[UserModel])
+def all_users(paginate: ReadOnlyPaginate):
+    return paginate(select(User))
+
+@router.get("/async_users", response_model=Page[UserModel])
+async def async_all_users(paginate: AsyncReadOnlyPaginate):
+    return await paginate(select(User))
+
+
+# Alternative
+
+@router.get("/users/alt", response_model=Page[UserModel])
+def all_users_alt(
+    paginate: PaginateSignature = Depends(
+        Pagination(session_key="read_only")
+    ),
+):
+    return paginate(select(User))
+
+@router.get("/async_users/alt", response_model=Page[UserModel])
+async def async_all_users_alt(
+    paginate: AsyncPaginateSignature = Depends(
+        AsyncPagination(session_key="read_only")
+    ),
+):
     return await paginate(select(User))
 ```
 

fastapi_sqla-3.0.0.dist-info/RECORD

@@ -0,0 +1,16 @@
+fastapi_sqla/__init__.py,sha256=4-MHiYF6TclnQpgs2Z1KDwYe2OPHLHGEu5oJzV3494s,1128
+fastapi_sqla/_pytest_plugin.py,sha256=I7wmRmJHaws-wOHDBo2N50RUkmKd-zutMMpymI9Tg1w,5769
+fastapi_sqla/async_pagination.py,sha256=UA3KxkTa48utBFDHeAkOFxTGUIvtNEWKm4uH-bqzQH4,3114
+fastapi_sqla/async_sqla.py,sha256=BERggjeIjBW9hPqVxE7ry3iQyjBFc9KLvDGwV_LOfK8,5919
+fastapi_sqla/aws_aurora_support.py,sha256=4dxLKOqDccgLwFqlz81L6f4HzrOXMZkY7Zuf4t_310U,838
+fastapi_sqla/aws_rds_iam_support.py,sha256=YSJNhrxmhGN-GVk9PLMTmQSWTKZBvuorKkhc_XaoL44,1189
+fastapi_sqla/base.py,sha256=0X7Gbt49rBHPiSFmNy5S2PT0dA4UBNnwrAesYSkaHBc,1606
+fastapi_sqla/models.py,sha256=QhnPCX-Gz5exAZfWyCRyYSaZ7SM_8QY0Eir6b4_4oI8,1432
+fastapi_sqla/pagination.py,sha256=NsI4ZeOkgbiNNDBjtqZL1rF6j1ya--jjXiyf0GlLaXU,4459
+fastapi_sqla/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+fastapi_sqla/sqla.py,sha256=ih4JBAQzn-lT86jK5ZOX3tewT1UOrw2NUJauOso0mNg,5939
+fastapi_sqla-3.0.0.dist-info/LICENSE,sha256=8G0-nWLqi3xRYRrtRlTE8n1mkYJcnCRoZGUhv6ZE29c,1064
+fastapi_sqla-3.0.0.dist-info/METADATA,sha256=8l-HvdC7q5crdm7NJrSMg_bhf43jg73gx_gO0RM1eEc,19809
+fastapi_sqla-3.0.0.dist-info/WHEEL,sha256=Zb28QaM1gQi8f4VCBhsUklF61CTlNYfs9YAZn-TOGFk,88
+fastapi_sqla-3.0.0.dist-info/entry_points.txt,sha256=haa0EueKcRo8-AlJTpHBMn08wMBiULNGA53nkvaDWj0,53
+fastapi_sqla-3.0.0.dist-info/RECORD,,

fastapi_sqla-2.10.0.dist-info/RECORD

@@ -1,16 +0,0 @@
-fastapi_sqla/__init__.py,sha256=QKvBNGd3irKqWp9PQMZ_QpqoqpTzl8jI7sa-zbmOdn8,824
-fastapi_sqla/_pytest_plugin.py,sha256=K_kJRped_7-DhOoQgZVG7znA7Bb5zaDgbdKgDe4rthc,5534
-fastapi_sqla/async_pagination.py,sha256=YUNF6ORaFTiUisLNCKz996ptisr355UlEB6ZPJ06W48,2801
-fastapi_sqla/async_sqla.py,sha256=vfWeLJcZj8L3Z8uKuIQOF1yOEsCqo8rUYdZjfuTSu4M,3967
-fastapi_sqla/aws_aurora_support.py,sha256=4dxLKOqDccgLwFqlz81L6f4HzrOXMZkY7Zuf4t_310U,838
-fastapi_sqla/aws_rds_iam_support.py,sha256=hGs8erX4e7zYCqwandiokw9LHx_jlW83_fJBIYAb7c4,1090
-fastapi_sqla/base.py,sha256=jiL_4n0r_HhyAxkPhu9rwMfh110y2GLrqWDNgDmJk7E,933
-fastapi_sqla/models.py,sha256=QhnPCX-Gz5exAZfWyCRyYSaZ7SM_8QY0Eir6b4_4oI8,1432
-fastapi_sqla/pagination.py,sha256=wIDVTUYogE7LD1Gzt7AmR2uQU8FM2t2iMm7FP4A7Fe0,4239
-fastapi_sqla/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-fastapi_sqla/sqla.py,sha256=RgR-4VPKxfJw6QFC8oISL4htBrCyQHYc2EhjL8m0H_c,4955
-fastapi_sqla-2.10.0.dist-info/LICENSE,sha256=8G0-nWLqi3xRYRrtRlTE8n1mkYJcnCRoZGUhv6ZE29c,1064
-fastapi_sqla-2.10.0.dist-info/METADATA,sha256=XE56F8YbFEkOHF00SOC9CE8UFqb0h5jC4hIpmMY-Ukc,16597
-fastapi_sqla-2.10.0.dist-info/WHEEL,sha256=Zb28QaM1gQi8f4VCBhsUklF61CTlNYfs9YAZn-TOGFk,88
-fastapi_sqla-2.10.0.dist-info/entry_points.txt,sha256=haa0EueKcRo8-AlJTpHBMn08wMBiULNGA53nkvaDWj0,53
-fastapi_sqla-2.10.0.dist-info/RECORD,,