Generalize usage of request correlation id to all integrations

Author: leplatrem

docs/django.rst

@@ -386,11 +386,17 @@ configure **at least** the ``request.summary`` logger that way::
                 'logger_name': 'myproject'
             }
         },
+        'filters': {
+            'request_id': {
+                '()': 'dockerflow.logging.RequestIdFilter',
+            },
+        },
         'handlers': {
             'console': {
                 'level': 'DEBUG',
                 'class': 'logging.StreamHandler',
-                'formatter': 'json'
+                'formatter': 'json',
+                'filters': ['request_id']
             },
         },
         'loggers': {
@@ -408,6 +414,10 @@ In order to include querystrings in the request summary log, set this flag in se
     DOCKERFLOW_SUMMARY_LOG_QUERYSTRING = True
 
 
+A unique request ID is read from the `X-Request-ID` request header using the `RequestIdMiddleware` middleware (see :ref:`django-setup`), and a UUID4 value is generated if unset.
+
+Leveraging the `RequestIdFilter` in logging configuration as shown above will add a ``rid`` attribute to all log messages.
+
 .. _django-static:
 
 Static content

docs/fastapi.rst

@@ -29,6 +29,8 @@ The package ``dockerflow.fastapi`` package implements various tools to support
     For more information see the :doc:`API documentation <api/fastapi>` for
     the ``dockerflow.fastapi`` module.
 
+.. _fastapi-setup:
+
 Setup
 -----
 
@@ -39,12 +41,12 @@ To install ``python-dockerflow``'s FastAPI support please follow these steps:
 
      from fastapi import FastAPI
      from dockerflow.fastapi import router
-     from dockerflow.fastapi.middleware import MozlogRequestSummaryLogger, CorrelationIdMiddleware
+     from dockerflow.fastapi.middleware import MozlogRequestSummaryLogger, RequestIdMiddleware
 
      app = FastAPI()
      app.include_router(router)
      app.add_middleware(MozlogRequestSummaryLogger)
-     app.add_middleware(CorrelationIdMiddleware)  # see snok/asgi-correlation-id
+     app.add_middleware(RequestIdMiddleware)  # see snok/asgi-correlation-id
 
 #. Make sure the app root path is set correctly as this will be used
    to locate the ``version.json`` file that is generated by
@@ -298,7 +300,7 @@ for at least the ``request.summary`` logger:
         },
         'filters': {
             'request_id': {
-                '()': 'dockerflow.fastapi.RequestIdFilter',
+                '()': 'dockerflow.logging.RequestIdFilter',
             },
         },
         'handlers': {
@@ -324,7 +326,10 @@ In order to include querystrings in the request summary log, set this flag in th
 
     app.state.DOCKERFLOW_SUMMARY_LOG_QUERYSTRING = True
 
-A unique ID is added to each log message, using the `asgi-correlation-id <https://github.com/snok/asgi-correlation-id>`_ package and with the `RequestIdFilter` added in logging configuration as shown above.
+A unique request ID is read from the `X-Request-ID` request header using the `RequestIdMiddleware` middleware (see :ref:`fastapi-setup`), and a UUID4 value is generated if unset.
+
+Leveraging the `RequestIdFilter` in logging configuration as shown above will add a ``rid`` attribute to all log messages.
+
 
 .. _fastapi-static:
 

docs/flask.rst

@@ -441,11 +441,17 @@ for at least the ``request.summary`` logger::
                 'logger_name': 'myproject'
             }
         },
+        'filters': {
+            'request_id': {
+                '()': 'dockerflow.logging.RequestIdFilter',
+            },
+        },
         'handlers': {
             'console': {
                 'level': 'DEBUG',
                 'class': 'logging.StreamHandler',
-                'formatter': 'json'
+                'formatter': 'json',
+                'filters': ['request_id']
             },
         },
         'loggers': {
@@ -460,6 +466,10 @@ In order to include querystrings in the request summary log, set this flag in :r
 
     DOCKERFLOW_SUMMARY_LOG_QUERYSTRING = True
 
+A unique request ID is read from the `X-Request-ID` request header using the `RequestIdMiddleware` middleware (see :ref:`flask-setup`), and a UUID4 value is generated if unset.
+
+Leveraging the `RequestIdFilter` in logging configuration as shown above will add a ``rid`` attribute to all log messages.
+
 .. _flask-static:
 
 Static content

docs/sanic.rst

@@ -459,6 +459,10 @@ In order to include querystrings in the request summary log, set this flag in :r
 
     DOCKERFLOW_SUMMARY_LOG_QUERYSTRING = True
 
+A unique request ID is read from the `X-Request-ID` request header using the `RequestIdMiddleware` middleware (see :ref:`django-setup`), and a UUID4 value is generated if unset.
+
+Leveraging the `RequestIdFilter` in logging configuration as shown above will add a ``rid`` attribute to all log messages.
+
 
 .. _sanic-static:
 

setup.py

@@ -50,7 +50,7 @@ def read(*parts):
         "django": ["django"],
         "flask": ["flask", "blinker"],
         "sanic": ["sanic"],
-        "fastapi": ["fastapi", "asgiref", "asgi-correlation-id"],
+        "fastapi": ["fastapi", "asgiref"],
     },
     zip_safe=False,
     python_requires=">=3.7,<4",

src/dockerflow/django/middleware.py

@@ -2,11 +2,12 @@
 import re
 import time
 import urllib
-import uuid
 
 from django.conf import settings
 from django.utils.deprecation import MiddlewareMixin
 
+from dockerflow.logging import get_or_generate_request_id, request_id_context
+
 from . import views
 
 
