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..ba2a76aa6a --- /dev/null +++ b/backend/backend/celery_signals.py @@ -0,0 +1,82 @@ +"""Celery signal handlers carrying the HTTP ``request_id`` onto published tasks. + +The id travels in the task message headers under ``Common.REQUEST_ID``. Hooking +``before_task_publish`` rather than each producer is deliberate: it covers every +``send_task`` / ``.delay`` / ``.apply_async`` with no per-call-site change. +""" + +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. + + Relies on firing in the *producer* thread, where ``StateStore`` still holds + the id set by ``CustomRequestIDMiddleware``. No-ops without one (e.g. beat + publishes), leaving the worker to derive its own correlation 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: + 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 run by the backend's own workers. + + Two writes, both required: ``log_request_id``'s thread-local is the only + thing ``log_request_id.filters.RequestIDFilter`` reads (the HTTP middleware + is otherwise its sole writer, so a task would log ``request_id:-``), and + ``StateStore`` is what ``propagate_request_id`` reads, so tasks this worker + itself publishes carry the id onward. + """ + 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 -- worker threads are pooled and reused.""" + 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..439a8e9c25 100644 --- a/backend/backend/settings/base.py +++ b/backend/backend/settings/base.py @@ -276,10 +276,18 @@ 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 = "-" +# The frontend is a separate origin, so the browser hides every response header +# not named here -- without this the id the backend logged is unreadable in JS +# and the error toast falls back to showing an id that appears in no log. +CORS_EXPOSE_HEADERS = [REQUEST_ID_RESPONSE_HEADER] class OTelFieldFilter(logging.Filter): diff --git a/backend/backend/test_celery_signals.py b/backend/backend/test_celery_signals.py new file mode 100644 index 0000000000..9d426ee549 --- /dev/null +++ b/backend/backend/test_celery_signals.py @@ -0,0 +1,161 @@ +"""Unit checks for the backend's Celery request_id signal handlers. + +These are the producer half of the correlation chain: the HTTP request binds an +id (see ``middleware/test_request_id.py``), and these handlers carry it onto +every task the backend publishes and back off the message for tasks the +backend's *own* Celery workers run. + +The handlers are written to fail open -- correlation plumbing must never break +task publishing -- which means a regression here is silent by construction: the +header simply stops being set and every worker log line reverts to +``request_id:-``. Nothing raises, so only assertions catch it. + +Pure-logic: ``StateStore`` is a thread-local and the signals are called +directly, so no broker, worker or DB is involved. +""" + +import pytest +from account_v2.constants import Common +from log_request_id import local as log_request_id_local +from utils.local_context import StateStore + +from backend.celery_signals import ( + bind_request_id, + clear_request_id, + propagate_request_id, +) + +REQUEST_ID = "6f1c2d3e-4a5b-4c6d-8e9f-0a1b2c3d4e5f" + + +class _Request: + """Stands in for a Celery task ``Context``.""" + + def __init__(self, request_id=None, headers=None): + if request_id is not None: + self.request_id = request_id + self.headers = headers + + +class _Task: + def __init__(self, request=None): + self.request = request + + +@pytest.fixture(autouse=True) +def _clean_context(): + """Both stores are thread-local and shared across tests in a worker, so a + leaked id would make a later test pass for the wrong reason. + """ + yield + for store_key in (Common.REQUEST_ID,): + try: + StateStore.clear(store_key) + except AttributeError: + pass + if hasattr(log_request_id_local, "request_id"): + del log_request_id_local.request_id + + +# Producer side: before_task_publish + + +def test_publish_injects_request_id_from_state_store(): + """The id bound by the HTTP middleware rides out on the task message.""" + StateStore.set(Common.REQUEST_ID, REQUEST_ID) + headers = {} + + propagate_request_id(headers=headers) + + assert headers[Common.REQUEST_ID] == REQUEST_ID + + +def test_publish_is_a_noop_without_a_request_id(): + """Beat-scheduled publishes have no request in scope; the worker falls back + to its own execution_id/task_id rather than receiving an empty header. + """ + headers = {} + + propagate_request_id(headers=headers) + + assert headers == {} + + +def test_publish_does_not_clobber_an_explicit_header(): + """A caller that set the header deliberately outranks the ambient id.""" + StateStore.set(Common.REQUEST_ID, REQUEST_ID) + headers = {Common.REQUEST_ID: "caller-supplied"} + + propagate_request_id(headers=headers) + + assert headers[Common.REQUEST_ID] == "caller-supplied" + + +def test_publish_tolerates_missing_headers(): + """``before_task_publish`` must never raise -- it would break the publish + itself, taking the actual task down with the correlation plumbing. + """ + propagate_request_id(headers=None) # does not raise + + +# Consumer side: task_prerun / task_postrun on the backend's own workers + + +def test_prerun_binds_id_from_task_attribute(): + """Celery surfaces the custom header as an attribute on the Context.""" + bind_request_id(task=_Task(_Request(request_id=REQUEST_ID))) + + assert log_request_id_local.request_id == REQUEST_ID + assert StateStore.get(Common.REQUEST_ID) == REQUEST_ID + + +def test_prerun_falls_back_to_raw_headers_mapping(): + """Fallback for when Celery leaves the header only in the raw mapping.""" + bind_request_id(task=_Task(_Request(headers={Common.REQUEST_ID: REQUEST_ID}))) + + assert log_request_id_local.request_id == REQUEST_ID + + +def test_prerun_binding_makes_the_id_re_propagate(): + """Second-order effect that motivated the handler: a task the backend + worker itself publishes must carry the inherited id onward, or the chain + dies at the first backend-worker hop. + """ + bind_request_id(task=_Task(_Request(request_id=REQUEST_ID))) + headers = {} + + propagate_request_id(headers=headers) + + assert headers[Common.REQUEST_ID] == REQUEST_ID + + +def test_prerun_without_an_id_leaves_stores_untouched(): + """No header means no correlation to inherit -- and, importantly, no empty + string bound over whatever the logger would otherwise show. + """ + bind_request_id(task=_Task(_Request())) + + assert not hasattr(log_request_id_local, "request_id") + assert StateStore.get(Common.REQUEST_ID) is None + + +def test_prerun_tolerates_a_task_without_a_request(): + bind_request_id(task=_Task(request=None)) # does not raise + + assert not hasattr(log_request_id_local, "request_id") + + +def test_postrun_clears_the_bound_id(): + """Celery reuses worker threads; a surviving id would mislabel the *next* + task's logs with the previous task's correlation id. + """ + bind_request_id(task=_Task(_Request(request_id=REQUEST_ID))) + + clear_request_id() + + assert not hasattr(log_request_id_local, "request_id") + assert StateStore.get(Common.REQUEST_ID) is None + + +def test_postrun_is_safe_when_nothing_was_bound(): + clear_request_id() # does not raise diff --git a/backend/middleware/request_id.py b/backend/middleware/request_id.py index c7d0bdb82a..197fc3f19c 100644 --- a/backend/middleware/request_id.py +++ b/backend/middleware/request_id.py @@ -2,7 +2,40 @@ from log_request_id.middleware import RequestIDMiddleware +# The internal service boundary: our own workers call back here while executing +# a workflow, forwarding the request_id of the HTTP call that started it. +INTERNAL_PATH_PREFIX = "/internal/" + + +def _canonical_uuid(value: str | None) -> str | None: + """Return ``value`` only if it is a canonical hyphenated UUID.""" + try: + return value if str(uuid.UUID(value)) == value else None + except (AttributeError, TypeError, ValueError): + return None + class CustomRequestIDMiddleware(RequestIDMiddleware): + """Provisions the request id here rather than trusting the caller's. + + Adopting a caller-supplied id lets any client repeat one value across + unrelated requests. That does not merely lose correlation -- a log query for + that id returns other tenants' requests and reads as though it worked, which + is worse than returning nothing. The id also reaches every log line, Celery + message header and outbound internal-API call, so an unvalidated one can + forge a log record outright. + + The internal boundary is the exception: there the caller is our own worker + forwarding the id of the request that started the execution, which is the + hop that makes worker logs correlate back to the originating call. + """ + + def _get_request_id(self, request): + if request.path.startswith(INTERNAL_PATH_PREFIX): + forwarded = _canonical_uuid(request.META.get(self.request_id_header)) + if forwarded: + return forwarded + return self._generate_id() + def _generate_id(self): return str(uuid.uuid4()) diff --git a/backend/middleware/test_request_id.py b/backend/middleware/test_request_id.py new file mode 100644 index 0000000000..596f3a75a2 --- /dev/null +++ b/backend/middleware/test_request_id.py @@ -0,0 +1,203 @@ +"""Request-level tests for ``X-Request-ID`` provisioning at the HTTP boundary. + +Two properties are pinned here, and they pull in opposite directions. + +*The id is server-provisioned.* A caller-supplied id is ignored on public +routes. Honouring one would let a client send the same value on every request, +and a log query for it would then return unrelated requests across tenants -- +which reads as though correlation worked. The id also reaches every log line and +every published Celery message, so an unvalidated one can forge a log record. + +*Except across the internal boundary,* where the caller is our own worker +forwarding the id of the request that started the execution. That hop is what +makes worker logs correlate back to the originating call, and it hinges on a +single setting with a silent failure mode: when ``LOG_REQUEST_ID_HEADER`` is the +raw header name rather than the WSGI ``META`` key, ``django-log-request-id`` +looks up a key that never exists and mints a fresh id instead. Nothing errors -- +every service just logs a different id. Settings are deliberately *not* +overridden here, so reverting the production value fails these tests. + +No DB is touched (``SimpleTestCase`` + a bare middleware list), keeping this in +the fast unit tier. +""" + +import re +import uuid + +from django.http import HttpResponse +from django.test import SimpleTestCase, override_settings +from django.urls import path + +# Set by CustomRequestIDMiddleware; echoed back so the view can assert on it. +REQUEST_ID_ECHO_HEADER = "X-Seen-Request-Id" + + +def _echo_view(request): + """Reports the request id the middleware bound, so the test can compare it + against both the id sent in and the id echoed on the response. + """ + response = HttpResponse("ok") + response[REQUEST_ID_ECHO_HEADER] = getattr(request, "id", "") + return response + + +PUBLIC_URL = "/echo/" +INTERNAL_URL = "/internal/echo/" + +urlpatterns = [ + path("echo/", _echo_view), + path("internal/echo/", _echo_view), +] + +# Just the middleware under test: no auth, tenancy or session, so nothing here +# reaches the database. +_MIDDLEWARE = ["middleware.request_id.CustomRequestIDMiddleware"] + +_UUID4_RE = re.compile( + r"\A[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}\Z" +) + + +@override_settings(ROOT_URLCONF=__name__, MIDDLEWARE=_MIDDLEWARE) +class PublicRequestIDTest(SimpleTestCase): + """Public routes: the id is ours, whatever the caller sent.""" + + def test_caller_supplied_id_is_ignored(self): + """The frontend attaches a uuid4 of its own to every call. It is still + the backend's id that is authoritative -- the frontend reads the value + back off the response header, so nothing is lost by ignoring it. + """ + sent = "11111111-2222-4333-8444-555555555555" + + response = self.client.get(PUBLIC_URL, headers={"x-request-id": sent}) + + bound = response[REQUEST_ID_ECHO_HEADER] + self.assertNotEqual(bound, sent) + self.assertRegex(bound, _UUID4_RE) + + def test_repeated_caller_id_does_not_collapse_requests(self): + """The reason a caller's id is not adopted: one repeated across requests + would make a single log query return unrelated requests, and read as + though it had worked. + """ + sent = "aaaaaaaa-bbbb-4ccc-8ddd-eeeeeeeeeeee" + headers = {"x-request-id": sent} + + first = self.client.get(PUBLIC_URL, headers=headers) + second = self.client.get(PUBLIC_URL, headers=headers) + + self.assertNotEqual(first[REQUEST_ID_ECHO_HEADER], second[REQUEST_ID_ECHO_HEADER]) + + def test_hostile_id_never_reaches_a_log_line(self): + """An id lands unescaped in every log line, so one carrying terminal + control codes could erase the real prefix of a record and forge the + rest. gunicorn rejects only NUL/CR/LF, so the escape arrives intact. + """ + sent = "\x1b[2K\x1b[1000Ddeadbeef} :- SPOOFED: admin deleted org 42" + + response = self.client.get(PUBLIC_URL, headers={"x-request-id": sent}) + + self.assertRegex(response[REQUEST_ID_ECHO_HEADER], _UUID4_RE) + + def test_response_echoes_the_provisioned_id(self): + """The response carries the id back, which is how the frontend surfaces + it on an error and how anyone debugging a live call retrieves it. + """ + response = self.client.get(PUBLIC_URL) + + echoed = response["X-Request-ID"] + self.assertRegex(echoed, _UUID4_RE) + self.assertEqual(response[REQUEST_ID_ECHO_HEADER], echoed) + + def test_distinct_callers_get_distinct_ids(self): + """Two requests must not share an id, or correlation collapses instead + of merely missing. + """ + first = self.client.get(PUBLIC_URL)["X-Request-ID"] + second = self.client.get(PUBLIC_URL)["X-Request-ID"] + + self.assertNotEqual(first, second) + + def test_uuid_module_still_backs_the_generator(self): + """Guards the custom ``_generate_id`` override, which exists so ids are + plain uuid4 strings rather than the library's default hex form. + """ + response = self.client.get(PUBLIC_URL) + + parsed = uuid.UUID(response["X-Request-ID"]) + self.assertEqual(str(parsed), response["X-Request-ID"]) + + +@override_settings(ROOT_URLCONF=__name__, MIDDLEWARE=_MIDDLEWARE) +class InternalRequestIDTest(SimpleTestCase): + """The internal boundary: a worker's forwarded id is honoured.""" + + def test_forwarded_id_is_adopted(self): + """The hop that silently broke: workers send the originating request's + id on their callbacks, and before the fix the backend discarded it. + """ + sent = "11111111-2222-4333-8444-555555555555" + + response = self.client.get(INTERNAL_URL, headers={"x-request-id": sent}) + + self.assertEqual( + response[REQUEST_ID_ECHO_HEADER], + sent, + "backend minted a new id instead of adopting the worker's -- " + "LOG_REQUEST_ID_HEADER must be the WSGI META key HTTP_X_REQUEST_ID, " + "not the raw header name", + ) + self.assertEqual(response["X-Request-ID"], sent) + + def test_forwarded_id_must_be_a_uuid(self): + """The boundary is authenticated downstream, not here, so an unauthorised + caller can still reach this path. Only the shape our own services emit is + accepted, which bounds both length and character set. + """ + response = self.client.get(INTERNAL_URL, headers={"x-request-id": "edge-lb-7f"}) + + self.assertRegex(response[REQUEST_ID_ECHO_HEADER], _UUID4_RE) + + def test_overlong_forwarded_id_is_rejected(self): + """An id is re-stamped onto every published Celery message and every log + line of an execution, so an unbounded one amplifies: gunicorn accepts + ~8KB, which one call can fan out across N file tasks. + """ + response = self.client.get(INTERNAL_URL, headers={"x-request-id": "A" * 8000}) + + self.assertRegex(response[REQUEST_ID_ECHO_HEADER], _UUID4_RE) + + def test_id_is_provisioned_when_nothing_is_forwarded(self): + response = self.client.get(INTERNAL_URL) + + self.assertRegex(response[REQUEST_ID_ECHO_HEADER], _UUID4_RE) + + +@override_settings( + ROOT_URLCONF=__name__, + MIDDLEWARE=["corsheaders.middleware.CorsMiddleware"] + _MIDDLEWARE, +) +class RequestIDIsReadableByTheBrowserTest(SimpleTestCase): + """The frontend is a separate origin, so the echo only reaches JS if the + header is named in ``CORS_EXPOSE_HEADERS``. + + Without it the browser hides the header, ``getRequestIdFromError`` falls + through to the id the interceptor *sent*, and -- now that the backend + provisions its own -- the error toast shows an id that appears in no log. + Nothing errors; the id is simply wrong, which is why this is pinned. + """ + + def test_request_id_is_exposed_to_a_cross_origin_caller(self): + origin = "http://localhost:3000" + + response = self.client.get( + PUBLIC_URL, headers={"origin": origin, "x-request-id": "ignored"} + ) + + exposed = response.get("Access-Control-Expose-Headers", "") + self.assertIn( + "X-Request-ID", + [h.strip() for h in exposed.split(",")], + "the browser cannot read the id the backend logged", + ) + self.assertRegex(response["X-Request-ID"], _UUID4_RE) diff --git a/unstract/core/src/unstract/core/flask/logging.py b/unstract/core/src/unstract/core/flask/logging.py index d131cb92fe..7b588aa223 100644 --- a/unstract/core/src/unstract/core/flask/logging.py +++ b/unstract/core/src/unstract/core/flask/logging.py @@ -33,11 +33,17 @@ def setup_logging(log_level: int): "disable_existing_loggers": False, "formatters": { "default": { + # Canonical cross-service log format -- this module owns it. + # A single gcloud query parses request_id/trace_id/span_id + # across every service only while the copies agree; the known + # copies are the Django backend ``enriched`` formatter, the + # workers' ``WorkerLogger``, and ``x2text-service``. "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..dbde819c67 100644 --- a/unstract/core/src/unstract/core/flask/middleware.py +++ b/unstract/core/src/unstract/core/flask/middleware.py @@ -3,6 +3,25 @@ from flask import Flask, g, request +def _incoming_request_id() -> str: + """Adopt the caller's id, or mint one. + + Unlike the Django backend these services are not internet-facing: every + caller is another Unstract service forwarding an id, so honouring it is the + whole point. The id still lands unescaped in every log line and in the echoed + response header, so only the canonical UUID our services emit is accepted -- + which bounds length and rules out the control characters that would let a + caller forge a log record. + """ + request_id = request.headers.get("X-Request-ID") + try: + if request_id and str(uuid.UUID(request_id)) == request_id: + return request_id + except (AttributeError, TypeError, ValueError): + pass + return str(uuid.uuid4()) + + def register_request_id_middleware(app: Flask): """Adds request ID to each request @@ -12,4 +31,14 @@ 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())) + g.request_id = _incoming_request_id() + + @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..0b857063e6 100644 --- a/workers/shared/infrastructure/logging/logger.py +++ b/workers/shared/infrastructure/logging/logger.py @@ -33,6 +33,11 @@ class LogContext: organization_id: str | None = None correlation_id: str | None = None request_id: str | None = None + # Gates worker->worker re-propagation: only an id received from upstream may + # be stamped onto child tasks. A locally-derived fallback (task_id, or a + # payload id) must not be, or it would override the child's own + # file_execution_id correlation -- e.g. on beat-scheduled pipelines. + request_id_propagatable: bool = False class RequestIDFilter(logging.Filter): @@ -703,22 +708,72 @@ 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. + + Celery may surface a custom header either as a direct attribute on + ``task.request`` or only in its raw ``headers`` mapping, so both are tried. + """ + 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: message header, then a payload-derived id, then the Celery + ``task_id``. Only the header source is marked propagatable. + + The whole resolution stays inside the ``try``: an exception escaping here + would skip ``update_context`` and leave the *previous* task's id bound on a + pooled 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 an upstream correlation id flowing across worker->worker chains (e.g. + a file-processing task enqueuing a callback). Gated on + ``request_id_propagatable`` -- see ``LogContext`` for why a fallback id must + not fan out. 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 +783,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 +796,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/workers/tests/test_request_id_signals.py b/workers/tests/test_request_id_signals.py new file mode 100644 index 0000000000..1bc797a62c --- /dev/null +++ b/workers/tests/test_request_id_signals.py @@ -0,0 +1,240 @@ +"""Worker half of the ``X-Request-ID`` correlation chain. + +The backend's tests cover the producer side (injecting the id into a published +task's message headers). These cover what the worker does with it: preferring an +upstream header over a payload-derived id, refusing to fan out an id it derived +locally, clearing per-task state on a pooled thread, and re-emitting the id on +outbound internal-API calls. + +Two of these pin regressions that were caught by review rather than by a test -- +the propagatable gate, and the reset of that flag on task teardown. +""" + +import pytest + +from shared.clients.base_client import _current_request_id +from shared.infrastructure.logging.logger import ( + LogContext, + WorkerLogger, + _bind_task_context, + _clear_task_context, + _propagate_request_id_on_publish, + _request_id_from_message, +) + +HEADER_ID = "11111111-2222-4333-8444-555555555555" +PAYLOAD_ID = "99999999-8888-4777-8666-555555555555" +TASK_ID = "celery-task-id-0001" + + +class _Request: + """Stands in for ``celery.app.task.Context``. + + Celery may expose a custom message header as an attribute or leave it only + in the raw ``headers`` mapping, so both shapes are constructible here. + """ + + def __init__(self, request_id=None, headers=None, **payload): + if request_id is not None: + self.request_id = request_id + self.headers = headers + self.__dict__.update(payload) + + +class _Task: + def __init__(self, request): + self.request = request + self.name = "workers.test.task" + + +@pytest.fixture(autouse=True) +def _isolate_context(): + """Each test starts and ends on a clean thread-local context.""" + WorkerLogger.clear_context() + yield + WorkerLogger.clear_context() + + +# --------------------------------------------------------------------------- +# Reading the id off the message +# --------------------------------------------------------------------------- + + +def test_header_attribute_is_read(): + assert _request_id_from_message(_Task(_Request(request_id=HEADER_ID))) == HEADER_ID + + +def test_raw_headers_mapping_is_the_fallback(): + task = _Task(_Request(headers={"request_id": HEADER_ID})) + + assert _request_id_from_message(task) == HEADER_ID + + +def test_absent_header_reads_as_none(): + assert _request_id_from_message(_Task(_Request())) is None + + +def test_missing_request_object_does_not_raise(): + class _Bare: + request = None + + assert _request_id_from_message(_Bare()) is None + + +# --------------------------------------------------------------------------- +# Binding: precedence and the propagatable gate +# --------------------------------------------------------------------------- + + +def test_header_id_wins_over_payload_and_is_propagatable(): + """An upstream id is the authoritative correlation key, so it takes + precedence over anything derivable from the payload. + """ + task = _Task(_Request(request_id=HEADER_ID)) + + _bind_task_context(TASK_ID, task, (), {"file_execution_id": PAYLOAD_ID}) + + ctx = WorkerLogger.get_context() + assert ctx.request_id == HEADER_ID + assert ctx.request_id_propagatable is True + + +def test_payload_id_is_used_but_is_not_propagatable(): + """A payload-derived id correlates this task only. Marking it propagatable + would stamp it onto child tasks and override *their* own file_execution_id. + """ + task = _Task(_Request()) + + _bind_task_context(TASK_ID, task, (), {"file_execution_id": PAYLOAD_ID}) + + ctx = WorkerLogger.get_context() + assert ctx.request_id == PAYLOAD_ID + assert ctx.request_id_propagatable is False + + +def test_task_id_is_the_last_resort_and_is_not_propagatable(): + _bind_task_context(TASK_ID, _Task(_Request()), (), {}) + + ctx = WorkerLogger.get_context() + assert ctx.request_id == TASK_ID + assert ctx.request_id_propagatable is False + + +def test_a_raising_request_object_still_binds_the_task_id(): + """The whole resolution sits inside a try precisely so a malformed message + cannot leave the *previous* task's id bound on a reused thread. + """ + + class _Exploding: + @property + def request(self): + raise RuntimeError("malformed message") + + WorkerLogger.set_context(LogContext(request_id="stale-previous-task")) + + _bind_task_context(TASK_ID, _Exploding(), (), {}) + + ctx = WorkerLogger.get_context() + assert ctx.request_id == TASK_ID + assert ctx.request_id_propagatable is False + + +# --------------------------------------------------------------------------- +# Re-publishing to child tasks +# --------------------------------------------------------------------------- + + +def test_propagatable_id_is_stamped_onto_a_child_task(): + WorkerLogger.set_context( + LogContext(request_id=HEADER_ID, request_id_propagatable=True) + ) + headers = {} + + _propagate_request_id_on_publish(headers=headers) + + assert headers["request_id"] == HEADER_ID + + +def test_locally_derived_id_is_not_stamped_onto_a_child_task(): + """The regression the gate exists for: without it a scheduler-originated + task fans its own task_id out over every child's own correlation id. + """ + WorkerLogger.set_context( + LogContext(request_id=TASK_ID, request_id_propagatable=False) + ) + headers = {} + + _propagate_request_id_on_publish(headers=headers) + + assert headers == {} + + +def test_an_id_the_caller_already_set_is_left_alone(): + WorkerLogger.set_context( + LogContext(request_id=HEADER_ID, request_id_propagatable=True) + ) + headers = {"request_id": "explicitly-set-by-caller"} + + _propagate_request_id_on_publish(headers=headers) + + assert headers["request_id"] == "explicitly-set-by-caller" + + +def test_no_headers_mapping_is_a_no_op(): + WorkerLogger.set_context( + LogContext(request_id=HEADER_ID, request_id_propagatable=True) + ) + + _propagate_request_id_on_publish(headers=None) # does not raise + + +# --------------------------------------------------------------------------- +# Teardown +# --------------------------------------------------------------------------- + + +def test_teardown_resets_the_propagatable_flag(): + """Prefork/thread workers reuse the thread. A flag left True would let the + next task's locally-derived id fan out as though it came from upstream. + """ + _bind_task_context(TASK_ID, _Task(_Request(request_id=HEADER_ID)), (), {}) + + _clear_task_context() + + ctx = WorkerLogger.get_context() + assert ctx.request_id is None + assert ctx.request_id_propagatable is False + + +def test_a_child_published_after_teardown_inherits_nothing(): + _bind_task_context(TASK_ID, _Task(_Request(request_id=HEADER_ID)), (), {}) + _clear_task_context() + headers = {} + + _propagate_request_id_on_publish(headers=headers) + + assert headers == {} + + +# --------------------------------------------------------------------------- +# Outbound internal-API calls +# --------------------------------------------------------------------------- + + +def test_bound_id_is_offered_to_outbound_calls(): + _bind_task_context(TASK_ID, _Task(_Request(request_id=HEADER_ID)), (), {}) + + assert _current_request_id() == HEADER_ID + + +def test_placeholder_id_sends_no_header(): + """``"-"`` is the formatter's empty rendering, not an id -- sending it would + put a meaningless X-Request-ID on the wire. + """ + WorkerLogger.set_context(LogContext(request_id="-")) + + assert _current_request_id() is None + + +def test_no_context_sends_no_header(): + assert _current_request_id() is None 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..d2b51066f7 --- /dev/null +++ b/x2text-service/app/logging_util.py @@ -0,0 +1,108 @@ +"""Request-id-aware logging for the x2text-service. + +Deliberately a self-contained copy of ``unstract.core.flask``'s logging rather +than an import: this service does not take the ``unstract-core`` dependency. +""" + +import logging +import uuid +from logging.config import dictConfig + +from flask import Flask, g, has_request_context, request + +# Copy of the canonical format owned by ``unstract.core.flask.logging``; a +# divergence silently splits this service out of the cross-service log query. +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: + # Only the canonical UUID our services emit is adopted; see + # ``unstract.core.flask.middleware`` for why the shape is checked at all. + request_id = request.headers.get("X-Request-ID") + try: + if not (request_id and str(uuid.UUID(request_id)) == request_id): + request_id = str(uuid.uuid4()) + except (AttributeError, TypeError, ValueError): + request_id = str(uuid.uuid4()) + g.request_id = request_id + + @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