fix: make lock_tables generic over session type (#270)

src/fastapi_toolsets/db.py

@@ -151,14 +151,13 @@ class LockMode(str, Enum):
     ACCESS_EXCLUSIVE = "ACCESS EXCLUSIVE"
 
 
-@asynccontextmanager
+def lock_tables(
-async def lock_tables(
+    session_maker: async_sessionmaker[_SessionT],
-    session_maker: async_sessionmaker[AsyncSession],
     tables: list[type[DeclarativeBase]],
     *,
     mode: LockMode = LockMode.SHARE_UPDATE_EXCLUSIVE,
     timeout: str = "5s",
-) -> AsyncGenerator[AsyncSession, None]:
+) -> AbstractAsyncContextManager[_SessionT]:
     """Lock PostgreSQL tables for the duration of a transaction.
 
     Args:
@@ -190,15 +189,19 @@ async def lock_tables(
     """
     table_names = ",".join(table.__tablename__ for table in tables)
 
-    async with session_maker() as session:
+    @asynccontextmanager
-        try:
+    async def _lock() -> AsyncGenerator[_SessionT, None]:
-            await session.execute(text(f"SET LOCAL lock_timeout='{timeout}'"))
+        async with session_maker() as session:
-            await session.execute(text(f"LOCK {table_names} IN {mode.value} MODE"))
+            try:
-            yield session
+                await session.execute(text(f"SET LOCAL lock_timeout='{timeout}'"))
-            await session.commit()
+                await session.execute(text(f"LOCK {table_names} IN {mode.value} MODE"))
-        except BaseException:
+                yield session
-            await session.rollback()
+                await session.commit()
-            raise
+            except BaseException:
+                await session.rollback()
+                raise
+
+    return _lock()
 
 
 async def create_database(