[DOP-19796] - add log_url before sedning run task to celery worker (#110)

maxim-lixakov · web-flow · commit 890b312e2b21 · 2024-11-01T16:46:27.000+03:00
diff --git a/.env.docker b/.env.docker
@@ -4,10 +4,14 @@ ENV=LOCAL
 # Debug
 SYNCMASTER__SERVER__DEBUG=true
 
-# Logging
+# Logging Backend
 SYNCMASTER__SERVER__LOGGING__SETUP=True
 SYNCMASTER__SERVER__LOGGING__PRESET=colored
 
+# Logging Worker
+SYNCMASTER__WORKER__LOGGING__SETUP=True
+SYNCMASTER__WORKER__LOGGING__PRESET=json
+
 # Postgres
 SYNCMASTER__DATABASE__URL=postgresql+asyncpg://syncmaster:changeme@db:5432/syncmaster
 
diff --git a/.env.local b/.env.local
@@ -8,6 +8,10 @@ export SYNCMASTER__SERVER__DEBUG=true
 export SYNCMASTER__SERVER__LOGGING__SETUP=True
 export SYNCMASTER__SERVER__LOGGING__PRESET=colored
 
+# Logging Worker
+export SYNCMASTER__WORKER__LOGGING__SETUP=True
+export SYNCMASTER__WORKER__LOGGING__PRESET=json
+
 # Postgres
 export SYNCMASTER__DATABASE__URL=postgresql+asyncpg://syncmaster:changeme@localhost:5432/syncmaster
 
diff --git a/docs/backend/configuration/logging.rst b/docs/backend/configuration/logging.rst
@@ -1,7 +1,7 @@
-.. _backend-configuration-logging:
+.. _configuration-logging:
 
 Logging settings
 ================
 
 
-.. autopydantic_model:: syncmaster.settings.server.log.LoggingSettings
+.. autopydantic_model:: syncmaster.settings.log.LoggingSettings
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -98,8 +98,9 @@ worker = [
     "jinja2",
     "psycopg",
     "uuid6",
+    "coloredlogs",
+    "python-json-logger",
 ]
-
 [tool.poetry.group.test.dependencies]
 pandas-stubs = "^2.2.2.240909"
 pytest = "^8.3.3"
diff --git a/syncmaster/backend/api/v1/runs.py b/syncmaster/backend/api/v1/runs.py
@@ -1,10 +1,14 @@
 # SPDX-FileCopyrightText: 2023-2024 MTS PJSC
 # SPDX-License-Identifier: Apache-2.0
 from datetime import datetime
+from typing import Annotated
 
+from asgi_correlation_id import correlation_id
 from fastapi import APIRouter, Depends, Query
+from jinja2 import Template
 from kombu.exceptions import KombuError
 
+from syncmaster.backend.dependencies import Stub
 from syncmaster.backend.services import UnitOfWork, get_user
 from syncmaster.db.models import Status, User
 from syncmaster.db.utils import Permission
@@ -18,6 +22,7 @@
     ReadRunSchema,
     RunPageSchema,
 )
+from syncmaster.settings import Settings
 from syncmaster.worker.config import celery
 
 router = APIRouter(tags=["Runs"], responses=get_error_responses())
