[DOP-19790] Implement scheduler

Author: Ilyas Gasanov

syncmaster/scheduler/__main__.py

@@ -2,15 +2,17 @@
 # SPDX-License-Identifier: Apache-2.0
 import asyncio
 
+from syncmaster.config import Settings
 from syncmaster.scheduler.transfer_fetcher import TransferFetcher
 from syncmaster.scheduler.transfer_job_manager import TransferJobManager
 
 TIMEOUT = 180  # seconds
 
 
 async def main():
-    transfer_fetcher = TransferFetcher()
-    transfer_job_manager = TransferJobManager()
+    settings = Settings()
+    transfer_fetcher = TransferFetcher(settings)
+    transfer_job_manager = TransferJobManager(settings)
     transfer_job_manager.scheduler.start()
 
     while True:

syncmaster/scheduler/transfer_fetcher.py

@@ -1,28 +1,24 @@
 # SPDX-FileCopyrightText: 2023-2024 MTS PJSC
 # SPDX-License-Identifier: Apache-2.0
 from sqlalchemy import select
-from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
-from sqlalchemy.orm import sessionmaker
 
 from syncmaster.config import Settings
 from syncmaster.db.models import Transfer
-
-DATABASE_URL = Settings().build_db_connection_uri()
-engine = create_async_engine(DATABASE_URL)
-AsyncSessionLocal = sessionmaker(bind=engine, class_=AsyncSession, expire_on_commit=False)
+from syncmaster.scheduler.utils import get_async_session
 
 
 class TransferFetcher:
-    def __init__(self):
+    def __init__(self, settings: Settings):
+        self.settings = settings
         self.last_updated_at = None
 
     async def fetch_updated_jobs(self):
 
-        async with AsyncSessionLocal() as session:
-            if self.last_updated_at is None:
-                query = select(Transfer).filter(Transfer.is_scheduled)
-            else:
-                query = select(Transfer).filter(Transfer.updated_at > self.last_updated_at)
+        async with get_async_session(self.settings) as session:
+            query = select(Transfer)
+            if self.last_updated_at is not None:
+                query = query.filter(Transfer.updated_at > self.last_updated_at)
+
             result = await session.execute(query)
             transfers = result.scalars().all()
 

syncmaster/scheduler/transfer_job_manager.py

@@ -10,15 +10,16 @@
 from syncmaster.config import Settings
 from syncmaster.db.models import RunType, Status, Transfer
 from syncmaster.exceptions.run import CannotConnectToTaskQueueError
-from syncmaster.scheduler.transfer_fetcher import AsyncSessionLocal
+from syncmaster.scheduler.transfer_fetcher import get_async_session
 from syncmaster.schemas.v1.connections.connection import ReadAuthDataSchema
 from syncmaster.worker.config import celery
 
 
 class TransferJobManager:
-    def __init__(self):
-        self.scheduler = AsyncIOScheduler(timezone=Settings().TZ)
-        self.scheduler.add_jobstore("sqlalchemy", url=Settings().build_db_connection_uri(driver="psycopg2"))
+    def __init__(self, settings: Settings):
+        self.scheduler = AsyncIOScheduler(timezone=settings.TZ)
+        self.scheduler.add_jobstore("sqlalchemy", url=settings.build_db_connection_uri(driver="psycopg2"))
+        self.settings = settings
 
     def update_jobs(self, transfers: list[Transfer]) -> None:
         for transfer in transfers:
@@ -34,21 +35,21 @@ def update_jobs(self, transfers: list[Transfer]) -> None:
                 self.scheduler.modify_job(
                     job_id=job_id,
                     trigger=CronTrigger.from_crontab(transfer.schedule),
-                    args=(transfer.id,),
+                    args=(transfer.id, self.settings),
                 )
             else:
                 self.scheduler.add_job(
                     func=TransferJobManager.send_job_to_celery,
                     id=job_id,
                     trigger=CronTrigger.from_crontab(transfer.schedule),
-                    args=(transfer.id,),
+                    args=(transfer.id, self.settings),
                 )
 
     @staticmethod
-    async def send_job_to_celery(transfer_id: int) -> None:
+    async def send_job_to_celery(transfer_id: int, settings: Settings) -> None:
 
-        async with AsyncSessionLocal() as session:
-            unit_of_work = UnitOfWork(session=session, settings=Settings())
+        async with get_async_session(settings) as session:
+            unit_of_work = UnitOfWork(session=session, settings=settings)
 
             transfer = await unit_of_work.transfer.read_by_id(transfer_id)
             credentials_source = await unit_of_work.credentials.read(transfer.source_connection_id)

syncmaster/scheduler/utils.py

@@ -0,0 +1,12 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS PJSC
+# SPDX-License-Identifier: Apache-2.0
+from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
+
+from syncmaster.config import Settings
+from syncmaster.db.factory import create_session_factory
+
+
+def get_async_session(settings: Settings) -> AsyncSession:
+    engine = create_async_engine(settings.build_db_connection_uri())
+    session_factory = create_session_factory(engine=engine)
+    return session_factory()