diff --git a/backend/backend/celery_service.py b/backend/backend/celery_service.py index 9e4f697170..4b0484bb66 100644 --- a/backend/backend/celery_service.py +++ b/backend/backend/celery_service.py @@ -27,4 +27,8 @@ app.config_from_object("backend.celery_config.CeleryConfig") app.autodiscover_tasks() +# Register signal handlers (e.g. request_id propagation onto published tasks). +# Importing the module connects the @before_task_publish handler. +import backend.celery_signals # noqa: E402, F401 + logger.debug(f"Celery Configuration:\n {pformat(app.conf.table(with_defaults=True))}") diff --git a/backend/backend/celery_signals.py b/backend/backend/celery_signals.py new file mode 100644 index 0000000000..cf9a27b801 --- /dev/null +++ b/backend/backend/celery_signals.py @@ -0,0 +1,43 @@ +"""Celery signal handlers for the backend (producer side). + +Propagates the HTTP ``request_id`` (correlation ID assigned by +``CustomRequestIDMiddleware``) onto every published Celery task so that worker +logs can be correlated back to the originating request. + +The value is placed in the task message headers under ``request_id``. Workers +read it from ``task.request`` in ``task_prerun`` and bind it onto their log +context -- see ``workers/shared/infrastructure/logging/logger.py``. Using the +``before_task_publish`` signal means this works for *every* ``send_task`` / +``.delay`` / ``.apply_async`` call with no per-call-site changes. +""" + +import logging + +from account_v2.constants import Common +from celery.signals import before_task_publish +from utils.local_context import StateStore + +logger = logging.getLogger(__name__) + + +@before_task_publish.connect +def propagate_request_id(headers=None, **kwargs): + """Inject the current request_id into the outgoing task's message headers. + + Fires in the producer thread (the web request thread for API-triggered + tasks), where ``StateStore`` still holds the request_id set by + ``CustomRequestIDMiddleware``. No-ops when there is no request_id in scope + (e.g. beat-scheduled publishes), leaving the worker to fall back to its + own correlation id (execution_id / task_id). + """ + if headers is None: + return + try: + request_id = StateStore.get(Common.REQUEST_ID) + except Exception: + # StateStore can raise if CONCURRENCY_MODE is misconfigured; never let + # correlation plumbing break task publishing. + logger.debug("Unable to read request_id from StateStore", exc_info=True) + return + if request_id and not headers.get(Common.REQUEST_ID): + headers[Common.REQUEST_ID] = request_id diff --git a/unstract/core/src/unstract/core/flask/logging.py b/unstract/core/src/unstract/core/flask/logging.py index d131cb92fe..d6848a4a7e 100644 --- a/unstract/core/src/unstract/core/flask/logging.py +++ b/unstract/core/src/unstract/core/flask/logging.py @@ -33,11 +33,15 @@ def setup_logging(log_level: int): "disable_existing_loggers": False, "formatters": { "default": { + # Canonical format shared with the Django backend (``enriched``), + # the workers (``WorkerLogger``) and the x2text-service so a single + # gcloud query parses request_id/trace_id/span_id uniformly. "format": ( "%(levelname)s : [%(asctime)s]" - "{pid:%(process)d tid:%(thread)d request_id:%(request_id)s " - + "trace_id:%(otelTraceID)s span_id:%(otelSpanID)s " - + "%(name)s}:- %(message)s" + "{module:%(module)s process:%(process)d thread:%(thread)d " + "request_id:%(request_id)s " + "trace_id:%(otelTraceID)s span_id:%(otelSpanID)s}" + " :- %(message)s" ), }, }, diff --git a/workers/shared/clients/base_client.py b/workers/shared/clients/base_client.py index 017a1d62ec..33beb79337 100644 --- a/workers/shared/clients/base_client.py +++ b/workers/shared/clients/base_client.py @@ -32,6 +32,20 @@ APPLICATION_JSON = "application/json" +def _current_request_id() -> str | None: + """Return the request_id bound on the current worker log context, if any. + + Bound by the ``task_prerun`` handler in the logging module; used to + propagate ``X-Request-ID`` onto outbound calls to the backend internal API. + Returns ``None`` for the ``"-"`` placeholder so no empty header is sent. + """ + ctx = WorkerLogger.get_context() + request_id = getattr(ctx, "request_id", None) if ctx else None + if not request_id or request_id == "-": + return None + return request_id + + # Single PG-queue rollout flag (same key as pg_queue.flags / executor_rpc). _PG_QUEUE_FLAG_KEY = "pg_queue_enabled" @@ -316,6 +330,12 @@ def _make_request( if current_org_id: headers["X-Organization-ID"] = current_org_id + # Propagate the correlation id back to the backend so worker + # callbacks share the originating request's request_id in logs. + request_id = _current_request_id() + if request_id: + headers["X-Request-ID"] = request_id + if headers: kwargs["headers"] = headers diff --git a/workers/shared/infrastructure/logging/logger.py b/workers/shared/infrastructure/logging/logger.py index c82619acf9..74acaea349 100644 --- a/workers/shared/infrastructure/logging/logger.py +++ b/workers/shared/infrastructure/logging/logger.py @@ -703,24 +703,69 @@ def _extract_request_id( return None +def _request_id_from_message(task: Any) -> str | None: + """Read an explicit request_id propagated via Celery message headers. + + The task producer (backend ``before_task_publish`` handler, or a worker + re-publishing a downstream task) injects ``request_id`` into the message + headers. Celery exposes custom headers on ``task.request`` -- as a direct + attribute under protocol v2, and via the raw ``headers`` mapping as a + version-safe fallback. This is the authoritative cross-service correlation + id and takes precedence over payload-derived ids (file_execution_id, etc.). + """ + request = getattr(task, "request", None) + if request is None: + return None + value = getattr(request, "request_id", None) + if not value: + headers = getattr(request, "headers", None) + if isinstance(headers, Mapping): + value = headers.get("request_id") + return _coerce_id(value) + + def _bind_task_context(task_id, task, args, kwargs, **_): """Celery ``task_prerun`` handler: bind request_id onto the log context. + Resolution order: an explicit request_id propagated on the message headers + (cross-service correlation), then a payload-derived id + (``_extract_request_id``), then the Celery ``task_id``. + Catches any extraction failure so a malformed payload can never leave the previous task's id bound on the thread. """ - try: - request_id = _extract_request_id(args or (), kwargs or {}, task) or task_id - except Exception: - logging.getLogger(__name__).debug( - "request_id extraction failed for task %s; falling back to task_id", - task_id, - exc_info=True, - ) - request_id = task_id + request_id = _request_id_from_message(task) + if not request_id: + try: + request_id = _extract_request_id(args or (), kwargs or {}, task) + except Exception: + logging.getLogger(__name__).debug( + "request_id extraction failed for task %s; falling back to task_id", + task_id, + exc_info=True, + ) + request_id = None + request_id = request_id or task_id WorkerLogger.update_context(request_id=request_id, task_id=task_id) +def _propagate_request_id_on_publish(headers=None, **_): + """Celery ``before_task_publish`` handler (worker side): forward the current + request_id onto tasks this worker publishes. + + Keeps the correlation id flowing across worker->worker task chains (e.g. a + file-processing task enqueuing a callback). Reads the request_id bound onto + the thread-local log context by ``_bind_task_context``; no-ops when absent + or when the caller already set the header. + """ + if headers is None or headers.get("request_id"): + return + ctx = WorkerLogger.get_context() + request_id = _coerce_id(getattr(ctx, "request_id", None)) if ctx else None + if request_id: + headers["request_id"] = request_id + + def _clear_task_context(**_): """Celery ``task_postrun`` handler: reset task-scoped fields only. @@ -739,13 +784,14 @@ def _install_celery_request_id_signals() -> None: debug log if Celery is not importable (e.g. unit tests). """ try: - from celery.signals import task_postrun, task_prerun + from celery.signals import before_task_publish, task_postrun, task_prerun except ImportError as exc: logging.getLogger(__name__).debug( "celery.signals not importable; request_id signal install skipped: %s", exc, ) return + before_task_publish.connect(_propagate_request_id_on_publish, weak=False) task_prerun.connect(_bind_task_context, weak=False) task_postrun.connect(_clear_task_context, weak=False) diff --git a/x2text-service/app/config.py b/x2text-service/app/config.py index 2fca6a6c4d..5823a27c71 100644 --- a/x2text-service/app/config.py +++ b/x2text-service/app/config.py @@ -1,17 +1,25 @@ +import logging from os import environ as env from dotenv import load_dotenv from flask import Flask from app.controllers import api +from app.logging_util import register_request_id_middleware, setup_logging from app.models import X2TextAudit, be_db load_dotenv() def create_app() -> Flask: + log_level = getattr(logging, env.get("LOG_LEVEL", "INFO").upper(), logging.INFO) + setup_logging(log_level) + app = Flask(__name__) + # Assign/propagate a request_id (X-Request-ID) for cross-service log correlation. + register_request_id_middleware(app) + api_url_prefix = env.get("API_URL_PREFIX", "/api/v1") app.register_blueprint(api, url_prefix=api_url_prefix) diff --git a/x2text-service/app/logging_util.py b/x2text-service/app/logging_util.py new file mode 100644 index 0000000000..de9d0a0c66 --- /dev/null +++ b/x2text-service/app/logging_util.py @@ -0,0 +1,92 @@ +"""Request-id-aware logging for the x2text-service. + +Self-contained mirror of the shared ``unstract.core.flask`` logging pattern so +the service participates in cross-service correlation (a single ``request_id`` +in every log line) without taking on the ``unstract-core`` dependency. + +The format string is kept identical to the Django backend and the workers so a +single gcloud query parses ``request_id`` / ``trace_id`` / ``span_id`` uniformly +across every service. +""" + +import logging +import uuid +from logging.config import dictConfig + +from flask import Flask, g, request + +# Canonical log format shared with the Django backend (``enriched``) and the +# workers (``WorkerLogger``). Keep these in sync. +LOG_FORMAT = ( + "%(levelname)s : [%(asctime)s]" + "{module:%(module)s process:%(process)d thread:%(thread)d " + "request_id:%(request_id)s trace_id:%(otelTraceID)s span_id:%(otelSpanID)s}" + " :- %(message)s" +) + + +class RequestIDFilter(logging.Filter): + """Inject the current request's ``request_id`` into log records.""" + + def filter(self, record: logging.LogRecord) -> bool: + record.request_id = getattr(g, "request_id", "-") if g else "-" + return True + + +class OTelFieldFilter(logging.Filter): + """Default OpenTelemetry id fields to ``"-"`` when not populated.""" + + def filter(self, record: logging.LogRecord) -> bool: + for attr in ("otelTraceID", "otelSpanID"): + if not getattr(record, attr, None): + setattr(record, attr, "-") + return True + + +def setup_logging(log_level: int = logging.INFO) -> None: + """Configure root/werkzeug/gunicorn loggers with the standardized format.""" + dictConfig( + { + "version": 1, + "disable_existing_loggers": False, + "formatters": {"default": {"format": LOG_FORMAT}}, + "filters": { + "request_id": {"()": RequestIDFilter}, + "otel_ids": {"()": OTelFieldFilter}, + }, + "handlers": { + "wsgi": { + "class": "logging.StreamHandler", + "stream": "ext://flask.logging.wsgi_errors_stream", + "formatter": "default", + "filters": ["request_id", "otel_ids"], + }, + }, + "loggers": { + "werkzeug": { + "level": log_level, + "handlers": ["wsgi"], + "propagate": False, + }, + "gunicorn.access": { + "level": log_level, + "handlers": ["wsgi"], + "propagate": False, + }, + "gunicorn.error": { + "level": log_level, + "handlers": ["wsgi"], + "propagate": False, + }, + }, + "root": {"level": log_level, "handlers": ["wsgi"]}, + } + ) + + +def register_request_id_middleware(app: Flask) -> None: + """Read ``X-Request-ID`` from each request (or mint one) onto Flask ``g``.""" + + @app.before_request + def _assign_request_id() -> None: + g.request_id = request.headers.get("X-Request-ID", str(uuid.uuid4()))