atomic_db_session and set_test_context

Author: krylosov-aa

context_async_sqlalchemy/session.py

@@ -1,5 +1,5 @@
 from contextlib import asynccontextmanager
-from typing import AsyncGenerator
+from typing import AsyncGenerator, Literal
 
 from sqlalchemy.ext.asyncio import AsyncSession
 
@@ -26,13 +26,31 @@ async def db_session(connect: DBConnect) -> AsyncSession:
     return session
 
 
+_current_transaction_choices = Literal[
+    "commit",
+    "rollback",
+    "append",
+    "raise",
+]
+
+
 @asynccontextmanager
 async def atomic_db_session(
     connect: DBConnect,
+    current_transaction: _current_transaction_choices = "commit",
 ) -> AsyncGenerator[AsyncSession, None]:
     """
-    Autocommit or autorollback in place to avoid waiting for the end of the
-        request.
+    A context manager that can be used to wrap another function which
+        uses a context session, making that call isolated within its
+        own transaction.
+
+    There are several options that define how the function will handle
+        an already open transaction.
+    current_transaction:
+        "commit" - commits the open transaction and starts a new one
+        "rollback" - rolls back the open transaction and starts a new one
+        "append" - continues using the current transaction and commits it
+        "raise" - raises an InvalidRequestError
 
     example of use:
         async with atomic_db_session(connect) as session
@@ -42,9 +60,22 @@ async def atomic_db_session(
     """
     session = await db_session(connect)
     if session.in_transaction():
-        await session.commit()
-    async with session.begin():
-        yield session
+        if current_transaction == "commit":
+            await session.commit()
+        if current_transaction == "rollback":
+            await session.rollback()
+
+    if current_transaction == "append":
+        try:
+            yield session
+        except Exception:
+            await session.rollback()
+            raise
+        else:
+            await session.commit()
+    else:
+        async with session.begin():
+            yield session
 
 
 async def commit_db_session(connect: DBConnect) -> None:

context_async_sqlalchemy/test_utils.py

