Generalize usage of request correlation id to all integrations (#101)

* Generalize usage of request correlation id to all integrations

* Use json.loads()

* Leverage caplog handler as suggested by @grahamalama

Author: leplatrem

docs/changelog.rst

@@ -1,6 +1,14 @@
 Changelog
 ---------
 
+2024.2.1
+~~~~~~~~~~~~~~~~~~~~~
+
+- Add request correlation ID support (#101).
+  In order to add a ``rid`` field to all log messages, add the ``dockerflow.logging.RequestIdLogFilter`` filter to your logging handlers.
+  See: :ref:`Django <django-logging>`, :ref:`FastAPI <fastapi-logging>`, :ref:`Flask <flask-logging>`, :ref:`Sanic <sanic-logging>` for details.
+
+
 2024.2.0
 ~~~~~~~~~~~~~~~~~~~~~
 

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,17 @@ 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, 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.
+
+The header name to obtain the request ID can be customized in settings:
+
+.. code-block:: python
+
+    DOCKERFLOW_REQUEST_ID_HEADER_NAME = "X-Cloud-Trace-Context"
+
+
 .. _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,16 @@ 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.
+
+The header name to obtain the request ID can be customized in settings:
+
+.. code-block:: python
+
+    app.state.DOCKERFLOW_REQUEST_ID_HEADER_NAME = "X-Cloud-Trace-Context"
+
 
 .. _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,17 @@ 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, 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.
+
+The header name to obtain the request ID can be customized in settings:
+
+.. code-block:: python
+
+    DOCKERFLOW_REQUEST_ID_HEADER_NAME = "X-Cloud-Trace-Context"
+
+
 .. _flask-static:
 
 Static content

docs/logging.rst

@@ -39,11 +39,17 @@ this::
                 'logger_name': 'myproject'
             }
         },
+        'filters': {
+            'request_id': {
+                '()': 'dockerflow.logging.RequestIdFilter',
+            },
+        },
         'handlers': {
             'console': {
                 'level': 'DEBUG',
                 'class': 'logging.StreamHandler',
-                'formatter': 'json'
+                'formatter': 'json',
+                'filters': ['request_id']
             },
         },
         'loggers': {
@@ -80,6 +86,12 @@ thing as the dictionary based configuratio above:
     [formatters]
     keys = json
 
+    [filters]
+    keys = request_id
+
+    [filter_request_id]
+    class = dockerflow.logging.RequestIdFilter
+
     [logger_root]
     level = INFO
     handlers = console
@@ -94,6 +106,7 @@ thing as the dictionary based configuratio above:
     level = DEBUG
     args = (sys.stdout,)
     formatter = json
+    filters = request_id
 
     [formatter_json]
     class = dockerflow.logging.JsonLogFormatter

docs/sanic.rst

@@ -421,11 +421,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': {
@@ -459,6 +465,16 @@ 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, 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.
+
+The header name to obtain the request ID can be customized in settings:
+
+.. code-block:: python
+
+    DOCKERFLOW_REQUEST_ID_HEADER_NAME = "X-Cloud-Trace-Context"
+
 
 .. _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

@@ -3,14 +3,25 @@
 import time
 import typing
 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
 
 
+def extract_request_id(request):
+    """Extract request ID from request."""
+    rid = get_or_generate_request_id(
+        request.headers,
+        header_name=getattr(settings, "DOCKERFLOW_REQUEST_ID_HEADER_NAME", None),
+    )
+    request_id_context.set(rid)
+    request._id = rid  # Used in tests.
+
+
 class DockerflowMiddleware(MiddlewareMixin):
     """
     Emits a request.summary type log entry for every request.
@@ -34,7 +45,8 @@ def process_request(self, request):
             if pattern.match(request.path_info):
                 return view(request)
 
-        request._id = str(uuid.uuid4())
+        extract_request_id(request)
+
         request._start_timestamp = time.time()
         return None
 
@@ -57,8 +69,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,37 @@
     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
+
+        rid = get_or_generate_request_id(
+            headers,
+            header_name=getattr(
+                scope["app"].state, "DOCKERFLOW_REQUEST_ID_HEADER_NAME", None
+            ),
+        )
+        request_id_context.set(rid)
+
+        await self.app(scope, receive, send)
 
 
 class MozlogRequestSummaryLogger:
@@ -75,7 +104,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):