@@ -33,7 +34,10 @@ def process_request(self, request):
             if pattern.match(request.path_info):
                 return view(request)
 
-        request._id = str(uuid.uuid4())
+        rid = get_or_generate_request_id(request.headers)
+        request_id_context.set(rid)
+        request._id = rid  # Used in tests.
+
         request._start_timestamp = time.time()
         return None
 
@@ -56,8 +60,7 @@ def _build_extra_meta(self, request):
         # attributes before trying to use them.
         if hasattr(request, "user"):
             out["uid"] = request.user.is_authenticated and request.user.pk or ""
-        if hasattr(request, "_id"):
-            out["rid"] = request._id
+        out["rid"] = request_id_context.get()
         if hasattr(request, "_start_timestamp"):
             # Duration of request, in milliseconds.
             out["t"] = int(1000 * (time.time() - request._start_timestamp))

src/dockerflow/fastapi/__init__.py

@@ -1,6 +1,3 @@
-from logging import Filter, LogRecord
-
-from asgi_correlation_id import correlation_id
 from fastapi import APIRouter
 from fastapi.routing import APIRoute
 
@@ -15,14 +12,3 @@
     ],
 )
 """This router adds the Dockerflow views."""
-
-
-class RequestIdLogFilter(Filter):
-    """Logging filter to attach request IDs to log records"""
-
-    def filter(self, record: "LogRecord") -> bool:
-        """
-        Attach the request ID to the log record.
-        """
-        record.rid = correlation_id.get(None)
-        return True

src/dockerflow/fastapi/middleware.py

@@ -6,7 +6,6 @@
 import urllib
 from typing import Any, Dict
 
-from asgi_correlation_id import CorrelationIdMiddleware, correlation_id  # noqa
 from asgiref.typing import (
     ASGI3Application,
     ASGIReceiveCallable,
@@ -15,7 +14,30 @@
     HTTPScope,
 )
 
-from ..logging import JsonLogFormatter
+from ..logging import JsonLogFormatter, get_or_generate_request_id, request_id_context
+
+
+class RequestIdMiddleware:
+    def __init__(
+        self,
+        app: ASGI3Application,
+    ) -> None:
+        self.app = app
+
+    async def __call__(
+        self, scope: HTTPScope, receive: ASGIReceiveCallable, send: ASGISendCallable
+    ) -> None:
+        if scope["type"] != "http":
+            return await self.app(scope, receive, send)
+
+        headers = {}
+        for name, value in scope["headers"]:
+            header_key = name.decode("latin1").lower()
+            header_val = value.decode("latin1")
+            headers[header_key] = header_val
+
+        request_id_context.set(get_or_generate_request_id(headers))
+        await self.app(scope, receive, send)
 
 
 class MozlogRequestSummaryLogger:
@@ -75,7 +97,7 @@ def _format(self, scope: HTTPScope, info) -> Dict[str, Any]:
             "code": info["response"]["status"],
             "lang": info["request_headers"].get("accept-language"),
             "t": int(request_duration_ms),
-            "rid": correlation_id.get(),
+            "rid": request_id_context.get(),
         }
 
         if getattr(scope["app"].state, "DOCKERFLOW_SUMMARY_LOG_QUERYSTRING", False):

src/dockerflow/flask/app.py

@@ -4,13 +4,13 @@
 import logging
 import os
 import time
-import uuid
 import warnings
 
 import flask
 from werkzeug.exceptions import InternalServerError
 
 from dockerflow import checks
+from dockerflow.logging import get_or_generate_request_id, request_id_context
 
 from .. import version
 from .checks import (
@@ -188,9 +188,12 @@ def _before_request(self):
         """
         The before_request callback.
         """
-        flask.g._request_id = str(uuid.uuid4())
+        rid = get_or_generate_request_id(flask.request.headers)
+        request_id_context.set(rid)
+        flask.g._request_id = rid  # For retro-compatibility and tests.
         if not hasattr(flask.g, "request_id"):
-            flask.g.request_id = flask.g._request_id
+            flask.g.request_id = rid
+
         flask.g._start_timestamp = time.time()
 
     def _after_request(self, response):
@@ -269,9 +272,7 @@ def summary_extra(self):
         out["uid"] = user_id
 
         # the rid value to the current request ID
-        request_id = flask.g.get("_request_id", None)
-        if request_id is not None:
-            out["rid"] = request_id
+        out["rid"] = request_id_context.get()
 
         # and the t value to the time it took to render
         start_timestamp = flask.g.get("_start_timestamp", None)

src/dockerflow/logging.py

@@ -8,6 +8,9 @@
 import socket
 import sys
 import traceback
+import uuid
+from contextvars import ContextVar
+from typing import Optional
 
 
 class SafeJSONEncoder(json.JSONEncoder):
@@ -153,3 +156,25 @@ def safer_format_traceback(exc_typ, exc_val, exc_tb):
     lines.append("%r\n" % (exc_typ,))
     lines.append("%r\n" % (exc_val,))
     return "".join(lines)
+
+
+request_id_context: ContextVar[Optional[str]] = ContextVar("request_id", default=None)
+
+
+def get_or_generate_request_id(headers: dict) -> str:
+    header_name = "x-request-id"
+    rid = headers.get(header_name, "")
+    if not rid:
+        rid = str(uuid.uuid4())
+    return rid
+
+
+class RequestIdLogFilter(logging.Filter):
+    """Logging filter to attach request IDs to log records"""
+
+    def filter(self, record: "logging.LogRecord") -> bool:
+        """
+        Attach the request ID to the log record.
+        """
+        record.rid = request_id_context.get(None)
+        return True