From 4bad6254429d9145a904c69c6c06c8bccf918cf6 Mon Sep 17 00:00:00 2001 From: Mathieu Leplatre Date: Wed, 28 Feb 2024 16:46:30 +0100 Subject: [PATCH] (fastapi) Log request ID when set in headers (#100) * (fastapi) Log request ID when set in headers * Implement option 1 * Add mention in docs * Add and document a logging filter --- docs/fastapi.rst | 11 ++++-- setup.py | 2 +- src/dockerflow/fastapi/__init__.py | 14 ++++++++ src/dockerflow/fastapi/middleware.py | 2 ++ tests/fastapi/test_fastapi.py | 53 +++++++++++++++++++++++++++- tests/requirements/fastapi.txt | 1 + 6 files changed, 79 insertions(+), 4 deletions(-) diff --git a/docs/fastapi.rst b/docs/fastapi.rst index c48e202..523bdc3 100644 --- a/docs/fastapi.rst +++ b/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 `_ package and with the `RequestIdFilter` added in logging configuration as shown above. .. _fastapi-static: diff --git a/setup.py b/setup.py index 1c1c3cb..d26aac2 100644 --- a/setup.py +++ b/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", diff --git a/src/dockerflow/fastapi/__init__.py b/src/dockerflow/fastapi/__init__.py index 19596b0..102279f 100644 --- a/src/dockerflow/fastapi/__init__.py +++ b/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 diff --git a/src/dockerflow/fastapi/middleware.py b/src/dockerflow/fastapi/middleware.py index 7533a38..dbf15c6 100644 --- a/src/dockerflow/fastapi/middleware.py +++ b/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): diff --git a/tests/fastapi/test_fastapi.py b/tests/fastapi/test_fastapi.py index 4c51da9..d0815c7 100644 --- a/tests/fastapi/test_fastapi.py +++ b/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") diff --git a/tests/requirements/fastapi.txt b/tests/requirements/fastapi.txt index 18476e7..99cb858 100644 --- a/tests/requirements/fastapi.txt +++ b/tests/requirements/fastapi.txt @@ -1,3 +1,4 @@ fastapi asgiref +asgi-correlation-id httpx \ No newline at end of file