(fastapi) Log request ID when set in headers

docs/fastapi.rst

@@ -39,11 +39,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
+     from dockerflow.fastapi.middleware import MozlogRequestSummaryLogger, CorrelationIdMiddleware
 
      app = FastAPI()
      app.include_router(router)
      app.add_middleware(MozlogRequestSummaryLogger)
+     app.add_middleware(CorrelationIdMiddleware)  # 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
@@ -295,10 +296,16 @@ for at least the ``request.summary`` logger:
                 'logger_name': 'myproject'
             }
         },
+        'filters': {
+            'request_id': {
+                '()': 'dockerflow.fastapi.RequestIdFilter',
+            },
+        },
         'handlers': {
             'console': {
                 'level': 'DEBUG',
                 'class': 'logging.StreamHandler',
+                'filters': ['request_id'],
                 'formatter': 'json'
             },
         },
@@ -317,7 +324,7 @@ 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.
 
 .. _fastapi-static:
 

setup.py

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

src/dockerflow/fastapi/__init__.py

@@ -1,3 +1,6 @@
+from logging import Filter, LogRecord
+
+from asgi_correlation_id import correlation_id
 from fastapi import APIRouter
 from fastapi.routing import APIRoute
 
@@ -12,3 +15,14 @@
     ],
 )
 """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,6 +6,7 @@
 import urllib
 from typing import Any, Dict
 
+from asgi_correlation_id import CorrelationIdMiddleware, correlation_id  # noqa
 from asgiref.typing import (
     ASGI3Application,
     ASGIReceiveCallable,
@@ -74,6 +75,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(),
         }
 
         if getattr(scope["app"].state, "DOCKERFLOW_SUMMARY_LOG_QUERYSTRING", False):

tests/fastapi/test_fastapi.py

@@ -9,14 +9,19 @@
 from fastapi.testclient import TestClient
 
 from dockerflow import checks
+from dockerflow.fastapi import RequestIdLogFilter
 from dockerflow.fastapi import router as dockerflow_router
-from dockerflow.fastapi.middleware import MozlogRequestSummaryLogger
+from dockerflow.fastapi.middleware import (
+    CorrelationIdMiddleware,
+    MozlogRequestSummaryLogger,
+)
 
 
 def create_app():
     app = FastAPI()
     app.include_router(dockerflow_router)
     app.add_middleware(MozlogRequestSummaryLogger)
+    app.add_middleware(CorrelationIdMiddleware, validator=None)
     return app
 
 
@@ -64,6 +69,52 @@ def test_mozlog(app, client, caplog):
     assert isinstance(record.t, int)
 
 
+def test_mozlog_request_id(client, caplog):
+    client.get(
+        "/__lbheartbeat__",
+        headers={
+            "X-Request-ID": "tracked-value",
+        },
+    )
+    record = caplog.records[0]
+
+    assert record.rid == "tracked-value"
+
+
+def test_mozlog_without_correlation_id_middleware(client, caplog):
+    app = FastAPI()
+    app.include_router(dockerflow_router)
+    app.add_middleware(MozlogRequestSummaryLogger)
+    client = TestClient(app)
+
+    client.get("/__lbheartbeat__")
+    record = caplog.records[0]
+
+    assert record.rid is None
+
+
+def test_request_id_passed_to_all_log_messages(caplog):
+    handler = logging.StreamHandler()
+    handler.addFilter(RequestIdLogFilter())
+    _logger = logging.getLogger("some_logger")
+    _logger.addHandler(handler)
+
+    app = create_app()
+
+    @app.get("/ping")
+    def ping():
+        logger = logging.getLogger("some_logger")
+        logger.info("returning pong")
+        return "pong"
+
+    client = TestClient(app)
+
+    client.get("/ping")
+
+    log_message = next(r for r in caplog.records if r.name == "some_logger")
+    assert log_message.rid is not None
+
+
 def test_mozlog_failure(client, mocker, caplog):
     mocker.patch(
         "dockerflow.fastapi.views.get_version", side_effect=ValueError("crash")

tests/requirements/fastapi.txt

@@ -1,3 +1,4 @@
 fastapi
 asgiref
+asgi-correlation-id
 httpx