[DOP-21799] Refactor celery initialization

Author: Ilyas Gasanov

docker/entrypoint_worker.sh

@@ -6,4 +6,3 @@ set -e
 
 # exec is required to forward all signals to the main process
 exec python -m celery -A syncmaster.worker.celery worker --max-tasks-per-child=1 "$@"
-

syncmaster/backend/__init__.py

@@ -1,5 +1,6 @@
 # SPDX-FileCopyrightText: 2023-2024 MTS PJSC
 # SPDX-License-Identifier: Apache-2.0
+from celery import Celery
 from fastapi import FastAPI, HTTPException
 from fastapi.exceptions import RequestValidationError
 from pydantic import ValidationError
@@ -18,6 +19,7 @@
 from syncmaster.backend.settings import ServerAppSettings as Settings
 from syncmaster.db.factory import create_session_factory, get_uow
 from syncmaster.exceptions import SyncmasterError
+from syncmaster.worker import celery_factory
 
 
 def application_factory(settings: Settings) -> FastAPI:
@@ -30,6 +32,7 @@ def application_factory(settings: Settings) -> FastAPI:
         redoc_url=None,
     )
     application.state.settings = settings
+    application.state.celery = celery_factory(settings)
     application.include_router(api_router)
     application.exception_handler(RequestValidationError)(validation_exception_handler)
     application.exception_handler(ValidationError)(validation_exception_handler)
@@ -44,6 +47,7 @@ def application_factory(settings: Settings) -> FastAPI:
         {
             Settings: lambda: settings,
             UnitOfWork: get_uow(session_factory, settings=settings),
+            Celery: lambda: application.state.celery,
         },
     )
 

syncmaster/backend/api/v1/runs.py

@@ -5,6 +5,7 @@
 from typing import Annotated
 
 from asgi_correlation_id import correlation_id
+from celery import Celery
 from fastapi import APIRouter, Depends, Query
 from jinja2 import Template
 from kombu.exceptions import KombuError
@@ -24,7 +25,6 @@
     ReadRunSchema,
     RunPageSchema,
 )
-from syncmaster.worker import celery
 
 router = APIRouter(tags=["Runs"], responses=get_error_responses())
 
@@ -84,6 +84,7 @@ async def read_run(
 async def start_run(
     create_run_data: CreateRunSchema,
     settings: Annotated[Settings, Depends(Stub(Settings))],
+    celery: Annotated[Celery, Depends(Stub(Celery))],
     unit_of_work: UnitOfWork = Depends(UnitOfWork),
     current_user: User = Depends(get_user(is_active=True)),
 ) -> ReadRunSchema:

syncmaster/scheduler/celery.py

@@ -0,0 +1,6 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS PJSC
+# SPDX-License-Identifier: Apache-2.0
+from syncmaster.scheduler.settings import SchedulerAppSettings
+from syncmaster.worker import celery_factory
+
+celery = celery_factory(SchedulerAppSettings())

syncmaster/scheduler/transfer_job_manager.py

@@ -9,10 +9,10 @@
 from syncmaster.backend.services.unit_of_work import UnitOfWork
 from syncmaster.db.models import RunType, Status, Transfer
 from syncmaster.exceptions.run import CannotConnectToTaskQueueError
+from syncmaster.scheduler.celery import celery
 from syncmaster.scheduler.settings import SchedulerAppSettings as Settings
 from syncmaster.scheduler.utils import get_async_session
 from syncmaster.schemas.v1.connections.connection import ReadAuthDataSchema
-from syncmaster.worker import celery
 
 
 class TransferJobManager:

syncmaster/worker/__init__.py

@@ -3,10 +3,9 @@
 from celery import Celery
 
 from syncmaster.worker.base import WorkerTask
-from syncmaster.worker.settings import WorkerAppSettings
 
 
-def create_celery_app(settings) -> Celery:
+def celery_factory(settings) -> Celery:
     app = Celery(
         __name__,
         broker=settings.broker.url,
@@ -17,8 +16,3 @@ def create_celery_app(settings) -> Celery:
         ],
     )
     return app
-
-
-# TODO: initialize celery app in __name__ == "__main__"
-#       then initialize celery app in backend via dependency injection and initialize in scheduler
-celery = create_celery_app(WorkerAppSettings())

syncmaster/worker/celery.py

@@ -0,0 +1,6 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS PJSC
+# SPDX-License-Identifier: Apache-2.0
+from syncmaster.worker import celery_factory
+from syncmaster.worker.settings import WorkerAppSettings
+
+celery = celery_factory(WorkerAppSettings())

syncmaster/worker/transfer.py

@@ -16,7 +16,7 @@
 from syncmaster.db.models import AuthData, Run, Status, Transfer
 from syncmaster.db.repositories.utils import decrypt_auth_data
 from syncmaster.exceptions.run import RunNotFoundError
-from syncmaster.worker import celery
+from syncmaster.worker.celery import celery
 from syncmaster.worker.controller import TransferController
 from syncmaster.worker.settings import WorkerAppSettings
 

tests/conftest.py

@@ -23,7 +23,7 @@
 from syncmaster.backend.utils.jwt import sign_jwt
 from syncmaster.db.models import Base
 from syncmaster.scheduler.settings import SchedulerAppSettings
-from syncmaster.worker import create_celery_app
+from syncmaster.worker import celery_factory
 from syncmaster.worker.settings import WorkerAppSettings
 from tests.mocks import UserTestRoles
 from tests.settings import TestSettings
@@ -140,7 +140,7 @@ async def client(settings: Settings) -> AsyncGenerator:
 
 @pytest.fixture(scope="session", params=[{}])
 def celery(worker_settings: WorkerAppSettings) -> Celery:
-    celery_app = create_celery_app(worker_settings)
+    celery_app = celery_factory(worker_settings)
     return celery_app
 
 

tests/test_integration/celery_test.py

@@ -1,3 +1,3 @@
-from syncmaster.worker import celery
+from syncmaster.scheduler.celery import celery
 
 celery.conf.update(imports=list(celery.conf.imports) + ["tests.test_integration.test_scheduler.test_task"])