diff --git a/README.md b/README.md index 6df054b..d48e728 100644 --- a/README.md +++ b/README.md @@ -99,6 +99,8 @@ results = ( - Direct `session` — use your own session, commit manually - Context manager — full lifecycle control with auto rollback - `with_session()` — override session per query +- `transaction()` — one transaction shared by commands across models +- `session_context()` — one shared session with manual commit, rollback, and flush - Works with Flask, FastAPI, and any framework **Async** @@ -296,6 +298,120 @@ item = qm.get(id=1) > Accessing lazy relationship attributes on detached objects raises `DetachedInstanceError`. > Use `select_related()` or `prefetch_related()` to load relationships upfront, or use a direct `session`. +##### Multi-model transaction + +Use `transaction()` when commands for different models must succeed or fail +together. Pass a `sessionmaker`; the context creates and closes the session itself. +Every query manager inside the block automatically uses that session, so individual +commands do not need a `session=` argument. + +```python +from sqlalchemy_query_manager import transaction + +with transaction(Session): + owner = Owner.query_manager.create( + first_name="John", + last_name="Doe", + ) + group = Group.query_manager.create( + name="Backend", + owner_id=owner.id, + ) + Item.query_manager.create( + name="First task", + group_id=group.id, + ) +``` + +The outer transaction commits on normal exit and rolls back on an exception. Manual +`commit()` and `rollback()` are intentionally unavailable because they would break +the atomic boundary. Explicit `flush()` remains available: + +```python +with transaction(Session) as control: + item = Item.query_manager.create(name="Needs an id") + control.flush() + print(item.id) +``` + +Nested transaction blocks reuse the same session and create a database SAVEPOINT. +The nested transaction does not need the session source again: + +```python +with transaction(Session): + Owner.query_manager.create(first_name="John", last_name="Doe") + + try: + with transaction(): + Group.query_manager.create(name="Rolled back") + raise ValueError("cancel group") + except ValueError: + pass + + Item.query_manager.create(name="Still committed") +``` + +When the nested block fails, only its SAVEPOINT is rolled back. If the exception +leaves the outer block, the complete transaction is rolled back. + +##### Manual session context + +Use `session_context()` for commit-as-you-go workflows. It creates one ambient session +for all models but leaves transaction boundaries under your control. + +```python +from sqlalchemy_query_manager import session_context + +with session_context(Session) as work: + owner = Owner.query_manager.create( + first_name="John", + last_name="Doe", + ) + Group.query_manager.create(name="Committed", owner_id=owner.id) + + work.flush() + work.commit() + + Item.query_manager.create(name="Discarded") + work.rollback() +``` + +After `commit()` or `rollback()`, the next database command starts a new transaction +in the same session. For a session created by the context, any transaction still +pending on exit is rolled back; the context never commits implicitly. + +Both APIs also accept an existing `Session` or a session context-manager factory. +An existing session is borrowed and is never committed, rolled back, or closed by +`session_context()` on exit; its lifecycle remains the caller's responsibility. + +##### Async contexts + +The same API works with an async sessionmaker. Context and control methods become +awaitable, while all models still share one `AsyncSession`: + +```python +async with transaction(AsyncSessionMaker) as control: + owner = await Owner.query_manager.create( + first_name="John", + last_name="Doe", + ) + await Group.query_manager.create(name="Backend", owner_id=owner.id) + await control.flush() + +async with session_context(AsyncSessionMaker) as work: + await Item.query_manager.create(name="Committed manually") + await work.commit() +``` + +An `AsyncSession` cannot be shared by concurrent tasks. Do not run query-manager +commands from one context through `asyncio.gather()` or child tasks. Open a separate +context with its own sessionmaker inside each concurrent task instead. + +A synchronous `Session` likewise cannot be shared across threads. Each worker thread +must open its own `transaction(SessionMaker)` or `session_context(SessionMaker)`. +If a `ContextVar` is explicitly propagated to another thread, the package detects +the ownership mismatch and raises instead of reusing the ambient session. + --- #### Filtering diff --git a/pyproject.toml b/pyproject.toml index a2a87af..beb2237 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "sqlalchemy-query-manager" -version = "0.6.1" +version = "0.7.0" description = "Django-style ORM interface for SQLAlchemy: Q filters, eager loading, bulk operations, and native async support." authors = ["ViAchKoN"] readme = "README.md" diff --git a/sqlalchemy_query_manager/__init__.py b/sqlalchemy_query_manager/__init__.py index aef1752..5cda984 100644 --- a/sqlalchemy_query_manager/__init__.py +++ b/sqlalchemy_query_manager/__init__.py @@ -1,6 +1,7 @@ """Public API for SQLAlchemy Query Manager.""" from sqlalchemy_query_manager.core.async_query_manager import AsyncQueryManager +from sqlalchemy_query_manager.core.contexts import session_context, transaction from sqlalchemy_query_manager.core.exceptions import ( DoesNotExist, MultipleObjectsReturned, @@ -28,5 +29,7 @@ "MultipleObjectsReturned", "Q", "QueryManager", + "session_context", "Sum", + "transaction", ] diff --git a/sqlalchemy_query_manager/core/async_query_manager.py b/sqlalchemy_query_manager/core/async_query_manager.py index 271cb22..a35276f 100644 --- a/sqlalchemy_query_manager/core/async_query_manager.py +++ b/sqlalchemy_query_manager/core/async_query_manager.py @@ -3,12 +3,15 @@ from sqlalchemy import delete, func, inspect, select, text, update from sqlalchemy.orm import sessionmaker +from sqlalchemy_query_manager.core.contexts import get_current_context from sqlalchemy_query_manager.core.helpers import AggregateFunc from sqlalchemy_query_manager.core.sync_query_manager import QueryManager from sqlalchemy_query_manager.core.utils import get_async_session class AsyncQueryManager(QueryManager): + def _should_commit(self): + return isinstance(self.session, sessionmaker) and get_current_context() is None @get_async_session async def first(self, session=None): @@ -104,7 +107,7 @@ async def create(self, session=None, **kwargs): new_obj = self.ConverterConfig.model(**kwargs) session.add(new_obj) - if isinstance(self.session, sessionmaker): + if self._should_commit(): await session.commit() else: await session.flush() @@ -125,7 +128,7 @@ async def bulk_create( objects = [self.ConverterConfig.model(**item) for item in data] session.add_all(objects) - if isinstance(self.session, sessionmaker): + if self._should_commit(): await session.commit() else: await session.flush() @@ -193,7 +196,7 @@ async def update(self, session=None, expunge=True, **kwargs): ) returned_pks = [row[0] for row in result] - if isinstance(self.session, sessionmaker): + if self._should_commit(): await session.commit() else: await session.flush() @@ -217,7 +220,7 @@ async def _update_legacy_path(self, session, expunge=True, **kwargs): await session.execute(update_query) - if isinstance(self.session, sessionmaker): + if self._should_commit(): await session.commit() else: await session.flush() @@ -256,7 +259,7 @@ async def update_raw(self, session=None, **kwargs): result = await session.execute(update_query) - if isinstance(self.session, sessionmaker): + if self._should_commit(): await session.commit() else: await session.flush() @@ -279,7 +282,7 @@ async def update_or_create(self, session=None, defaults=None, **kwargs): if hasattr(existing, key): setattr(existing, key, value) - if isinstance(self.session, sessionmaker): + if self._should_commit(): await session.commit() else: await session.flush() @@ -349,7 +352,7 @@ async def delete(self, session=None, synchronize_session=True): delete_query, execution_options={"synchronize_session": False} ) - if isinstance(self.session, sessionmaker): + if self._should_commit(): await session.commit() else: await session.flush() diff --git a/sqlalchemy_query_manager/core/contexts.py b/sqlalchemy_query_manager/core/contexts.py new file mode 100644 index 0000000..4f5c65e --- /dev/null +++ b/sqlalchemy_query_manager/core/contexts.py @@ -0,0 +1,353 @@ +import asyncio +import contextvars +import dataclasses +import sys +import threading +import typing + +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.orm import Session + +from sqlalchemy_query_manager.core.transaction_context_manager import ( + AsyncTransactionSessionContextManager, + TransactionSessionContextManager, +) + + +@dataclasses.dataclass +class _ContextState: + session: typing.Union[Session, AsyncSession] + source: typing.Any + is_async: bool + owner_task: typing.Any = None + owner_thread_id: typing.Optional[int] = None + transaction_depth: int = 0 + + +_current_context = contextvars.ContextVar( + "sqlalchemy_query_manager_context", + default=None, +) + + +def _get_asyncio_task(): + try: + return asyncio.current_task() + except RuntimeError: + return None + + +def _validate_context_owner(state): + if state.is_async: + current_task = _get_asyncio_task() + if current_task is not state.owner_task: + raise RuntimeError( + "An ambient AsyncSession cannot be shared across concurrent tasks. " + "Open a separate context inside each task." + ) + elif threading.get_ident() != state.owner_thread_id: + raise RuntimeError( + "An ambient Session cannot be shared across threads. " + "Open a separate context inside each thread." + ) + + +def get_current_context(): + state = _current_context.get() + if state is not None: + _validate_context_owner(state) + return state + + +def _get_context_for_entry(source, is_async): + state = _current_context.get() + if state is None or state.is_async != is_async: + return state + + if is_async: + is_owner = _get_asyncio_task() is state.owner_task + worker = "task" + session_type = "AsyncSession" + else: + is_owner = threading.get_ident() == state.owner_thread_id + worker = "thread" + session_type = "Session" + + if is_owner: + return state + + if source is None or source is state.session: + raise RuntimeError( + "A child {worker} cannot inherit the ambient {session_type}. Pass a " + "session factory or context-manager provider to open an independent " + "context inside the child {worker}.".format( + worker=worker, + session_type=session_type, + ) + ) + return None + + +def _validate_nested_source(state, source): + if source is None or source is state.source or source is state.session: + return + raise ValueError( + "A nested context must use the ambient session. Omit the session source " + "or pass the same source used by the outer context." + ) + + +class SessionContextHandle: + def __init__(self, state): + self._state = state + self._active = True + + def _ensure_active(self): + if not self._active: + raise RuntimeError("The session context is no longer active.") + _validate_context_owner(self._state) + + def flush(self): + self._ensure_active() + return self._state.session.flush() + + def commit(self): + self._ensure_active() + if self._state.transaction_depth: + raise RuntimeError("Manual commit is not allowed inside transaction().") + return self._state.session.commit() + + def rollback(self): + self._ensure_active() + if self._state.transaction_depth: + raise RuntimeError("Manual rollback is not allowed inside transaction().") + return self._state.session.rollback() + + +class TransactionHandle(SessionContextHandle): + def commit(self): + raise RuntimeError( + "Manual commit is not allowed inside transaction(). " + "Use session_context() for commit-as-you-go workflows." + ) + + def rollback(self): + raise RuntimeError( + "Manual rollback is not allowed inside transaction(). " + "Use a nested transaction() savepoint or session_context()." + ) + + +class _SessionContext: + def __init__(self, source): + self.source = source + self._state = None + self._token = None + self._manager = None + self._handle = None + self._is_outermost = False + + def __enter__(self): + current = _get_context_for_entry(self.source, is_async=False) + if current is not None: + if current.is_async: + raise TypeError("Use 'async with' inside an async session context.") + _validate_nested_source(current, self.source) + self._state = current + else: + if self.source is None: + raise ValueError( + "An outer session_context() requires a session source." + ) + self._manager = TransactionSessionContextManager( + session=self.source, + ) + session = self._manager.__enter__() + self._state = _ContextState( + session=session, + source=self.source, + is_async=False, + owner_thread_id=threading.get_ident(), + ) + self._token = _current_context.set(self._state) + self._is_outermost = True + + self._handle = SessionContextHandle(self._state) + return self._handle + + def __exit__(self, exc_type, exc, tb): + self._handle._active = False + if not self._is_outermost: + return False + + try: + if self._manager.owns_session and self._state.session.in_transaction(): + self._state.session.rollback() + return self._manager.__exit__(exc_type, exc, tb) + finally: + _current_context.reset(self._token) + + async def __aenter__(self): + current = _get_context_for_entry(self.source, is_async=True) + if current is not None: + if not current.is_async: + raise TypeError("Use 'with' inside a synchronous session context.") + _validate_nested_source(current, self.source) + self._state = current + else: + if self.source is None: + raise ValueError( + "An outer session_context() requires a session source." + ) + self._manager = AsyncTransactionSessionContextManager( + session=self.source, + ) + session = await self._manager.__aenter__() + self._state = _ContextState( + session=session, + source=self.source, + is_async=True, + owner_task=_get_asyncio_task(), + ) + self._token = _current_context.set(self._state) + self._is_outermost = True + + self._handle = SessionContextHandle(self._state) + return self._handle + + async def __aexit__(self, exc_type, exc, tb): + self._handle._active = False + if not self._is_outermost: + return False + + try: + if self._manager.owns_session and self._state.session.in_transaction(): + await self._state.session.rollback() + return await self._manager.__aexit__(exc_type, exc, tb) + finally: + _current_context.reset(self._token) + + +class _TransactionContext: + def __init__(self, source): + self.source = source + self._state = None + self._token = None + self._manager = None + self._transaction = None + self._handle = None + self._is_outermost = False + + def __enter__(self): + current = _get_context_for_entry(self.source, is_async=False) + if current is not None: + if current.is_async: + raise TypeError("Use 'async with' inside an async transaction context.") + _validate_nested_source(current, self.source) + self._state = current + self._transaction = self._state.session.begin_nested() + else: + if self.source is None: + raise ValueError("An outer transaction() requires a session source.") + self._manager = TransactionSessionContextManager( + session=self.source, + ) + session = self._manager.__enter__() + self._state = _ContextState( + session=session, + source=self.source, + is_async=False, + owner_thread_id=threading.get_ident(), + ) + self._token = _current_context.set(self._state) + self._is_outermost = True + if session.in_transaction(): + self._transaction = session.begin_nested() + else: + self._transaction = session.begin() + + self._handle = TransactionHandle(self._state) + self._state.transaction_depth += 1 + return self._handle + + def __exit__(self, exc_type, exc, tb): + self._handle._active = False + manager_exception = (exc_type, exc, tb) + try: + if exc_type is None: + self._transaction.commit() + else: + self._transaction.rollback() + except BaseException: + manager_exception = sys.exc_info() + raise + finally: + self._state.transaction_depth -= 1 + if self._is_outermost: + try: + self._manager.__exit__(*manager_exception) + finally: + _current_context.reset(self._token) + return False + + async def __aenter__(self): + current = _get_context_for_entry(self.source, is_async=True) + if current is not None: + if not current.is_async: + raise TypeError("Use 'with' inside a synchronous transaction context.") + _validate_nested_source(current, self.source) + self._state = current + self._transaction = await self._state.session.begin_nested() + else: + if self.source is None: + raise ValueError("An outer transaction() requires a session source.") + self._manager = AsyncTransactionSessionContextManager( + session=self.source, + ) + session = await self._manager.__aenter__() + self._state = _ContextState( + session=session, + source=self.source, + is_async=True, + owner_task=_get_asyncio_task(), + ) + self._token = _current_context.set(self._state) + self._is_outermost = True + if session.in_transaction(): + self._transaction = await session.begin_nested() + else: + self._transaction = await session.begin() + + self._handle = TransactionHandle(self._state) + self._state.transaction_depth += 1 + return self._handle + + async def __aexit__(self, exc_type, exc, tb): + self._handle._active = False + manager_exception = (exc_type, exc, tb) + try: + if exc_type is None: + await self._transaction.commit() + else: + await self._transaction.rollback() + except BaseException: + manager_exception = sys.exc_info() + raise + finally: + self._state.transaction_depth -= 1 + if self._is_outermost: + try: + await self._manager.__aexit__(*manager_exception) + finally: + _current_context.reset(self._token) + return False + + +def session_context(session_source=None): + """Create an ambient session context with explicit transaction control.""" + return _SessionContext(session_source) + + +def transaction(session_source=None): + """Create an atomic transaction context, using SAVEPOINTs when nested.""" + return _TransactionContext(session_source) diff --git a/sqlalchemy_query_manager/core/sync_query_manager.py b/sqlalchemy_query_manager/core/sync_query_manager.py index 9461038..21529da 100644 --- a/sqlalchemy_query_manager/core/sync_query_manager.py +++ b/sqlalchemy_query_manager/core/sync_query_manager.py @@ -9,6 +9,7 @@ from sqlalchemy.ext.hybrid import hybrid_property from sqlalchemy.orm import InstrumentedAttribute, Session, sessionmaker +from sqlalchemy_query_manager.core.contexts import get_current_context from sqlalchemy_query_manager.core.helpers import AggregateFunc, E, Q, _format_sql_value from sqlalchemy_query_manager.core.types import JoinConfig, JoinType from sqlalchemy_query_manager.core.utils import get_session @@ -22,10 +23,7 @@ def __init__(self, model, session=None): self.session: typing.Union[Session, AsyncSession, sessionmaker] = session - self._to_commit = False - - if isinstance(self.session, sessionmaker): - self._to_commit = True + self._to_commit = isinstance(self.session, sessionmaker) self.fields = None @@ -78,6 +76,9 @@ def _clone(self): return new_manager + def _should_commit(self): + return self._to_commit and get_current_context() is None + def join_models( self, query, @@ -858,7 +859,7 @@ def create(self, session=None, expunge=True, **kwargs): new_obj = self.ConverterConfig.model(**kwargs) session.add(new_obj) - if self._to_commit: + if self._should_commit(): session.commit() else: session.flush() @@ -889,7 +890,7 @@ def bulk_create(self, data: typing.List[typing.Dict], session=None, expunge=True objects = [self.ConverterConfig.model(**item) for item in data] session.add_all(objects) - if self._to_commit: + if self._should_commit(): session.commit() else: session.flush() @@ -990,7 +991,7 @@ def update(self, session=None, expunge=True, **kwargs): session.flush() if not returned_pks: - if self._to_commit: + if self._should_commit(): session.commit() return [] @@ -1003,7 +1004,7 @@ def update(self, session=None, expunge=True, **kwargs): if expunge: session.expunge_all() - if self._to_commit: + if self._should_commit(): session.commit() if not updated_objects: @@ -1026,7 +1027,7 @@ def _update_legacy_path(self, session, expunge=True, **kwargs): if expunge: session.expunge_all() - if self._to_commit: + if self._should_commit(): session.commit() if not updated_objects: @@ -1070,7 +1071,7 @@ def update_raw(self, session=None, **kwargs): result = session.execute(update_query) - if self._to_commit: + if self._should_commit(): session.commit() else: session.flush() @@ -1106,7 +1107,7 @@ def update_or_create(self, session=None, expunge=True, defaults=None, **kwargs): if hasattr(existing, key): setattr(existing, key, value) - if self._to_commit: + if self._should_commit(): session.commit() else: session.flush() @@ -1205,7 +1206,7 @@ def delete(self, session=None, expunge=True): result = session.execute(delete_query) - if self._to_commit: + if self._should_commit(): session.commit() else: session.flush() diff --git a/sqlalchemy_query_manager/core/transaction_context_manager.py b/sqlalchemy_query_manager/core/transaction_context_manager.py index 9c7560d..fa49286 100644 --- a/sqlalchemy_query_manager/core/transaction_context_manager.py +++ b/sqlalchemy_query_manager/core/transaction_context_manager.py @@ -1,90 +1,95 @@ -from contextlib import _AsyncGeneratorContextManager, _GeneratorContextManager - from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.orm import Session, sessionmaker - - -def is_generator_context_manager(obj): - """ - Helper function to check if the object is a generator-based context manager - (created by @contextmanager). - """ - # Ensure the object is callable, and if called, returns a generator. - if callable(obj): - try: - # Check if calling it returns a generator (context manager) - result = obj() - if isinstance(result, _GeneratorContextManager) or isinstance( - result, _AsyncGeneratorContextManager - ): - return True - except Exception: - pass - return False - - -def is_async_context_manager(obj): - return callable(getattr(obj, "__aenter__", None)) and callable( - getattr(obj, "__aexit__", None) - ) +from sqlalchemy.orm import Session class BaseSessionContextManager: - def __init__( - self, - session, - ) -> None: # type: ignore + def __init__(self, session) -> None: # type: ignore self.session = session # if a session is passed, and we need already existing one self.is_session_already_set = False + self.owns_session = False self._to_exit = False self._ctx = None # if a session passed as a context manager + def _create_resource(self): + if callable(self.session): + return self.session() + return self.session + class TransactionSessionContextManager(BaseSessionContextManager): def __enter__(self): # type: ignore - if isinstance(self.session, sessionmaker): - self.resource = self.session().__enter__() - self._to_exit = True - elif isinstance(self.session, Session): + if isinstance(self.session, Session): self.resource = self.session self.is_session_already_set = True - elif is_generator_context_manager(self.session): - self._ctx = self.session() - self.resource = self._ctx.__enter__() - self._to_exit = True else: - raise NotImplementedError + resource = self._create_resource() + + if isinstance(resource, Session): + self.resource = resource + self.owns_session = True + self._to_exit = True + elif callable(getattr(resource, "__enter__", None)): + self._ctx = resource + self.resource = self._ctx.__enter__() + if not isinstance(self.resource, Session): + raise TypeError( + "The session context manager must yield a SQLAlchemy Session." + ) + self.owns_session = True + self._to_exit = True + else: + raise NotImplementedError + return self.resource def __exit__(self, exc_type, exc, tb): # type: ignore - if self._to_exit: - if self._ctx: - self._ctx.__exit__(exc_type, exc, tb) - else: - self.resource.__exit__(exc_type, exc, tb) + if not self._to_exit: + return False + + if self._ctx: + return self._ctx.__exit__(exc_type, exc, tb) + + self.resource.close() + + return False class AsyncTransactionSessionContextManager(BaseSessionContextManager): async def __aenter__(self): # type: ignore - if isinstance(self.session, sessionmaker): - self.resource = await self.session().__aenter__() - self._to_exit = True - elif isinstance(self.session, AsyncSession): + if isinstance(self.session, AsyncSession): self.resource = self.session self.is_session_already_set = True - elif is_generator_context_manager(self.session): - self._ctx = self.session() - self.resource = await self._ctx.__aenter__() - self._to_exit = True else: - raise NotImplementedError + resource = self._create_resource() + + if isinstance(resource, AsyncSession): + self.resource = resource + self.owns_session = True + self._to_exit = True + elif callable(getattr(resource, "__aenter__", None)): + self._ctx = resource + self.resource = await self._ctx.__aenter__() + if not isinstance(self.resource, AsyncSession): + raise TypeError( + "The async session context manager must yield a SQLAlchemy " + "AsyncSession." + ) + self.owns_session = True + self._to_exit = True + else: + raise NotImplementedError + return self.resource async def __aexit__(self, exc_type, exc, tb): # type: ignore - if self._to_exit: - if self._ctx: - await self._ctx.__aexit__(exc_type, exc, tb) - else: - await self.resource.__aexit__(exc_type, exc, tb) + if not self._to_exit: + return False + + if self._ctx: + return await self._ctx.__aexit__(exc_type, exc, tb) + + await self.resource.close() + + return False diff --git a/sqlalchemy_query_manager/core/utils.py b/sqlalchemy_query_manager/core/utils.py index e7c87e9..1772ccc 100644 --- a/sqlalchemy_query_manager/core/utils.py +++ b/sqlalchemy_query_manager/core/utils.py @@ -3,19 +3,23 @@ from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import Session +from sqlalchemy_query_manager.core.contexts import get_current_context from sqlalchemy_query_manager.core.transaction_context_manager import ( AsyncTransactionSessionContextManager, TransactionSessionContextManager, ) -def _get_explicit_session(self, session): +def _get_explicit_session(self, session, current_context=None): if session is not None: return session if getattr(self, "_session_is_explicit", False): return self.session + if current_context is not None: + return current_context.session + return None @@ -23,11 +27,11 @@ def _has_for_update(self): return getattr(self, "_for_update", None) is not None -def _validate_sync_for_update_session(self, session): +def _validate_sync_for_update_session(self, session, current_context=None): if not _has_for_update(self): return - explicit_session = _get_explicit_session(self, session) + explicit_session = _get_explicit_session(self, session, current_context) if not isinstance(explicit_session, Session): raise ValueError( "select_for_update() requires an explicit SQLAlchemy Session. " @@ -35,11 +39,11 @@ def _validate_sync_for_update_session(self, session): ) -def _validate_async_for_update_session(self, session): +def _validate_async_for_update_session(self, session, current_context=None): if not _has_for_update(self): return - explicit_session = _get_explicit_session(self, session) + explicit_session = _get_explicit_session(self, session, current_context) if not isinstance(explicit_session, AsyncSession): raise ValueError( "select_for_update() requires an explicit SQLAlchemy AsyncSession. " @@ -52,15 +56,33 @@ def get_session(func): @wraps(func) def wrapper(self, *args, session=None, **kwargs): - _validate_sync_for_update_session(self, session) - - ctx_manager = TransactionSessionContextManager( - session=session or self.session, - ) + current_context = get_current_context() + if current_context is not None and current_context.is_async: + raise TypeError("A synchronous query cannot use an async session context.") + + explicit_session = _get_explicit_session(self, session) + if ( + current_context is not None + and explicit_session is not None + and explicit_session is not current_context.session + ): + raise ValueError( + "An explicit session cannot override the active session context." + ) + + _validate_sync_for_update_session(self, session, current_context) + + session_source = explicit_session + if session_source is None and current_context is not None: + session_source = current_context.session + if session_source is None: + session_source = self.session + + ctx_manager = TransactionSessionContextManager(session=session_source) with ctx_manager as managed_session: expunge = True - if session or ctx_manager.is_session_already_set: + if explicit_session is not None or ctx_manager.is_session_already_set: expunge = False return func(self, session=managed_session, expunge=expunge, *args, **kwargs) @@ -73,10 +95,30 @@ def get_async_session(func): @wraps(func) async def wrapper(self, *args, session=None, **kwargs): - _validate_async_for_update_session(self, session) + current_context = get_current_context() + if current_context is not None and not current_context.is_async: + raise TypeError("An async query cannot use a synchronous session context.") + + explicit_session = _get_explicit_session(self, session) + if ( + current_context is not None + and explicit_session is not None + and explicit_session is not current_context.session + ): + raise ValueError( + "An explicit session cannot override the active session context." + ) + + _validate_async_for_update_session(self, session, current_context) + + session_source = explicit_session + if session_source is None and current_context is not None: + session_source = current_context.session + if session_source is None: + session_source = self.session async with AsyncTransactionSessionContextManager( - session=session or self.session + session=session_source, ) as managed_session: return await func(self, session=managed_session, *args, **kwargs) diff --git a/tests/transactions/test_async_session_context.py b/tests/transactions/test_async_session_context.py new file mode 100644 index 0000000..4217138 --- /dev/null +++ b/tests/transactions/test_async_session_context.py @@ -0,0 +1,208 @@ +from contextlib import asynccontextmanager + +import pytest +from sqlalchemy import event, select +from sqlalchemy.ext.asyncio import AsyncSession + +from sqlalchemy_query_manager import AsyncQueryManager, session_context, transaction +from tests.models import Group, Item, Owner + + +@pytest.mark.asyncio +async def test_async_session_context__supports_manual_commit_and_rollback( + create_tables, + async_db_sessionmaker, +): + owner_manager = AsyncQueryManager(Owner) + item_manager = AsyncQueryManager(Item) + + async with session_context(async_db_sessionmaker) as work: + await owner_manager.create(first_name="John", last_name="Doe") + await work.flush() + await work.commit() + + await item_manager.create(name="rolled back") + await work.rollback() + + async with async_db_sessionmaker() as session: + assert len((await session.execute(select(Owner))).scalars().all()) == 1 + assert len((await session.execute(select(Item))).scalars().all()) == 0 + + +@pytest.mark.asyncio +async def test_async_session_context__does_not_commit_pending_work_on_exit( + create_tables, + async_db_sessionmaker, +): + async with session_context(async_db_sessionmaker): + await AsyncQueryManager(Item).create(name="not committed") + + async with async_db_sessionmaker() as session: + assert len((await session.execute(select(Item))).scalars().all()) == 0 + + +@pytest.mark.asyncio +async def test_async_session_context__leaves_borrowed_session_lifecycle_to_caller( + create_tables, + async_db_sessionmaker, +): + async with async_db_sessionmaker() as session: + async with session_context(session): + await AsyncQueryManager(Item).create(name="pending") + + assert session.in_transaction() + result = await session.execute(select(Item)) + assert result.scalars().one().name == "pending" + + await session.rollback() + result = await session.execute(select(Item)) + assert result.scalars().all() == [] + + +@pytest.mark.asyncio +async def test_async_session_context__supports_context_manager_provider( + create_tables, + async_db_sessionmaker, +): + @asynccontextmanager + async def session_provider(): + async with async_db_sessionmaker() as session: + yield session + + async with session_context(session_provider) as work: + await AsyncQueryManager(Item).create(name="committed") + await work.commit() + + async with async_db_sessionmaker() as session: + result = await session.execute(select(Item)) + assert result.scalars().one().name == "committed" + + +@pytest.mark.asyncio +async def test_async_session_context__nested_context_reuses_session( + create_tables, + async_db_sessionmaker, +): + async with session_context(async_db_sessionmaker) as outer: + owner = await AsyncQueryManager(Owner).create( + first_name="John", + last_name="Doe", + ) + + async with session_context() as inner: + await AsyncQueryManager(Group).create(name="team", owner_id=owner.id) + await inner.flush() + + await outer.commit() + + async with async_db_sessionmaker() as session: + owners = (await session.execute(select(Owner))).scalars().all() + groups = (await session.execute(select(Group))).scalars().all() + assert len(owners) == 1 + assert len(groups) == 1 + + +@pytest.mark.asyncio +async def test_async_session_context__transaction_does_not_commit_outer_transaction( + create_tables, + async_db_sessionmaker, +): + async with session_context(async_db_sessionmaker) as work: + await AsyncQueryManager(Owner).create( + first_name="John", + last_name="Doe", + ) + async with transaction(): + await AsyncQueryManager(Group).create(name="savepoint") + await work.rollback() + + async with async_db_sessionmaker() as session: + owners = (await session.execute(select(Owner))).scalars().all() + groups = (await session.execute(select(Group))).scalars().all() + assert owners == [] + assert groups == [] + + +@pytest.mark.asyncio +async def test_async_session_context__nested_context_rejects_different_source( + create_tables, + async_db_sessionmaker, +): + def different_source(): + return async_db_sessionmaker() + + async with session_context(async_db_sessionmaker): + with pytest.raises(ValueError, match="must use the ambient session"): + async with session_context(different_source): + pass + + +@pytest.mark.asyncio +async def test_async_session_context__handle_is_inactive_after_exit( + create_tables, + async_db_sessionmaker, +): + async with session_context(async_db_sessionmaker) as work: + pass + + for method_name in ("flush", "commit", "rollback"): + with pytest.raises(RuntimeError, match="no longer active"): + await getattr(work, method_name)() + + +@pytest.mark.asyncio +async def test_async_session_context__outer_context_requires_source(): + with pytest.raises(ValueError, match="requires a session source"): + async with session_context(): + pass + + +@pytest.mark.asyncio +async def test_async_session_context__rejects_sync_usage_inside_async_context( + create_tables, + async_db_sessionmaker, +): + async with session_context(async_db_sessionmaker): + with pytest.raises(TypeError, match="async session context"): + with session_context(): + pass + + +@pytest.mark.asyncio +async def test_async_read_with_sessionmaker__does_not_commit( + create_tables, + async_db_sessionmaker, +): + commits = [] + + def record_commit(session): + commits.append(session) + + event.listen(AsyncSession.sync_session_class, "after_commit", record_commit) + try: + assert await AsyncQueryManager(Item, async_db_sessionmaker).count() == 0 + finally: + event.remove(AsyncSession.sync_session_class, "after_commit", record_commit) + + assert commits == [] + + +@pytest.mark.asyncio +async def test_async_explicit_session__preserves_configured_sessionmaker_auto_commit( + create_tables, + async_db_sessionmaker, + async_item_sql_query_manager, +): + async with async_db_sessionmaker() as session: + await async_item_sql_query_manager.query_manager.create( + session=session, + name="committed", + ) + + async with async_db_sessionmaker() as other_session: + result = await other_session.execute(select(Item)) + assert len(result.scalars().all()) == 1 + + async with async_db_sessionmaker() as verification_session: + result = await verification_session.execute(select(Item)) + assert len(result.scalars().all()) == 1 diff --git a/tests/transactions/test_async_transaction.py b/tests/transactions/test_async_transaction.py new file mode 100644 index 0000000..ed5a32a --- /dev/null +++ b/tests/transactions/test_async_transaction.py @@ -0,0 +1,414 @@ +import asyncio +import contextvars +from concurrent.futures import ThreadPoolExecutor +from contextlib import asynccontextmanager + +import pytest +from sqlalchemy import select + +from sqlalchemy_query_manager import ( + AsyncQueryManager, + QueryManager, + session_context, + transaction, +) +from tests.models import Group, Item, Owner + + +WRITE_OPERATIONS = ( + "create", + "bulk_create", + "get_or_create", + "update", + "update_raw", + "update_or_create", + "bulk_update", + "delete", +) +OPERATIONS_REQUIRING_EXISTING_ITEM = { + "update", + "update_raw", + "update_or_create", + "bulk_update", + "delete", +} + + +async def _run_write_operation(operation, item_id=None): + manager = AsyncQueryManager(Item) + + if operation == "create": + await manager.create(name="created") + elif operation == "bulk_create": + await manager.bulk_create([{"name": "first"}, {"name": "second"}]) + elif operation == "get_or_create": + await manager.get_or_create(name="created") + elif operation == "update": + await manager.where(id=item_id).update(name="updated") + elif operation == "update_raw": + await manager.where(id=item_id).update_raw(name="updated") + elif operation == "update_or_create": + await manager.update_or_create(id=item_id, defaults={"name": "updated"}) + elif operation == "bulk_update": + await manager.bulk_update([{"id": item_id, "name": "updated"}]) + elif operation == "delete": + await manager.where(id=item_id).delete() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", WRITE_OPERATIONS) +async def test_async_transaction__write_operation_is_rolled_back( + operation, + create_tables, + async_db_sessionmaker, +): + item_id = None + if operation in OPERATIONS_REQUIRING_EXISTING_ITEM: + async with async_db_sessionmaker() as session: + item = Item(name="original") + session.add(item) + await session.commit() + item_id = item.id + + with pytest.raises(RuntimeError, match="abort"): + async with transaction(async_db_sessionmaker): + await _run_write_operation(operation, item_id) + raise RuntimeError("abort") + + async with async_db_sessionmaker() as session: + items = (await session.execute(select(Item))).scalars().all() + if operation in OPERATIONS_REQUIRING_EXISTING_ITEM: + assert len(items) == 1 + assert items[0].name == "original" + else: + assert items == [] + + +@pytest.mark.asyncio +async def test_async_transaction__uses_one_session_across_models_and_commits( + create_tables, + async_db_sessionmaker, +): + owner_manager = AsyncQueryManager(Owner) + group_manager = AsyncQueryManager(Group) + item_manager = AsyncQueryManager(Item) + + async with transaction(async_db_sessionmaker) as control: + owner = await owner_manager.create(first_name="John", last_name="Doe") + group = await group_manager.create(name="team", owner_id=owner.id) + await item_manager.create(name="item", group_id=group.id) + + await control.flush() + + assert await owner_manager.count() == 1 + assert await group_manager.count() == 1 + assert await item_manager.count() == 1 + + async with async_db_sessionmaker() as session: + assert len((await session.execute(select(Owner))).scalars().all()) == 1 + assert len((await session.execute(select(Group))).scalars().all()) == 1 + assert len((await session.execute(select(Item))).scalars().all()) == 1 + + +@pytest.mark.asyncio +async def test_async_transaction__exception_rolls_back_all_models( + create_tables, + async_db_sessionmaker, +): + with pytest.raises(RuntimeError, match="abort"): + async with transaction(async_db_sessionmaker): + await AsyncQueryManager(Owner).create( + first_name="John", + last_name="Doe", + ) + await AsyncQueryManager(Group).create(name="team") + raise RuntimeError("abort") + + async with async_db_sessionmaker() as session: + assert len((await session.execute(select(Owner))).scalars().all()) == 0 + assert len((await session.execute(select(Group))).scalars().all()) == 0 + + +@pytest.mark.asyncio +async def test_async_transaction__nested_exception_rolls_back_savepoint_only( + create_tables, + async_db_sessionmaker, +): + async with transaction(async_db_sessionmaker): + await AsyncQueryManager(Owner).create( + first_name="John", + last_name="Doe", + ) + + with pytest.raises(RuntimeError, match="nested abort"): + async with transaction(): + await AsyncQueryManager(Group).create(name="discarded") + raise RuntimeError("nested abort") + + await AsyncQueryManager(Item).create(name="survives") + + async with async_db_sessionmaker() as session: + assert len((await session.execute(select(Owner))).scalars().all()) == 1 + assert len((await session.execute(select(Group))).scalars().all()) == 0 + assert len((await session.execute(select(Item))).scalars().all()) == 1 + + +@pytest.mark.asyncio +async def test_async_transaction__outer_rollback_includes_released_savepoint( + create_tables, + async_db_sessionmaker, +): + with pytest.raises(RuntimeError, match="outer abort"): + async with transaction(async_db_sessionmaker): + async with transaction(): + await AsyncQueryManager(Group).create(name="released savepoint") + raise RuntimeError("outer abort") + + async with async_db_sessionmaker() as session: + assert len((await session.execute(select(Group))).scalars().all()) == 0 + + +@pytest.mark.asyncio +async def test_async_transaction__manual_commit_and_rollback_are_rejected( + create_tables, + async_db_sessionmaker, +): + async with transaction(async_db_sessionmaker) as control: + with pytest.raises(RuntimeError, match="Manual commit"): + await control.commit() + with pytest.raises(RuntimeError, match="Manual rollback"): + await control.rollback() + + +@pytest.mark.asyncio +async def test_async_transaction__handle_is_inactive_after_exit( + create_tables, + async_db_sessionmaker, +): + async with transaction(async_db_sessionmaker) as control: + pass + + with pytest.raises(RuntimeError, match="no longer active"): + await control.flush() + + +@pytest.mark.asyncio +async def test_async_nested_session_context__cannot_commit_transaction( + create_tables, + async_db_sessionmaker, +): + async with transaction(async_db_sessionmaker): + async with session_context() as work: + with pytest.raises(RuntimeError, match="Manual commit"): + await work.commit() + with pytest.raises(RuntimeError, match="Manual rollback"): + await work.rollback() + + +@pytest.mark.asyncio +async def test_async_transaction__supports_context_manager_provider( + create_tables, + async_db_sessionmaker, +): + @asynccontextmanager + async def session_provider(): + async with async_db_sessionmaker() as session: + yield session + + async with transaction(session_provider): + await AsyncQueryManager(Owner).create( + first_name="John", + last_name="Doe", + ) + + async with async_db_sessionmaker() as session: + assert len((await session.execute(select(Owner))).scalars().all()) == 1 + + +@pytest.mark.asyncio +async def test_async_transaction__supports_borrowed_session_without_closing_it( + create_tables, + async_db_sessionmaker, +): + async with async_db_sessionmaker() as session: + async with transaction(session): + await AsyncQueryManager(Item).create(name="created") + + result = await session.execute(select(Item)) + assert result.scalars().one().name == "created" + + +@pytest.mark.asyncio +async def test_async_transaction__uses_savepoint_for_active_borrowed_session( + create_tables, + async_db_sessionmaker, +): + async with async_db_sessionmaker() as session: + await session.begin() + session.add(Owner(first_name="John", last_name="Doe")) + + async with transaction(session): + await AsyncQueryManager(Group).create(name="nested") + + await session.rollback() + + async with async_db_sessionmaker() as verification_session: + owners = (await verification_session.execute(select(Owner))).scalars().all() + groups = (await verification_session.execute(select(Group))).scalars().all() + assert owners == [] + assert groups == [] + + +@pytest.mark.asyncio +async def test_async_transaction__rejects_explicit_session_override( + create_tables, + async_db_sessionmaker, +): + async with async_db_sessionmaker() as other_session: + async with transaction(async_db_sessionmaker): + with pytest.raises(ValueError, match="cannot override"): + await AsyncQueryManager(Item).create( + session=other_session, + name="wrong session", + ) + + +@pytest.mark.asyncio +async def test_async_transaction__select_for_update_uses_ambient_session( + create_tables, + async_db_sessionmaker, +): + manager = AsyncQueryManager(Item, async_db_sessionmaker) + item = await manager.create(name="lock me") + + async with transaction(async_db_sessionmaker): + locked = await manager.where(id=item.id).select_for_update().get() + assert locked.id == item.id + + +@pytest.mark.asyncio +async def test_async_transaction__context_is_reset_after_exit( + create_tables, + async_db_sessionmaker, +): + async with transaction(async_db_sessionmaker): + pass + + with pytest.raises(ValueError, match="requires a session source"): + async with transaction(): + pass + + with pytest.raises(RuntimeError, match="abort"): + async with transaction(async_db_sessionmaker): + raise RuntimeError("abort") + + with pytest.raises(ValueError, match="requires a session source"): + async with transaction(): + pass + + +@pytest.mark.asyncio +async def test_async_transaction__context_is_reset_when_commit_fails( + create_tables, + async_db_sessionmaker, + monkeypatch, +): + context = transaction(async_db_sessionmaker) + + async def fail_commit(transaction): + raise RuntimeError("commit failed") + + with pytest.raises(RuntimeError, match="commit failed"): + async with context: + monkeypatch.setattr(type(context._transaction), "commit", fail_commit) + + with pytest.raises(ValueError, match="requires a session source"): + async with transaction(): + pass + + +@pytest.mark.asyncio +async def test_async_transaction__rejects_sync_usage_inside_async_context( + create_tables, + async_db_sessionmaker, +): + async with transaction(async_db_sessionmaker): + with pytest.raises(TypeError, match="async transaction context"): + with transaction(): + pass + + +@pytest.mark.asyncio +async def test_async_transaction__provider_must_yield_async_session(): + @asynccontextmanager + async def invalid_provider(): + yield object() + + with pytest.raises(TypeError, match="must yield a SQLAlchemy AsyncSession"): + async with transaction(invalid_provider): + pass + + +@pytest.mark.asyncio +async def test_async_context__cannot_share_session_with_child_task( + create_tables, + async_db_sessionmaker, +): + async with transaction(async_db_sessionmaker): + task = asyncio.create_task(AsyncQueryManager(Item).count()) + with pytest.raises(RuntimeError, match="cannot be shared"): + await task + + +@pytest.mark.asyncio +async def test_async_context__cannot_share_session_with_child_thread( + create_tables, + async_db_sessionmaker, +): + def use_inherited_context(): + QueryManager(Item).count() + + loop = asyncio.get_running_loop() + async with transaction(async_db_sessionmaker): + inherited_context = contextvars.copy_context() + with ThreadPoolExecutor(max_workers=1) as executor: + future = loop.run_in_executor( + executor, + inherited_context.run, + use_inherited_context, + ) + with pytest.raises(RuntimeError, match="cannot be shared"): + await future + + +@pytest.mark.asyncio +async def test_async_context__child_task_cannot_reuse_ambient_session( + create_tables, + async_db_sessionmaker, +): + async with async_db_sessionmaker() as session: + + async def reuse_ambient_session(): + async with transaction(session): + pass + + async with transaction(session): + task = asyncio.create_task(reuse_ambient_session()) + with pytest.raises(RuntimeError, match="cannot inherit"): + await task + + +@pytest.mark.asyncio +async def test_async_context__child_task_can_open_independent_context( + create_tables, + async_db_sessionmaker, +): + async def create_in_independent_context(): + async with transaction(async_db_sessionmaker): + await AsyncQueryManager(Item).create(name="child") + + async with session_context(async_db_sessionmaker): + await asyncio.create_task(create_in_independent_context()) + + async with async_db_sessionmaker() as session: + assert len((await session.execute(select(Item))).scalars().all()) == 1 diff --git a/tests/transactions/test_sync_session_context.py b/tests/transactions/test_sync_session_context.py new file mode 100644 index 0000000..680fd13 --- /dev/null +++ b/tests/transactions/test_sync_session_context.py @@ -0,0 +1,194 @@ +import asyncio +from contextlib import contextmanager + +import pytest +from sqlalchemy import event + +from sqlalchemy_query_manager import QueryManager, session_context, transaction +from tests.models import Group, Item, Owner + + +def test_session_context__supports_manual_flush_commit_and_rollback( + create_tables, + sync_db_sessionmaker, +): + owner_manager = QueryManager(Owner) + group_manager = QueryManager(Group) + item_manager = QueryManager(Item) + + with session_context(sync_db_sessionmaker) as work: + owner = owner_manager.create(first_name="John", last_name="Doe") + group_manager.create(name="committed", owner_id=owner.id) + work.flush() + work.commit() + + item_manager.create(name="rolled back") + work.rollback() + + assert owner_manager.count() == 1 + assert group_manager.count() == 1 + assert item_manager.count() == 0 + + with sync_db_sessionmaker() as session: + assert session.query(Owner).count() == 1 + assert session.query(Group).count() == 1 + assert session.query(Item).count() == 0 + + +def test_session_context__does_not_commit_pending_work_on_exit( + create_tables, + sync_db_sessionmaker, +): + with session_context(sync_db_sessionmaker): + QueryManager(Item).create(name="not committed") + + with sync_db_sessionmaker() as session: + assert session.query(Item).count() == 0 + + +def test_session_context__leaves_borrowed_session_lifecycle_to_caller( + create_tables, + sync_db_sessionmaker, +): + session = sync_db_sessionmaker() + try: + with session_context(session): + QueryManager(Item).create(name="pending") + + assert session.in_transaction() + assert session.query(Item).one().name == "pending" + + session.rollback() + assert session.query(Item).count() == 0 + finally: + session.close() + + +def test_session_context__supports_context_manager_provider( + create_tables, + sync_db_sessionmaker, +): + @contextmanager + def session_provider(): + with sync_db_sessionmaker() as session: + yield session + + with session_context(session_provider) as work: + QueryManager(Item).create(name="committed") + work.commit() + + with sync_db_sessionmaker() as session: + assert session.query(Item).one().name == "committed" + + +def test_session_context__nested_context_reuses_session( + create_tables, + sync_db_sessionmaker, +): + with session_context(sync_db_sessionmaker) as outer: + owner = QueryManager(Owner).create(first_name="John", last_name="Doe") + + with session_context() as inner: + QueryManager(Group).create(name="team", owner_id=owner.id) + inner.flush() + + outer.commit() + + with sync_db_sessionmaker() as session: + assert session.query(Owner).count() == 1 + assert session.query(Group).count() == 1 + + +def test_session_context__transaction_does_not_commit_outer_transaction( + create_tables, + sync_db_sessionmaker, +): + with session_context(sync_db_sessionmaker) as work: + QueryManager(Owner).create(first_name="John", last_name="Doe") + with transaction(): + QueryManager(Group).create(name="savepoint") + work.rollback() + + with sync_db_sessionmaker() as session: + assert session.query(Owner).count() == 0 + assert session.query(Group).count() == 0 + + +def test_session_context__nested_context_rejects_different_source( + create_tables, + sync_db_sessionmaker, +): + def different_source(): + return sync_db_sessionmaker() + + with session_context(sync_db_sessionmaker): + with pytest.raises(ValueError, match="must use the ambient session"): + with session_context(different_source): + pass + + +def test_session_context__handle_is_inactive_after_exit( + create_tables, + sync_db_sessionmaker, +): + with session_context(sync_db_sessionmaker) as work: + pass + + for method_name in ("flush", "commit", "rollback"): + with pytest.raises(RuntimeError, match="no longer active"): + getattr(work, method_name)() + + +def test_session_context__outer_context_requires_source(): + with pytest.raises(ValueError, match="requires a session source"): + with session_context(): + pass + + +def test_session_context__rejects_async_usage_inside_sync_context( + create_tables, + sync_db_sessionmaker, +): + async def enter_async_context(): + async with session_context(): + pass + + with session_context(sync_db_sessionmaker): + with pytest.raises(TypeError, match="synchronous session context"): + asyncio.run(enter_async_context()) + + +def test_explicit_session__preserves_configured_sessionmaker_auto_commit( + create_tables, + sync_db_sessionmaker, + item_sql_query_manager, +): + with sync_db_sessionmaker() as session: + item_sql_query_manager.query_manager.create( + session=session, + name="committed", + ) + + with sync_db_sessionmaker() as other_session: + assert other_session.query(Item).count() == 1 + + with sync_db_sessionmaker() as verification_session: + assert verification_session.query(Item).count() == 1 + + +def test_read_with_sessionmaker__does_not_commit( + create_tables, + sync_db_sessionmaker, +): + commits = [] + + def record_commit(session): + commits.append(session) + + event.listen(sync_db_sessionmaker.class_, "after_commit", record_commit) + try: + assert QueryManager(Item, sync_db_sessionmaker).count() == 0 + finally: + event.remove(sync_db_sessionmaker.class_, "after_commit", record_commit) + + assert commits == [] diff --git a/tests/transactions/test_sync_transaction.py b/tests/transactions/test_sync_transaction.py new file mode 100644 index 0000000..f09769c --- /dev/null +++ b/tests/transactions/test_sync_transaction.py @@ -0,0 +1,377 @@ +import asyncio +import contextvars +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager + +import pytest +from sqlalchemy import select + +from sqlalchemy_query_manager import QueryManager, session_context, transaction +from tests.models import Group, Item, Owner + + +WRITE_OPERATIONS = ( + "create", + "bulk_create", + "get_or_create", + "update", + "update_raw", + "update_or_create", + "bulk_update", + "delete", +) +OPERATIONS_REQUIRING_EXISTING_ITEM = { + "update", + "update_raw", + "update_or_create", + "bulk_update", + "delete", +} + + +def _run_write_operation(operation, item_id=None): + manager = QueryManager(Item) + + if operation == "create": + manager.create(name="created") + elif operation == "bulk_create": + manager.bulk_create([{"name": "first"}, {"name": "second"}]) + elif operation == "get_or_create": + manager.get_or_create(name="created") + elif operation == "update": + manager.where(id=item_id).update(name="updated") + elif operation == "update_raw": + manager.where(id=item_id).update_raw(name="updated") + elif operation == "update_or_create": + manager.update_or_create(id=item_id, defaults={"name": "updated"}) + elif operation == "bulk_update": + manager.bulk_update([{"id": item_id, "name": "updated"}]) + elif operation == "delete": + manager.where(id=item_id).delete() + + +@pytest.mark.parametrize("operation", WRITE_OPERATIONS) +def test_transaction__write_operation_is_rolled_back( + operation, + create_tables, + sync_db_sessionmaker, +): + item_id = None + if operation in OPERATIONS_REQUIRING_EXISTING_ITEM: + with sync_db_sessionmaker() as session: + item = Item(name="original") + session.add(item) + session.commit() + item_id = item.id + + with pytest.raises(RuntimeError, match="abort"): + with transaction(sync_db_sessionmaker): + _run_write_operation(operation, item_id) + raise RuntimeError("abort") + + with sync_db_sessionmaker() as session: + items = session.query(Item).all() + if operation in OPERATIONS_REQUIRING_EXISTING_ITEM: + assert len(items) == 1 + assert items[0].name == "original" + else: + assert items == [] + + +def test_transaction__uses_one_session_across_models_and_commits( + create_tables, + sync_db_sessionmaker, +): + owner_manager = QueryManager(Owner) + group_manager = QueryManager(Group) + item_manager = QueryManager(Item) + + with transaction(sync_db_sessionmaker) as control: + owner = owner_manager.create(first_name="John", last_name="Doe") + group = group_manager.create(name="team", owner_id=owner.id) + item_manager.create(name="item", group_id=group.id) + + control.flush() + + assert owner_manager.count() == 1 + assert group_manager.count() == 1 + assert item_manager.count() == 1 + + with sync_db_sessionmaker() as session: + assert session.query(Owner).count() == 1 + assert session.query(Group).count() == 1 + assert session.query(Item).count() == 1 + + +def test_transaction__exception_rolls_back_all_models( + create_tables, + sync_db_sessionmaker, +): + with pytest.raises(RuntimeError, match="abort"): + with transaction(sync_db_sessionmaker): + QueryManager(Owner).create(first_name="John", last_name="Doe") + QueryManager(Group).create(name="team") + raise RuntimeError("abort") + + with sync_db_sessionmaker() as session: + assert session.query(Owner).count() == 0 + assert session.query(Group).count() == 0 + + +def test_transaction__nested_exception_rolls_back_savepoint_only( + create_tables, + sync_db_sessionmaker, +): + with transaction(sync_db_sessionmaker): + QueryManager(Owner).create(first_name="John", last_name="Doe") + + with pytest.raises(RuntimeError, match="nested abort"): + with transaction(): + QueryManager(Group).create(name="discarded") + raise RuntimeError("nested abort") + + QueryManager(Item).create(name="survives") + + with sync_db_sessionmaker() as session: + assert session.query(Owner).count() == 1 + assert session.query(Group).count() == 0 + assert session.query(Item).count() == 1 + + +def test_transaction__outer_rollback_includes_released_savepoint( + create_tables, + sync_db_sessionmaker, +): + with pytest.raises(RuntimeError, match="outer abort"): + with transaction(sync_db_sessionmaker): + with transaction(): + QueryManager(Group).create(name="released savepoint") + raise RuntimeError("outer abort") + + with sync_db_sessionmaker() as session: + assert session.query(Group).count() == 0 + + +def test_transaction__manual_commit_and_rollback_are_rejected( + create_tables, + sync_db_sessionmaker, +): + with transaction(sync_db_sessionmaker) as control: + with pytest.raises(RuntimeError, match="Manual commit"): + control.commit() + with pytest.raises(RuntimeError, match="Manual rollback"): + control.rollback() + + +def test_transaction__handle_is_inactive_after_exit( + create_tables, + sync_db_sessionmaker, +): + with transaction(sync_db_sessionmaker) as control: + pass + + with pytest.raises(RuntimeError, match="no longer active"): + control.flush() + + +def test_transaction__context_is_reset_after_exit( + create_tables, + sync_db_sessionmaker, +): + with transaction(sync_db_sessionmaker): + pass + + with pytest.raises(ValueError, match="requires a session source"): + with transaction(): + pass + + with pytest.raises(RuntimeError, match="abort"): + with transaction(sync_db_sessionmaker): + raise RuntimeError("abort") + + with pytest.raises(ValueError, match="requires a session source"): + with transaction(): + pass + + +def test_transaction__context_is_reset_when_commit_fails( + create_tables, + sync_db_sessionmaker, + monkeypatch, +): + context = transaction(sync_db_sessionmaker) + + def fail_commit(): + raise RuntimeError("commit failed") + + with pytest.raises(RuntimeError, match="commit failed"): + with context: + monkeypatch.setattr(context._transaction, "commit", fail_commit) + + with pytest.raises(ValueError, match="requires a session source"): + with transaction(): + pass + + +def test_transaction__rejects_async_usage_inside_sync_context( + create_tables, + sync_db_sessionmaker, +): + async def enter_async_transaction(): + async with transaction(): + pass + + with transaction(sync_db_sessionmaker): + with pytest.raises(TypeError, match="synchronous transaction context"): + asyncio.run(enter_async_transaction()) + + +def test_transaction__provider_must_yield_session(sync_db_sessionmaker): + @contextmanager + def invalid_provider(): + yield object() + + with pytest.raises(TypeError, match="must yield a SQLAlchemy Session"): + with transaction(invalid_provider): + pass + + +def test_nested_session_context__cannot_commit_transaction( + create_tables, + sync_db_sessionmaker, +): + with transaction(sync_db_sessionmaker): + with session_context() as work: + with pytest.raises(RuntimeError, match="Manual commit"): + work.commit() + with pytest.raises(RuntimeError, match="Manual rollback"): + work.rollback() + + +def test_transaction__supports_borrowed_session_without_closing_it( + create_tables, + sync_db_sessionmaker, +): + session = sync_db_sessionmaker() + try: + with transaction(session): + QueryManager(Item).create(name="created") + + assert session.execute(select(Item)).scalars().one().name == "created" + finally: + session.close() + + +def test_transaction__uses_savepoint_for_active_borrowed_session( + create_tables, + sync_db_sessionmaker, +): + session = sync_db_sessionmaker() + try: + session.begin() + session.add(Owner(first_name="John", last_name="Doe")) + + with transaction(session): + QueryManager(Group).create(name="nested") + + session.rollback() + + with sync_db_sessionmaker() as verification_session: + assert verification_session.query(Owner).count() == 0 + assert verification_session.query(Group).count() == 0 + finally: + session.close() + + +def test_transaction__supports_context_manager_provider( + create_tables, + sync_db_sessionmaker, +): + @contextmanager + def session_provider(): + with sync_db_sessionmaker() as session: + yield session + + with transaction(session_provider): + QueryManager(Owner).create(first_name="John", last_name="Doe") + QueryManager(Group).create(name="team") + + with sync_db_sessionmaker() as session: + assert session.query(Owner).count() == 1 + assert session.query(Group).count() == 1 + + +def test_transaction__rejects_explicit_session_override( + create_tables, + sync_db_sessionmaker, +): + with sync_db_sessionmaker() as other_session: + with transaction(sync_db_sessionmaker): + with pytest.raises(ValueError, match="cannot override"): + QueryManager(Item).create(session=other_session, name="wrong session") + + +def test_transaction__select_for_update_uses_ambient_session( + create_tables, + sync_db_sessionmaker, + item_sql_query_manager, +): + item = item_sql_query_manager.query_manager.create(name="lock me") + + with transaction(sync_db_sessionmaker): + locked = ( + item_sql_query_manager.query_manager.where(id=item.id) + .select_for_update() + .get() + ) + assert locked.id == item.id + + +def test_sync_context__cannot_share_session_with_child_thread( + create_tables, + sync_db_sessionmaker, +): + with transaction(sync_db_sessionmaker): + context = contextvars.copy_context() + with ThreadPoolExecutor(max_workers=1) as executor: + result = executor.submit(context.run, QueryManager(Item).count) + with pytest.raises(RuntimeError, match="cannot be shared across threads"): + result.result() + + +def test_sync_context__child_thread_can_open_independent_context( + create_tables, + sync_db_sessionmaker, +): + def create_in_independent_context(): + with transaction(sync_db_sessionmaker): + QueryManager(Item).create(name="child") + + with session_context(sync_db_sessionmaker): + context = contextvars.copy_context() + with ThreadPoolExecutor(max_workers=1) as executor: + executor.submit(context.run, create_in_independent_context).result() + + with sync_db_sessionmaker() as session: + assert session.query(Item).count() == 1 + + +def test_sync_context__child_thread_cannot_reuse_ambient_session( + create_tables, + sync_db_sessionmaker, +): + session = sync_db_sessionmaker() + + def reuse_ambient_session(): + with transaction(session): + pass + + try: + with transaction(session): + context = contextvars.copy_context() + with ThreadPoolExecutor(max_workers=1) as executor: + result = executor.submit(context.run, reuse_ambient_session) + with pytest.raises(RuntimeError, match="cannot inherit"): + result.result() + finally: + session.close()