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..d14f6618ab --- /dev/null +++ b/backend/backend/celery_signals.py @@ -0,0 +1,93 @@ +"""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, task_postrun, task_prerun +from log_request_id import local as log_request_id_local +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 + + +def _request_id_from_task(task) -> str | None: + """Read a propagated request_id off a Celery task's message context.""" + request = getattr(task, "request", None) + if request is None: + return None + request_id = getattr(request, Common.REQUEST_ID, None) + if not request_id: + task_headers = getattr(request, "headers", None) + if isinstance(task_headers, dict): + request_id = task_headers.get(Common.REQUEST_ID) + return request_id or None + + +@task_prerun.connect +def bind_request_id(task=None, **kwargs): + """Bind the propagated request_id for tasks executed by the backend's OWN + Celery workers (beat, dashboard-metric tasks, etc.). + + The separate ``workers/`` fleet has its own ``task_prerun`` reader; the + backend Celery app previously injected the header (``propagate_request_id``) + but never consumed it, so backend-executed tasks logged ``request_id:-``. + Binding it onto ``log_request_id``'s thread-local makes + ``log_request_id.filters.RequestIDFilter`` emit it, and onto ``StateStore`` + so any task this worker itself publishes re-propagates it. + """ + request_id = _request_id_from_task(task) + if not request_id: + return + log_request_id_local.request_id = request_id + try: + StateStore.set(Common.REQUEST_ID, request_id) + except Exception: + logger.debug("Unable to set request_id on StateStore", exc_info=True) + + +@task_postrun.connect +def clear_request_id(**kwargs): + """Clear the task-scoped request_id bound in ``bind_request_id``.""" + if hasattr(log_request_id_local, "request_id"): + try: + del log_request_id_local.request_id + except AttributeError: + pass + try: + StateStore.clear(Common.REQUEST_ID) + except Exception: + logger.debug("Unable to clear request_id from StateStore", exc_info=True) diff --git a/backend/backend/settings/base.py b/backend/backend/settings/base.py index 2f28090240..75e2526461 100644 --- a/backend/backend/settings/base.py +++ b/backend/backend/settings/base.py @@ -276,7 +276,11 @@ def get_required_setting(setting_key: str, default: str | None = None) -> str | CORS_ALLOW_ALL_ORIGINS = False # Request ID middleware settings -LOG_REQUEST_ID_HEADER = "X-Request-ID" +# django-log-request-id resolves this via request.META.get(...), where WSGI exposes +# the incoming "X-Request-ID" header as the HTTP_-prefixed key HTTP_X_REQUEST_ID. +# It MUST be the META key, not the raw header name, or the incoming id is never read +# and a fresh one is minted on every request (breaking client/worker correlation). +LOG_REQUEST_ID_HEADER = "HTTP_X_REQUEST_ID" REQUEST_ID_RESPONSE_HEADER = "X-Request-ID" GENERATE_REQUEST_ID_IF_NOT_IN_HEADER = True NO_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/unstract/core/src/unstract/core/flask/middleware.py b/unstract/core/src/unstract/core/flask/middleware.py index 2f31fc0b36..8463426589 100644 --- a/unstract/core/src/unstract/core/flask/middleware.py +++ b/unstract/core/src/unstract/core/flask/middleware.py @@ -13,3 +13,13 @@ def register_request_id_middleware(app: Flask): @app.before_request def assign_request_id(): g.request_id = request.headers.get("X-Request-ID", str(uuid.uuid4())) + + @app.after_request + def echo_request_id(response): + # Echo the id back so a caller that did not supply one can learn the + # value this service minted and correlate its own logs (mirrors the + # Django backend's REQUEST_ID_RESPONSE_HEADER). + request_id = getattr(g, "request_id", None) + if request_id: + response.headers["X-Request-ID"] = request_id + return response 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..18b745cc16 100644 --- a/workers/shared/infrastructure/logging/logger.py +++ b/workers/shared/infrastructure/logging/logger.py @@ -33,6 +33,12 @@ class LogContext: organization_id: str | None = None correlation_id: str | None = None request_id: str | None = None + # True only when request_id came from an upstream message header (a genuine + # cross-service correlation id), not a locally-derived payload id or the + # task_id fallback. Gates worker->worker re-propagation so a fallback id is + # never stamped onto child tasks (which would override their own + # file_execution_id correlation). + request_id_propagatable: bool = False class RequestIDFilter(logging.Filter): @@ -703,22 +709,82 @@ 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. - Catches any extraction failure so a malformed payload can never leave - the previous task's id bound on the thread. + Resolution order: an explicit request_id propagated on the message headers + (genuine cross-service correlation), then a payload-derived id + (``_extract_request_id``), then the Celery ``task_id``. Only the first + (header) source is marked propagatable, so a locally-derived fallback is + never re-stamped onto child tasks. + + The whole resolution runs inside the ``try`` so a malformed payload -- or a + surprising ``task.request`` -- can never raise and leave the previous task's + id bound on the thread. """ + propagatable = False try: - request_id = _extract_request_id(args or (), kwargs or {}, task) or task_id + request_id = _request_id_from_message(task) + if request_id: + propagatable = True + else: + 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 = task_id - WorkerLogger.update_context(request_id=request_id, task_id=task_id) + request_id = None + request_id = request_id or task_id + WorkerLogger.update_context( + request_id=request_id, + task_id=task_id, + request_id_propagatable=propagatable, + ) + + +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 a genuine cross-service correlation id flowing across worker->worker + task chains (e.g. a file-processing task enqueuing a callback). Only + propagates when the current id came from an upstream header + (``request_id_propagatable``) -- never a locally-derived payload id or the + ``task_id`` fallback, which would otherwise override the child task's own + ``file_execution_id`` correlation (e.g. for beat/scheduler-originated + pipelines). 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() + if not ctx or not getattr(ctx, "request_id_propagatable", False): + return + request_id = _coerce_id(getattr(ctx, "request_id", None)) + if request_id: + headers["request_id"] = request_id def _clear_task_context(**_): @@ -728,7 +794,9 @@ def _clear_task_context(**_): ``WorkerLogger.configure()``; only nulls out the per-task fields bound in ``_bind_task_context``. """ - WorkerLogger.update_context(request_id=None, task_id=None) + WorkerLogger.update_context( + request_id=None, task_id=None, request_id_propagatable=False + ) @functools.lru_cache(maxsize=1) @@ -739,13 +807,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/controllers/controller.py b/x2text-service/app/controllers/controller.py index 195cef9682..e1fcd57f6f 100644 --- a/x2text-service/app/controllers/controller.py +++ b/x2text-service/app/controllers/controller.py @@ -15,10 +15,8 @@ from app.util import X2TextUtil basic = Blueprint("basic", __name__) -# Configure the logging format and level -logging.basicConfig( - level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" -) +# Logging is configured centrally in app.logging_util.setup_logging() (called from +# create_app) with the request_id-aware canonical format shared across services. UNSTRUCTURED_URL = "unstructured-url" UNSTRUCTURED_API_KEY = "unstructured-api-key" diff --git a/x2text-service/app/logging_util.py b/x2text-service/app/logging_util.py new file mode 100644 index 0000000000..fd83897bcd --- /dev/null +++ b/x2text-service/app/logging_util.py @@ -0,0 +1,105 @@ +"""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, has_request_context, 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: + # Only touch the request-scoped ``g`` inside an active request context; + # outside one (e.g. gunicorn startup logs) fall back to the placeholder. + record.request_id = ( + getattr(g, "request_id", "-") if has_request_context() 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())) + + @app.after_request + def _echo_request_id(response): + # Echo the id back so a caller that did not supply one can learn the + # value this service minted (mirrors the backend's response header). + request_id = getattr(g, "request_id", None) + if request_id: + response.headers["X-Request-ID"] = request_id + return response