fix: lock_tables acquires dedicated session to enforce RAII lock boundaries (#261)

src/fastapi_toolsets/db.py

@@ -153,7 +153,7 @@ class LockMode(str, Enum):
 
 @asynccontextmanager
 async def lock_tables(
-    session: AsyncSession,
+    session_maker: async_sessionmaker[AsyncSession],
     tables: list[type[DeclarativeBase]],
     *,
     mode: LockMode = LockMode.SHARE_UPDATE_EXCLUSIVE,
@@ -161,42 +161,44 @@ async def lock_tables(
 ) -> AsyncGenerator[AsyncSession, None]:
     """Lock PostgreSQL tables for the duration of a transaction.
 
-    Acquires table-level locks that are held until the transaction ends.
-    Useful for preventing concurrent modifications during critical operations.
-
     Args:
-        session: AsyncSession instance
-        tables: List of SQLAlchemy model classes to lock
-        mode: Lock mode (default: SHARE UPDATE EXCLUSIVE)
-        timeout: Lock timeout (default: "5s")
+        session_maker: Async session factory used to create the dedicated
+            session.
+        tables: List of SQLAlchemy model classes to lock.
+        mode: Lock mode (default: SHARE UPDATE EXCLUSIVE).
+        timeout: Lock timeout (default: "5s").
 
     Yields:
-        The session with locked tables
+        The dedicated session, open within the locked transaction.
 
     Raises:
-        SQLAlchemyError: If lock cannot be acquired within timeout
+        SQLAlchemyError: If the lock cannot be acquired within *timeout*.
 
     Example:
         ```python
         from fastapi_toolsets.db import lock_tables, LockMode
 
-        async with lock_tables(session, [User, Account]):
-            # Tables are locked with SHARE UPDATE EXCLUSIVE mode
+        async with lock_tables(session_maker, [User, Account]) as session:
+            # Tables are locked; changes are committed when the context exits.
             user = await UserCrud.get(session, [User.id == 1])
             user.balance += 100
 
         # With custom lock mode
-        async with lock_tables(session, [Order], mode=LockMode.EXCLUSIVE):
-            # Exclusive lock - no other transactions can access
+        async with lock_tables(session_maker, [Order], mode=LockMode.EXCLUSIVE) as session:
             await process_order(session, order_id)
         ```
     """
     table_names = ",".join(table.__tablename__ for table in tables)
 
-    async with get_transaction(session):
-        await session.execute(text(f"SET LOCAL lock_timeout='{timeout}'"))
-        await session.execute(text(f"LOCK {table_names} IN {mode.value} MODE"))
-        yield session
+    async with session_maker() as session:
+        try:
+            await session.execute(text(f"SET LOCAL lock_timeout='{timeout}'"))
+            await session.execute(text(f"LOCK {table_names} IN {mode.value} MODE"))
+            yield session
+            await session.commit()
+        except BaseException:
+            await session.rollback()
+            raise
 
 
 async def create_database(

tests/conftest.py

@@ -439,6 +439,21 @@ async def engine():
     await engine.dispose()
 
 
+@pytest.fixture(scope="function")
+async def session_maker(engine):
+    """Provide a session factory with tables created and dropped around the test."""
+    async with engine.begin() as conn:
+        await conn.run_sync(Base.metadata.create_all)
+
+    factory = async_sessionmaker(engine, expire_on_commit=False)
+
+    try:
+        yield factory
+    finally:
+        async with engine.begin() as conn:
+            await conn.run_sync(Base.metadata.drop_all)
+
+
 @pytest.fixture(scope="function")
 async def db_session(engine):
     """Create a test database session with tables.

tests/test_db.py

@@ -116,13 +116,8 @@ class TestCreateDbDependency:
         await engine.dispose()
 
     @pytest.mark.anyio
-    async def test_update_after_lock_tables_is_persisted(self):
-        """Changes made after lock_tables exits (before endpoint returns) are committed.
-
-        Regression: without the auto-begin fix, lock_tables would start and commit a
-        real outer transaction, leaving the session idle. Any modifications after that
-        point were silently dropped.
-        """
+    async def test_data_inside_lock_is_committed(self):
+        """Changes made inside lock_tables are committed when the context exits."""
         engine = create_async_engine(DATABASE_URL, echo=False)
         session_factory = async_sessionmaker(engine, expire_on_commit=False)
 
@@ -130,21 +125,12 @@ class TestCreateDbDependency:
             await conn.run_sync(Base.metadata.create_all)
 
         try:
-            get_db = create_db_dependency(session_factory)
-
-            async for session in get_db():
-                async with lock_tables(session, [Role]):
-                    role = Role(name="lock_then_update")
-                    session.add(role)
-                    await session.flush()
-                # lock_tables has exited — outer transaction must still be open
-                assert session.in_transaction()
-                role.name = "updated_after_lock"
+            async with lock_tables(session_factory, [Role]) as session:
+                role = Role(name="lock_committed")
+                session.add(role)
 
             async with session_factory() as verify:
-                result = await RoleCrud.first(
-                    verify, [Role.name == "updated_after_lock"]
-                )
+                result = await RoleCrud.first(verify, [Role.name == "lock_committed"])
                 assert result is not None
         finally:
             async with engine.begin() as conn:
@@ -287,54 +273,55 @@ class TestLockTables:
     """Tests for lock_tables context manager (PostgreSQL-specific)."""
 
     @pytest.mark.anyio
-    async def test_lock_single_table(self, db_session: AsyncSession):
-        """Lock a single table."""
-        async with lock_tables(db_session, [Role]):
-            # Inside the lock, we can still perform operations
+    async def test_lock_single_table(self, session_maker):
+        """Lock a single table; changes inside are committed on context exit."""
+        async with lock_tables(session_maker, [Role]) as session:
             role = Role(name="locked_role")
-            db_session.add(role)
-            await db_session.flush()
+            session.add(role)
 
-        # After lock is released, verify the data was committed
-        result = await RoleCrud.first(db_session, [Role.name == "locked_role"])
-        assert result is not None
+        async with session_maker() as verify:
+            result = await RoleCrud.first(verify, [Role.name == "locked_role"])
+            assert result is not None
 
     @pytest.mark.anyio
-    async def test_lock_multiple_tables(self, db_session: AsyncSession):
+    async def test_lock_multiple_tables(self, session_maker):
         """Lock multiple tables."""
-        async with lock_tables(db_session, [Role, User]):
+        async with lock_tables(session_maker, [Role, User]) as session:
             role = Role(name="multi_lock_role")
-            db_session.add(role)
-            await db_session.flush()
+            session.add(role)
 
-        result = await RoleCrud.first(db_session, [Role.name == "multi_lock_role"])
-        assert result is not None
+        async with session_maker() as verify:
+            result = await RoleCrud.first(verify, [Role.name == "multi_lock_role"])
+            assert result is not None
 
     @pytest.mark.anyio
-    async def test_lock_with_custom_mode(self, db_session: AsyncSession):
+    async def test_lock_with_custom_mode(self, session_maker):
         """Lock with custom lock mode."""
-        async with lock_tables(db_session, [Role], mode=LockMode.EXCLUSIVE):
+        async with lock_tables(
+            session_maker, [Role], mode=LockMode.EXCLUSIVE
+        ) as session:
             role = Role(name="exclusive_lock_role")
-            db_session.add(role)
-            await db_session.flush()
+            session.add(role)
 
-        result = await RoleCrud.first(db_session, [Role.name == "exclusive_lock_role"])
-        assert result is not None
+        async with session_maker() as verify:
+            result = await RoleCrud.first(verify, [Role.name == "exclusive_lock_role"])
+            assert result is not None
 
     @pytest.mark.anyio
-    async def test_lock_rollback_on_exception(self, db_session: AsyncSession):
+    async def test_lock_rollback_on_exception(self, session_maker):
         """Lock context rolls back on exception."""
         try:
-            async with lock_tables(db_session, [Role]):
+            async with lock_tables(session_maker, [Role]) as session:
                 role = Role(name="lock_rollback_role")
-                db_session.add(role)
-                await db_session.flush()
+                session.add(role)
+                await session.flush()
                 raise ValueError("Simulated error")
         except ValueError:
             pass
 
-        result = await RoleCrud.first(db_session, [Role.name == "lock_rollback_role"])
-        assert result is None
+        async with session_maker() as verify:
+            result = await RoleCrud.first(verify, [Role.name == "lock_rollback_role"])
+            assert result is None
 
 
 class TestWaitForRowChange:
@@ -643,29 +630,30 @@ class TestM2MAdd:
             await m2m_add(db_session, user, User.role, role)
 
     @pytest.mark.anyio
-    async def test_works_inside_lock_tables(self, db_session: AsyncSession):
-        """m2m_add works correctly inside a lock_tables nested transaction."""
-        user = User(username="m2m_lock_author", email="m2m_lock@test.com")
-        db_session.add(user)
-        await db_session.flush()
+    async def test_works_inside_lock_tables(self, session_maker):
+        """m2m_add works correctly inside a lock_tables context."""
+        async with lock_tables(session_maker, [Tag]) as session:
+            user = User(username="m2m_lock_author", email="m2m_lock@test.com")
+            session.add(user)
+            await session.flush()
 
-        async with lock_tables(db_session, [Tag]):
             tag = Tag(name="locked_tag")
-            db_session.add(tag)
-            await db_session.flush()
+            session.add(tag)
+            await session.flush()
 
             post = Post(title="Post Lock", author_id=user.id)
-            db_session.add(post)
-            await db_session.flush()
+            session.add(post)
+            await session.flush()
 
-            await m2m_add(db_session, post, Post.tags, tag)
+            await m2m_add(session, post, Post.tags, tag)
 
-        result = await db_session.execute(
-            select(Post).where(Post.id == post.id).options(selectinload(Post.tags))
-        )
-        loaded = result.scalar_one()
-        assert len(loaded.tags) == 1
-        assert loaded.tags[0].name == "locked_tag"
+        async with session_maker() as verify:
+            result = await verify.execute(
+                select(Post).where(Post.id == post.id).options(selectinload(Post.tags))
+            )
+            loaded = result.scalar_one()
+            assert len(loaded.tags) == 1
+            assert loaded.tags[0].name == "locked_tag"
 
 
 class _LocalBase(DeclarativeBase):

tests/test_models.py

@@ -7,6 +7,7 @@ from unittest.mock import patch
 
 import pytest
 from sqlalchemy import String
+from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
 from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
 
 import fastapi_toolsets.models.watched as _watched_module
@@ -20,6 +21,7 @@ from fastapi_toolsets.models import (
     listens_for,
 )
 from fastapi_toolsets.models.watched import (
+    EventSession,
     _EVENT_HANDLERS,
     _SESSION_CREATES,
     _SESSION_DELETES,
@@ -338,6 +340,23 @@ async def mixin_session_expire():
         yield session
 
 
+@pytest.fixture(scope="function")
+async def mixin_session_maker():
+    """Provide an EventSession-backed session factory with MixinBase tables."""
+    engine = create_async_engine(DATABASE_URL, echo=False)
+    async with engine.begin() as conn:
+        await conn.run_sync(MixinBase.metadata.create_all)
+
+    factory = async_sessionmaker(engine, expire_on_commit=False, class_=EventSession)
+
+    try:
+        yield factory
+    finally:
+        async with engine.begin() as conn:
+            await conn.run_sync(MixinBase.metadata.drop_all)
+        await engine.dispose()
+
+
 class TestUUIDMixin:
     @pytest.mark.anyio
     async def test_uuid_generated_by_db(self, mixin_session):
@@ -1559,15 +1578,13 @@ class TestEventSessionWithGetTransaction:
         assert creates[0]["obj_id"] == survivor.id
 
     @pytest.mark.anyio
-    async def test_lock_tables_with_events(self, mixin_session):
-        """Events fire correctly after lock_tables context."""
+    async def test_lock_tables_with_events(self, mixin_session_maker):
+        """Events fire correctly when lock_tables commits on context exit."""
         from fastapi_toolsets.db import lock_tables
 
-        async with lock_tables(mixin_session, [WatchedModel]):
+        async with lock_tables(mixin_session_maker, [WatchedModel]) as session:
             obj = WatchedModel(status="locked", other="x")
-            mixin_session.add(obj)
-
-        await mixin_session.commit()
+            session.add(obj)
 
         creates = [e for e in _test_events if e["event"] == "create"]
         assert len(creates) == 1