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
-async def lock_tables(
-    session_maker: async_sessionmaker[AsyncSession],
+def lock_tables(
+    session_maker: async_sessionmaker[_SessionT],
     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,6 +189,8 @@ async def lock_tables(
     """
     table_names = ",".join(table.__tablename__ for table in tables)
 
+    @asynccontextmanager
+    async def _lock() -> AsyncGenerator[_SessionT, None]:
         async with session_maker() as session:
             try:
                 await session.execute(text(f"SET LOCAL lock_timeout='{timeout}'"))
@@ -200,6 +201,8 @@ async def lock_tables(
                 await session.rollback()
                 raise
 
+    return _lock()
+
 
 async def create_database(
     db_name: str,