[DOP-19992] - create celery app with provided settings

Author: maxim-lixakov

syncmaster/backend/api/v1/runs.py

@@ -21,8 +21,8 @@
     ReadRunSchema,
     RunPageSchema,
 )
-from syncmaster.worker.config import celery
-from syncmaster.worker.settings import get_worker_settings
+from syncmaster.worker.settings import WorkerAppSettings
+from syncmaster.worker.tasks import celery
 
 router = APIRouter(tags=["Runs"], responses=get_error_responses())
 
@@ -117,7 +117,7 @@ async def start_run(
             type=RunType.MANUAL,
         )
 
-        log_url = Template(get_worker_settings().worker.LOG_URL_TEMPLATE).render(
+        log_url = Template(WorkerAppSettings().worker.LOG_URL_TEMPLATE).render(
             run=run,
             correlation_id=correlation_id.get(),
         )

syncmaster/scheduler/transfer_fetcher.py

@@ -3,7 +3,7 @@
 from sqlalchemy import select
 
 from syncmaster.db.models import Transfer
-from syncmaster.scheduler.settings import SchedulerSettings as Settings
+from syncmaster.scheduler.settings import SchedulerAppSettings as Settings
 from syncmaster.scheduler.utils import get_async_session
 
 

syncmaster/scheduler/transfer_job_manager.py

@@ -12,7 +12,7 @@
 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.config import celery
+from syncmaster.worker.tasks import celery
 
 
 class TransferJobManager:

syncmaster/worker/__init__.py

@@ -1,2 +1,18 @@
 # SPDX-FileCopyrightText: 2023-2024 MTS PJSC
 # SPDX-License-Identifier: Apache-2.0
+from celery import Celery
+
+from syncmaster.worker.base import WorkerTask
+
+
+def create_celery_app(settings) -> Celery:
+    app = Celery(
+        __name__,
+        broker=settings.broker.url,
+        backend="db+" + settings.database.sync_url,
+        task_cls=WorkerTask(settings=settings),
+        imports=[
+            "syncmaster.worker.transfer",
+        ],
+    )
+    return app

syncmaster/worker/base.py

@@ -3,12 +3,12 @@
 from celery import Task
 from sqlalchemy import create_engine
 
-from syncmaster.worker.settings import get_worker_settings
+from syncmaster.worker.settings import WorkerAppSettings
 
 
 class WorkerTask(Task):
-    def __init__(self) -> None:
-        self.settings = get_worker_settings()
+    def __init__(self, settings: WorkerAppSettings) -> None:
+        self.settings = settings
         self.engine = create_engine(
             url=self.settings.database.sync_url,
         )

syncmaster/worker/config.py

@@ -1,8 +1,5 @@
 # SPDX-FileCopyrightText: 2023-2024 MTS PJSC
 # SPDX-License-Identifier: Apache-2.0
-import importlib
-import os
-
 from pydantic import Field
 from pydantic.types import ImportString
 from pydantic_settings import BaseSettings
@@ -54,16 +51,3 @@ class WorkerAppSettings(BaseSettings):
     class Config:
         env_prefix = "SYNCMASTER__"
         env_nested_delimiter = "__"
-
-
-def get_worker_settings() -> WorkerAppSettings:
-    # TODO: add to worker documentation
-    worker_settings_path = os.environ.get("WORKER_SETTINGS", None)
-
-    if worker_settings_path:
-        module_name, class_name = worker_settings_path.rsplit(".", 1)
-        module = importlib.import_module(module_name)
-        settings_class = getattr(module, class_name)
-    else:
-        settings_class = WorkerAppSettings
-    return settings_class()

syncmaster/worker/settings/__init__.py

@@ -0,0 +1,13 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS PJSC
+# SPDX-License-Identifier: Apache-2.0
+from syncmaster.worker import create_celery_app
+from syncmaster.worker.settings import WorkerAppSettings
+from syncmaster.worker.transfer import run_transfer_task
+
+
+def register_tasks(celery_app):
+    celery_app.task(name="run_transfer_task", bind=True, track_started=True)(run_transfer_task)
+
+
+celery = create_celery_app(WorkerAppSettings())
+register_tasks(celery)

syncmaster/worker/tasks.py

@@ -16,16 +16,14 @@
 from syncmaster.db.repositories.utils import decrypt_auth_data
 from syncmaster.exceptions.run import RunNotFoundError
 from syncmaster.worker.base import WorkerTask
-from syncmaster.worker.config import celery
 from syncmaster.worker.controller import TransferController
-from syncmaster.worker.settings import get_worker_settings
+from syncmaster.worker.settings import WorkerAppSettings
 
 logger = get_task_logger(__name__)
 
-CORRELATION_CELERY_HEADER_ID = get_worker_settings().worker.CORRELATION_CELERY_HEADER_ID
+CORRELATION_CELERY_HEADER_ID = WorkerAppSettings().worker.CORRELATION_CELERY_HEADER_ID
 
 
-@celery.task(name="run_transfer_task", bind=True, track_started=True)
 def run_transfer_task(self: WorkerTask, run_id: int) -> None:
     onetl.log.setup_logging(level=logging.INFO)
     with Session(self.engine) as session:

syncmaster/worker/transfer.py

@@ -8,6 +8,7 @@
 import pytest
 import pytest_asyncio
 from alembic.config import Config as AlembicConfig
+from celery import Celery
 from httpx import AsyncClient
 from sqlalchemy.ext.asyncio import (
     AsyncEngine,
@@ -22,6 +23,8 @@
 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.settings import WorkerAppSettings
 from tests.mocks import UserTestRoles
 from tests.settings import TestSettings
 from tests.utils import prepare_new_database, run_async_migrations
@@ -77,6 +80,11 @@ def scheduler_settings(request: pytest.FixtureRequest) -> SchedulerAppSettings:
     return SchedulerAppSettings.parse_obj(request.param)
 
 
+@pytest.fixture(scope="session", params=[{}])
+def worker_settings(request: pytest.FixtureRequest) -> WorkerAppSettings:
+    return WorkerAppSettings.parse_obj(request.param)
+
+
 @pytest.fixture(scope="session")
 def test_settings():
     return TestSettings()
@@ -130,6 +138,12 @@ async def client(settings: Settings) -> AsyncGenerator:
         logger.info("END CLIENT FIXTURE")
 
 
+@pytest.fixture(scope="session", params=[{}])
+def celery(worker_settings: WorkerAppSettings) -> Celery:
+    celery_app = create_celery_app(worker_settings)
+    return celery_app
+
+
 @pytest_asyncio.fixture
 async def create_connection_data(request):
     if hasattr(request, "param"):