@@ -6,6 +6,7 @@
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from context_async_sqlalchemy import (
+    commit_all_sessions,
     DBConnect,
     init_db_session_ctx,
     put_db_session_to_context,
@@ -27,19 +28,27 @@ async def rollback_session(
 
 
 @asynccontextmanager
-async def set_test_context() -> AsyncGenerator[None]:
+async def set_test_context(auto_close: bool = False) -> AsyncGenerator[None]:
     """
-    Opens a context similar to middleware, but doesn't commit or
-        rollback automatically. This task falls to the fixture in tests.
+    Opens a context similar to middleware.
+
+    Use auto_close=False if you’re using a test session and transaction
+        that you close manually elsewhere in your code.
+
+    Use auto_close=True if, for example, you want to call a
+        function in a test that uses a context bypassing the
+        middleware, and you want all sessions to be closed automatically.
     """
     token = init_db_session_ctx()
     try:
         yield
     finally:
+        if auto_close:
+            await commit_all_sessions()
         await reset_db_session_ctx(
             token,
-            # Don't close the session here, as you opened in fixture.
-            with_close=False,
+            # Don't close the session here if you opened in fixture.
+            with_close=auto_close,
         )
 
 

docs/api/index.html

@@ -298,10 +298,18 @@ <h3 id="atomic_db_session">atomic_db_session</h3>
 <pre><code class="language-python">@asynccontextmanager
 async def atomic_db_session(
     connect: DBConnect,
+    current_transaction: Literal[&quot;commit&quot;, &quot;rollback&quot;, &quot;append&quot;, &quot;raise&quot;] = &quot;commit&quot;,
 ) -&gt; AsyncGenerator[AsyncSession, None]:
 </code></pre>
 <p>A context manager that can be used to wrap another function which
 uses a context session, making that call isolated within its own transaction.</p>
+<p>There are several options that define how the function will handle
+an already open transaction.</p>
+<p>current_transaction:
+- <code>commit</code> - commits the open transaction and starts a new one
+- <code>rollback</code> - rolls back the open transaction and starts a new one
+- <code>append</code> - continues using the current transaction and commits it
+- <code>raise</code> - raises an InvalidRequestError</p>
 <hr />
 <h3 id="commit_db_session">commit_db_session</h3>
 <pre><code class="language-python">async def commit_db_session(connect: DBConnect) -&gt; None:
@@ -377,12 +385,17 @@ <h3 id="rollback_session">rollback_session</h3>
 <hr />
 <h3 id="set_test_context">set_test_context</h3>
 <pre><code class="language-python">@asynccontextmanager
-async def set_test_context() -&gt; AsyncGenerator[None]:
+async def set_test_context(auto_close: bool = False) -&gt; AsyncGenerator[None]:
 </code></pre>
 <p>A context manager that creates a new context in which you can place a
 dedicated test session.
 It’s intended for use in tests where the test and the application share
 a single transaction.</p>
+<p>Use <code>auto_close=False</code> if you’re using a test session and transaction
+that you close manually elsewhere in your code.</p>
+<p>Use <code>auto_close=True</code> if, for example, you want to call a
+function in a test that uses a context bypassing the
+middleware, and you want all sessions to be closed automatically.</p>
 <hr />
 <h3 id="put_savepoint_session_in_ctx">put_savepoint_session_in_ctx</h3>
 <pre><code class="language-python">async def put_savepoint_session_in_ctx(

docs/index.html

@@ -218,5 +218,5 @@ <h2 id="how-it-works">How it works</h2>
 
 <!--
 MkDocs version : 1.6.1
-Build Date UTC : 2025-11-25 21:27:44.719156+00:00
+Build Date UTC : 2025-11-28 22:35:06.618073+00:00
 -->

docs/search/search_index.json

@@ -195,11 +195,21 @@ calls return the same session.
 @asynccontextmanager
 async def atomic_db_session(
     connect: DBConnect,
+    current_transaction: Literal["commit", "rollback", "append", "raise"] = "commit",
 ) -> AsyncGenerator[AsyncSession, None]:
 ```
 A context manager that can be used to wrap another function which
 uses a context session, making that call isolated within its own transaction.
 
+There are several options that define how the function will handle
+an already open transaction.
+
+current_transaction:
+- `commit` - commits the open transaction and starts a new one
+- `rollback` - rolls back the open transaction and starts a new one
+- `append` - continues using the current transaction and commits it
+- `raise` - raises an InvalidRequestError
+
 ---
 
 ### commit_db_session
@@ -306,13 +316,20 @@ It’s intended for use in fixtures to execute SQL queries during tests.
 ### set_test_context
 ```python
 @asynccontextmanager
-async def set_test_context() -> AsyncGenerator[None]:
+async def set_test_context(auto_close: bool = False) -> AsyncGenerator[None]:
 ```
 A context manager that creates a new context in which you can place a
 dedicated test session.
 It’s intended for use in tests where the test and the application share
 a single transaction.
 
+Use `auto_close=False` if you’re using a test session and transaction
+that you close manually elsewhere in your code.
+
+Use `auto_close=True` if, for example, you want to call a
+function in a test that uses a context bypassing the
+middleware, and you want all sessions to be closed automatically.
+
 ---
 
 ### put_savepoint_session_in_ctx

docs/sitemap.xml.gz

@@ -12,14 +12,34 @@ async def handler_with_db_session_and_atomic_2() -> None:
     You want to reuse this function, but you need to commit immediately,
         rather than wait for the request to complete.
     """
-    # This is a new transaction in a new connection
+    # Open transaction
     await _insert_1()
 
-    # If you want to wrap an operation in a new
-    # transaction, the current transaction will be committed automatically.
+    # autocommit current -> start and end new
     async with atomic_db_session(connection):
         await _insert_1()
 
+    # Open transaction
+    await _insert_1()
+
+    # rollback current -> start and end new
+    async with atomic_db_session(connection, "rollback"):
+        await _insert_1()
+
+    # Open transaction
+    await _insert_1()
+
+    # use current transaction and commit
+    async with atomic_db_session(connection, "append"):
+        await _insert_1()
+
+    # Open transaction
+    await _insert_1()
+
+    # raise InvalidRequestError error
+    async with atomic_db_session(connection, "raise"):
+        await _insert_1()
+
 
 async def _insert_1() -> None:
     session = await db_session(connection)

docs_sources/docs/api.md

@@ -10,6 +10,7 @@
 
 from httpx import AsyncClient
 from sqlalchemy import exists, select
+from sqlalchemy.exc import InvalidRequestError
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from examples.fastapi_example.models import ExampleTable
@@ -67,16 +68,19 @@ async def test_example_with_db_session_and_atomic_2(
         such a handler should only be tested in non-transactional tests.
     """
     # Act
-    response = await client.post(
-        "/example_with_db_session_and_atomic_2",
-    )
+    try:
+        await client.post(
+            "/example_with_db_session_and_atomic_2",
+        )
+    except InvalidRequestError:
+        ...
+    else:
+        raise Exception()
 
     # Assert
-    assert response.status_code == HTTPStatus.OK
-
     result = await db_session_test.execute(select(ExampleTable))
     rows = result.scalars().all()
-    assert len(rows) == 2
+    assert len(rows) == 5
     for row in rows:
         assert row.text == "example_with_db_session_and_atomic"
 

examples/fastapi_example/routes/atomic_usage_2.py

@@ -1,6 +1,6 @@
 [project]
 name = "context-async-sqlalchemy"
-version = "2.0.4"
+version = "2.1.4"
 description = "A convenient way to configure and work with an async SQLAlchemy session through context in asynchronous applications"
 readme = "README.md"
 authors = [