[DOP-23620] Fix scheduler not closing DB connections

Author: dolfinus

syncmaster/scheduler/transfer_job_manager.py

@@ -12,7 +12,7 @@
 from syncmaster.exceptions.transfer import TransferNotFoundError
 from syncmaster.scheduler.celery import app as celery
 from syncmaster.scheduler.settings import SchedulerAppSettings as Settings
-from syncmaster.scheduler.utils import get_async_session
+from syncmaster.scheduler.utils import get_async_engine, get_async_session
 from syncmaster.schemas.v1.connections.connection_base import ReadAuthDataSchema
 from syncmaster.server.services.unit_of_work import UnitOfWork
 
@@ -54,7 +54,7 @@ async def remove_orphan_jobs(self) -> None:
         settings = self.settings
         job_transfer_ids = [int(job.id) for job in all_jobs]
 
-        async with get_async_session(settings) as session:
+        async with get_async_engine(settings) as engine, get_async_session(engine) as session:
             result = await session.execute(
                 select(Transfer).where(Transfer.id == any_(job_transfer_ids)),  # type: ignore[arg-type]
             )
@@ -76,7 +76,7 @@ async def send_job_to_celery(transfer_id: int) -> None:  # noqa: WPS602, WPS217
         """
         settings = Settings()
 
-        async with get_async_session(settings) as session:
+        async with get_async_engine(settings) as engine, get_async_session(engine) as session:
             unit_of_work = UnitOfWork(session=session, settings=settings)
 
             try:

syncmaster/scheduler/utils.py

@@ -1,11 +1,27 @@
 # SPDX-FileCopyrightText: 2023-2024 MTS PJSC
 # SPDX-License-Identifier: Apache-2.0
-from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
+import contextlib
+from collections.abc import AsyncGenerator
+
+from sqlalchemy.ext.asyncio import (
+    AsyncEngine,
+    AsyncSession,
+    async_sessionmaker,
+    create_async_engine,
+)
 
 from syncmaster.scheduler.settings import SchedulerAppSettings as Settings
 
 
-def get_async_session(settings: Settings) -> AsyncSession:
+@contextlib.asynccontextmanager
+async def get_async_engine(settings: Settings) -> AsyncGenerator[AsyncEngine, None]:
     engine = create_async_engine(url=settings.database.sync_url)
+    try:
+        yield engine
+    finally:
+        await engine.dispose()
+
+
+def get_async_session(engine: AsyncEngine) -> AsyncSession:
     session_factory = async_sessionmaker(bind=engine, class_=AsyncSession, expire_on_commit=False)
     return session_factory()