@@ -77,6 +82,7 @@ async def read_run(
 @router.post("/runs")
 async def start_run(
     create_run_data: CreateRunSchema,
+    settings: Annotated[Settings, Depends(Stub(Settings))],
     unit_of_work: UnitOfWork = Depends(UnitOfWork),
     current_user: User = Depends(get_user(is_active=True)),
 ) -> ReadRunSchema:
@@ -112,6 +118,15 @@ async def start_run(
             target_creds=ReadAuthDataSchema(auth_data=credentials_target).dict(),
         )
 
+        log_url = Template(settings.worker.LOG_URL_TEMPLATE).render(
+            run=run,
+            correlation_id=correlation_id.get(),
+        )
+        run = await unit_of_work.run.update(
+            run_id=run.id,
+            log_url=log_url,
+        )
+
     try:
         celery.send_task(
             "run_transfer_task",
diff --git a/syncmaster/settings/__init__.py b/syncmaster/settings/__init__.py
@@ -9,6 +9,7 @@
 from syncmaster.settings.broker import RabbitMQSettings
 from syncmaster.settings.database import DatabaseSettings
 from syncmaster.settings.server import ServerSettings
+from syncmaster.settings.worker import WorkerSettings
 
 
 class EnvTypes(StrEnum):
@@ -48,16 +49,18 @@ class Settings(BaseSettings):
     SECURITY_ALGORITHM: str = "HS256"
     CRYPTO_KEY: str = "UBgPTioFrtH2unlC4XFDiGf5sYfzbdSf_VgiUSaQc94="
 
-    CORRELATION_CELERY_HEADER_ID: str = "CORRELATION_CELERY_HEADER_ID"
-
     TOKEN_EXPIRED_TIME: int = 60 * 60 * 10  # 10 hours
     CREATE_SPARK_SESSION_FUNCTION: ImportString = "syncmaster.worker.spark.get_worker_spark_session"
 
     database: DatabaseSettings = Field(description=":ref:`Database settings <backend-configuration-database>`")
     broker: RabbitMQSettings = Field(description=":ref:`Broker settings <backend-configuration-broker>`")
     server: ServerSettings = Field(
         default_factory=ServerSettings,
-        description=":ref:`Server settings <backend-configuration>`",
+        description="Server settings <backend-configuration",
+    )
+    worker: WorkerSettings = Field(
+        default_factory=WorkerSettings,
+        description="Celery worker settings",
     )
 
     class Config:
diff --git a/syncmaster/settings/log/__init__.py b/syncmaster/settings/log/__init__.py
@@ -20,16 +20,26 @@ class LoggingSettings(BaseModel):
 
     .. code-block:: bash
 
+        # backend logging
         SYNCMASTER__SERVER__LOGGING__SETUP=True
         SYNCMASTER__SERVER__LOGGING__PRESET=json
 
+        # celery logging
+        SYNCMASTER__WORKER__LOGGING__SETUP=True
+        SYNCMASTER__WORKER__LOGGING__PRESET=json
+
     Passing custom logging config file:
 
     .. code-block:: bash
 
+        # backend logging
         SYNCMASTER__SERVER__LOGGING__SETUP=True
         SYNCMASTER__SERVER__LOGGING__CUSTOM__CONFIG_PATH=/some/logging.yml
 
+        # celery logging
+        SYNCMASTER__WORKER__LOGGING__SETUP=True
+        SYNCMASTER__WORKER__LOGGING__CUSTOM__CONFIG_PATH=/some/logging.yml
+
     Setup logging in some other way, e.g. using `uvicorn args <https://www.uvicorn.org/settings/#logging>`_:
 
     .. code-block:: bash
diff --git a/syncmaster/settings/log/colored.yml b/syncmaster/settings/log/colored.yml
@@ -24,6 +24,11 @@ handlers:
     formatter: colored
     filters: [correlation_id]
     stream: ext://sys.stdout
+  celery:
+    class: logging.StreamHandler
+    formatter: colored
+    filters: [correlation_id]
+    stream: ext://sys.stdout
 
 loggers:
   '':
@@ -34,3 +39,7 @@ loggers:
     handlers: [main]
     level: INFO
     propagate: false
+  celery:
+    level: DEBUG
+    handlers: [celery]
+    propagate: false
diff --git a/syncmaster/settings/log/json.yml b/syncmaster/settings/log/json.yml
@@ -23,6 +23,11 @@ handlers:
     formatter: json
     filters: [correlation_id]
     stream: ext://sys.stdout
+  celery:
+    class: logging.StreamHandler
+    formatter: json
+    filters: [correlation_id]
+    stream: ext://sys.stdout
 
 loggers:
   '':
@@ -33,3 +38,7 @@ loggers:
     handlers: [main]
     level: INFO
     propagate: false
+  celery:
+    level: DEBUG
+    handlers: [celery]
+    propagate: false
diff --git a/syncmaster/settings/log/plain.yml b/syncmaster/settings/log/plain.yml
@@ -24,6 +24,11 @@ handlers:
     formatter: plain
     filters: [correlation_id]
     stream: ext://sys.stdout
+  celery:
+    class: logging.StreamHandler
+    formatter: plain
+    filters: [correlation_id]
+    stream: ext://sys.stdout
 
 loggers:
   '':
@@ -34,3 +39,7 @@ loggers:
     handlers: [main]
     level: INFO
     propagate: false
+    celery:
+      level: DEBUG
+      handlers: [celery]
+      propagate: false
diff --git a/syncmaster/settings/server/__init__.py b/syncmaster/settings/server/__init__.py
@@ -5,8 +5,8 @@
 
 from pydantic import BaseModel, Field
 
+from syncmaster.settings.log import LoggingSettings
 from syncmaster.settings.server.cors import CORSSettings
-from syncmaster.settings.server.log import LoggingSettings
 from syncmaster.settings.server.request_id import RequestIDSettings
 
 
diff --git a/syncmaster/settings/worker/__init__.py b/syncmaster/settings/worker/__init__.py
@@ -0,0 +1,24 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS PJSC
+# SPDX-License-Identifier: Apache-2.0
+from pydantic import BaseModel, Field
+
+from syncmaster.settings.log import LoggingSettings
+
+
+class WorkerSettings(BaseModel):
+    """Celery worker settings.
+
+    Examples
+    --------
+
+    .. code-block:: bash
+
+        SYNCMASTER__WORKER__LOGGING__PRESET=colored
+    """
+
+    logging: LoggingSettings = Field(
+        default_factory=LoggingSettings,
+        description=":ref:`Logging settings <configuration-logging>`",
+    )
+    CORRELATION_CELERY_HEADER_ID: str = "CORRELATION_CELERY_HEADER_ID"
+    LOG_URL_TEMPLATE: str = ""
diff --git a/syncmaster/worker/transfer.py b/syncmaster/worker/transfer.py
@@ -5,10 +5,12 @@
 
 import onetl
 from asgi_correlation_id import correlation_id
-from celery.signals import before_task_publish, task_prerun
+from celery.signals import after_setup_logger, before_task_publish, task_prerun
+from celery.utils.log import get_task_logger
 from sqlalchemy import select
 from sqlalchemy.orm import Session, selectinload
 
+from syncmaster.backend.middlewares.logging import setup_logging
 from syncmaster.db.models import AuthData, Run, Status, Transfer
 from syncmaster.db.repositories.utils import decrypt_auth_data
 from syncmaster.exceptions.run import RunNotFoundError
@@ -17,7 +19,7 @@
 from syncmaster.worker.config import celery
 from syncmaster.worker.controller import TransferController
 
-logger = logging.getLogger(__name__)
+logger = get_task_logger(__name__)
 
 # TODO: configure correlation id from settings: settings.CORRELATION_CELERY_HEADER_ID
 CORRELATION_CELERY_HEADER_ID = "CORRELATION_CELERY_HEADER_ID"
@@ -82,6 +84,12 @@ def run_transfer(session: Session, run_id: int, settings: Settings):
     session.commit()
 
 
+@after_setup_logger.connect
+def setup_loggers(*args, **kwargs):
+    # TODO: remove calling Settings
+    setup_logging(Settings().worker.logging.get_log_config_path())
+
+
 @before_task_publish.connect()
 def transfer_correlation_id(headers, *args, **kwargs) -> None:
     # This is called before task.delay() finishes
diff --git a/tests/test_unit/test_runs/test_create_run.py b/tests/test_unit/test_runs/test_create_run.py
@@ -4,6 +4,7 @@
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from syncmaster.db.models import Run, Status
+from syncmaster.settings import Settings
 from tests.mocks import MockGroup, MockTransfer, MockUser, UserTestRoles
 
 pytestmark = [pytest.mark.asyncio, pytest.mark.backend]
@@ -123,9 +124,13 @@ async def test_superuser_can_create_run(
     superuser: MockUser,
     group_transfer: MockTransfer,
     session: AsyncSession,
+    settings: Settings,
     mocker,
-    monkeypatch,
 ) -> None:
+    # Arrange
+    settings.worker.LOG_URL_TEMPLATE = (
+        "https://grafana.example.com?correlation_id={{ correlation_id }}&run_id={{ run.id }}"
+    )
     mocker.patch("syncmaster.worker.config.celery.send_task")
 
     # Act
@@ -141,7 +146,8 @@ async def test_superuser_can_create_run(
     ).first()
 
     # Assert
-    assert result.json() == {
+    response = result.json()
+    assert response == {
         "id": run.id,
         "transfer_id": run.transfer_id,
         "status": run.status.value,
@@ -151,6 +157,8 @@ async def test_superuser_can_create_run(
         "transfer_dump": run.transfer_dump,
     }
     assert result.status_code == 200
+    assert "correlation_id" in response.get("log_url")
+    assert "run_id" in response.get("log_url")
 
 
 async def test_unauthorized_user_cannot_create_run(
