- fix typo

- autocommit and autorollback in run_in_new_ctx

Author: krylosov-aa

.python-version

@@ -5,7 +5,7 @@ lint:
 	flake8 .
 
 test:
-	pytest --cov context_async_sqlalchemy exmaples/fastapi_example/tests --cov-report=term-missing
+	pytest --cov context_async_sqlalchemy examples/fastapi_example/tests --cov-report=term-missing
 
 uv:
 	uv sync

Makefile

@@ -36,15 +36,15 @@ async def some_func() -> None:
 
 The repository includes an example integration with FastAPI,
 which describes numerous workflows.
-[FastAPI example](https://github.com/krylosov-aa/context-async-sqlalchemy/tree/main/exmaples/fastapi_example/routes)
+[FastAPI example](https://github.com/krylosov-aa/context-async-sqlalchemy/tree/main/examples/fastapi_example/routes)
 
 
 It also includes two types of test setups you can use in your projects.
 The library currently has 90% test coverage. The tests are in the
 examples, as we want to test not in the abstract but in the context of a real
 asynchronous web application.
 
-[FastAPI tests example](https://github.com/krylosov-aa/context-async-sqlalchemy/tree/main/exmaples/fastapi_example/tests)
+[FastAPI tests example](https://github.com/krylosov-aa/context-async-sqlalchemy/tree/main/examples/fastapi_example/tests)
 
 ### The most basic example
 
@@ -278,7 +278,7 @@ achieved through fast transaction rollback.
 You can see the capabilities in the examples:
 
 [Here are tests with a common transaction between the
-application and the tests.](https://github.com/krylosov-aa/context-async-sqlalchemy/blob/main/exmaples/fastapi_example/tests/transactional/__init__.py)
+application and the tests.](https://github.com/krylosov-aa/context-async-sqlalchemy/blob/main/examples/fastapi_example/tests/transactional/__init__.py)
 
 
-[And here's an example with different transactions.](https://github.com/krylosov-aa/context-async-sqlalchemy/blob/main/exmaples/fastapi_example/tests/non_transactional/__init__.py)
+[And here's an example with different transactions.](https://github.com/krylosov-aa/context-async-sqlalchemy/blob/main/examples/fastapi_example/tests/non_transactional/__init__.py)

README.md

@@ -5,7 +5,6 @@
     get_db_session_from_context,
     put_db_session_to_context,
     pop_db_session_from_context,
-    run_in_new_ctx,
 )
 from .connect import DBConnect
 from .session import (
@@ -23,9 +22,15 @@
     rollback_all_sessions,
     close_all_sessions,
 )
+from .run_in_new_context import run_in_new_ctx
+from .starlette_utils import (
+    add_starlette_http_db_session_middleware,
+    starlette_http_db_session_middleware,
+)
+
 from .fastapi_utils import (
-    fastapi_db_session_middleware,
-    add_fastapi_db_session_middleware,
+    fastapi_http_db_session_middleware,
+    add_fastapi_http_db_session_middleware,
 )
 
 __all__ = [
@@ -48,6 +53,8 @@
     "commit_all_sessions",
     "rollback_all_sessions",
     "close_all_sessions",
-    "fastapi_db_session_middleware",
-    "add_fastapi_db_session_middleware",
+    "add_starlette_http_db_session_middleware",
+    "starlette_http_db_session_middleware",
+    "fastapi_http_db_session_middleware",
+    "add_fastapi_http_db_session_middleware",
 ]

context_async_sqlalchemy/__init__.py

@@ -66,10 +66,10 @@ async def change_host(self, host: str) -> None:
 
     async def create_session(self) -> AsyncSession:
         """Creates a new session"""
-        maker = await self.get_session_maker()
+        maker = await self.session_maker()
         return maker()
 
-    async def get_session_maker(self) -> async_sessionmaker[AsyncSession]:
+    async def session_maker(self) -> async_sessionmaker[AsyncSession]:
         """Gets the session maker"""
         if self._before_create_session_handler:
             await self._before_create_session_handler(self)

context_async_sqlalchemy/connect.py

@@ -1,16 +1,18 @@
-from contextvars import ContextVar, copy_context, Token
-from typing import Any, Awaitable, Callable, Generator, TypeVar
+from contextvars import ContextVar, Token
+from typing import Any, Generator
 
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from .connect import DBConnect
 
 
-def init_db_session_ctx() -> Token[dict[str, AsyncSession] | None]:
+def init_db_session_ctx(
+    force: bool = False,
+) -> Token[dict[str, AsyncSession] | None]:
     """
     Initiates a context for storing sessions
     """
-    if is_context_initiated():
+    if not force and is_context_initiated():
         raise Exception("Context already initiated")
 
     return _init_db_session_ctx()
@@ -73,30 +75,6 @@ def sessions_stream() -> Generator[AsyncSession, Any, None]:
         yield session
 
 
-AsyncCallableResult = TypeVar("AsyncCallableResult")
-AsyncCallable = Callable[..., Awaitable[AsyncCallableResult]]
-
-
-async def run_in_new_ctx(
-    callable_func: AsyncCallable[AsyncCallableResult],
-    *args: Any,
-    **kwargs: Any,
-) -> AsyncCallableResult:
-    """
-    Runs a function in a new context with new sessions that have their
-        own connection.
-    The intended use is to run multiple database queries concurrently.
-
-    example of use:
-        await asyncio.gather(
-          run_in_new_ctx(your_function_with_db_session, ...),
-          run_in_new_ctx(your_function_with_db_session, ...),
-        )
-    """
-    new_ctx = copy_context()
-    return await new_ctx.run(_new_ctx_wrapper, callable_func, *args, **kwargs)
-
-
 _db_session_ctx: ContextVar[dict[str, AsyncSession] | None] = ContextVar(
     "db_session_ctx", default=None
 )
@@ -112,15 +90,3 @@ def _get_initiated_context() -> dict[str, AsyncSession]:
 def _init_db_session_ctx() -> Token[dict[str, AsyncSession] | None]:
     session_ctx: dict[str, AsyncSession] | None = {}
     return _db_session_ctx.set(session_ctx)
-
-
-async def _new_ctx_wrapper(
-    callable_func: AsyncCallable[AsyncCallableResult],
-    *args: Any,
-    **kwargs: Any,
-) -> AsyncCallableResult:
-    token = _init_db_session_ctx()
-    try:
-        return await callable_func(*args, **kwargs)
-    finally:
-        await reset_db_session_ctx(token)

context_async_sqlalchemy/context.py

@@ -1,9 +1,9 @@
 from .middleware import (
-    fastapi_db_session_middleware,
-    add_fastapi_db_session_middleware,
+    fastapi_http_db_session_middleware,
+    add_fastapi_http_db_session_middleware,
 )
 
 __all__ = [
-    "fastapi_db_session_middleware",
-    "add_fastapi_db_session_middleware",
+    "fastapi_http_db_session_middleware",
+    "add_fastapi_http_db_session_middleware",
 ]

context_async_sqlalchemy/fastapi_utils/__init__.py

@@ -1,27 +1,22 @@
-from fastapi import FastAPI, Request
+from fastapi import FastAPI
 from starlette.middleware.base import (  # type: ignore[attr-defined]
+    Request,
     Response,
     RequestResponseEndpoint,
-    BaseHTTPMiddleware,
 )
 
 from context_async_sqlalchemy import (
-    init_db_session_ctx,
-    is_context_initiated,
-    reset_db_session_ctx,
-    auto_commit_by_status_code,
-    rollback_all_sessions,
+    add_starlette_http_db_session_middleware,
+    starlette_http_db_session_middleware,
 )
 
 
-def add_fastapi_db_session_middleware(app: FastAPI) -> None:
+def add_fastapi_http_db_session_middleware(app: FastAPI) -> None:
     """Adds middleware to the application"""
-    app.add_middleware(
-        BaseHTTPMiddleware, dispatch=fastapi_db_session_middleware
-    )
+    add_starlette_http_db_session_middleware(app)
 
 
-async def fastapi_db_session_middleware(
+async def fastapi_http_db_session_middleware(
     request: Request, call_next: RequestResponseEndpoint
 ) -> Response:
     """
@@ -33,22 +28,4 @@ async def fastapi_db_session_middleware(
 
     But you can commit or rollback manually in the handler.
     """
-    # Tests have different session management rules
-    # so if the context variable is already set, we do nothing
-    if is_context_initiated():
-        return await call_next(request)
-
-    # We set the context here, meaning all child coroutines will receive the
-    # same context. And even if a child coroutine requests the
-    # session first, the dictionary itself is shared, and this coroutine will
-    # add the session to dictionary = shared context.
-    token = init_db_session_ctx()
-    try:
-        response = await call_next(request)
-        await auto_commit_by_status_code(response.status_code)
-        return response
-    except Exception:
-        await rollback_all_sessions()
-        raise
-    finally:
-        await reset_db_session_ctx(token)
+    return await starlette_http_db_session_middleware(request, call_next)

context_async_sqlalchemy/fastapi_utils/middleware.py

@@ -0,0 +1,52 @@
+from contextvars import copy_context
+from typing import Any, Awaitable, Callable, TypeVar
+
+from .context import init_db_session_ctx, reset_db_session_ctx
+from .auto_commit import commit_all_sessions, rollback_all_sessions
+
+
+AsyncCallableResult = TypeVar("AsyncCallableResult")
+AsyncCallable = Callable[..., Awaitable[AsyncCallableResult]]
+
+
+async def run_in_new_ctx(
+    callable_func: AsyncCallable[AsyncCallableResult],
+    *args: Any,
+    **kwargs: Any,
+) -> AsyncCallableResult:
+    """
+    Runs a function in a new context with new sessions that have their
+        own connection.
+
+    It will commit the transaction automatically if callable_func does not
+        raise exceptions. Otherwise, the transaction will be rolled back.
+
+    The intended use is to run multiple database queries concurrently.
+
+    example of use:
+        await asyncio.gather(
+            run_in_new_ctx(
+                your_function_with_db_session, some_arg, some_kwarg=123,
+            ),
+            run_in_new_ctx(your_function_with_db_session, ...),
+        )
+    """
+    new_ctx = copy_context()
+    return await new_ctx.run(_new_ctx_wrapper, callable_func, *args, **kwargs)
+
+
+async def _new_ctx_wrapper(
+    callable_func: AsyncCallable[AsyncCallableResult],
+    *args: Any,
+    **kwargs: Any,
+) -> AsyncCallableResult:
+    token = init_db_session_ctx(force=True)
+    try:
+        result = await callable_func(*args, **kwargs)
+        await commit_all_sessions()
+        return result
+    except Exception:
+        await rollback_all_sessions()
+        raise
+    finally:
+        await reset_db_session_ctx(token)

context_async_sqlalchemy/run_in_new_context.py

@@ -99,7 +99,7 @@ async def new_non_ctx_session(
         async with new_non_ctx_session(connect) as session:
             await session.execute(...)
     """
-    session_maker = await connect.get_session_maker()
+    session_maker = await connect.session_maker()
     async with session_maker() as session:
         yield session