fastapi: let a custom request formatter read the request body

CHANGELOG.md

@@ -11,6 +11,10 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/).
 - Connexion 3 `AsyncApp` support: `init_connexion()` plus `init_request_instrument(app)` now accept an `AsyncApp`,
   which runs without Flask. Requests are logged by an ASGI middleware, including error responses produced by
   exception handlers, and log records from sync and async handlers share the request's correlation ID.
+- `init_request_instrument(..., read_request_body=True)` (FastAPI only) reads the body of each logged request before the
+  handler runs and exposes it as `request.state.json_logging_request_body` (bytes), so a custom request formatter can use
+  it without `await` (#84). It is off by default and the body is never added to the log automatically. The whole body is
+  held in memory. The handler can still read the body. It needs Starlette 0.28 or newer and raises `RuntimeError` otherwise. Other frameworks raise `ValueError` if it is requested.
 
 ### Fixed
 

README.md

@@ -287,6 +287,26 @@ handler.setFormatter(json_logging.formatters.JSONRequestLogFormatter())
 request_logger.addHandler(handler)
 ```
 
+### Read the request body in a custom request formatter (FastAPI)
+
+The formatter methods are not `async`, so they cannot `await request.body()`. Pass `read_request_body=True` and the
+body is read for you before the handler runs, and is available as bytes on `request.state.json_logging_request_body`.
+It needs Starlette 0.28 or newer (older versions hang), otherwise `init_request_instrument` raises `RuntimeError`.
+It is not logged unless your formatter adds it. Do this only for payloads you are sure may be logged, and note that
+the whole body is held in memory.
+
+```python
+class RequestFormatter(json_logging.JSONRequestLogFormatter):
+    def _format_log_object(self, record, request_util):
+        log_object = super()._format_log_object(record, request_util)
+        request = record.request_response_data._request
+        body = getattr(request.state, "json_logging_request_body", b"")
+        log_object["request_body"] = body.decode(errors="replace")
+        return log_object
+
+json_logging.init_request_instrument(app, custom_formatter=RequestFormatter, read_request_body=True)
+```
+
 ## 2.3 Get current correlation-id
 
 The current request's correlation-id can be retrieved and passed to downstream service calls as follows:

json_logging/__init__.py

@@ -1,4 +1,5 @@
 # coding=utf-8
+import inspect
 import json
 import logging
 import sys
@@ -146,7 +147,7 @@ def __init(framework_name=None, custom_formatter=None, enable_json=False):
 
 
 def init_request_instrument(app=None, custom_formatter=None, exclude_url_patterns=[],
-                            request_response_dto_class=DefaultRequestResponseDTO):
+                            request_response_dto_class=DefaultRequestResponseDTO, read_request_body=False):
     """
     Configure the request instrumentation logging configuration for given web app. Must be called after init method
 
@@ -155,6 +156,9 @@ def init_request_instrument(app=None, custom_formatter=None, exclude_url_pattern
     :param app: current web application instance
     :param custom_formatter: formatter to override default JSONRequestLogFormatter.
     :param request_response_dto_class: request_response_dto_class to override default json_logging.RequestResponseDataExtractor.
+    :param read_request_body: FastAPI only. Read each logged request's body before the handler runs and expose it as
+        ``request.state.json_logging_request_body`` (bytes), so a custom formatter can use it without ``await``.
+        Off by default and never added to the log automatically. The whole body is held in memory.
     """
 
     if _current_framework is None or _current_framework == '-':
@@ -172,7 +176,13 @@ def init_request_instrument(app=None, custom_formatter=None, exclude_url_pattern
     _validated_excluded_fields(EXCLUDED_FIELDS)
 
     configurator = _current_framework['app_request_instrumentation_configurator']()
-    configurator.config(app, request_response_dto_class, exclude_url_patterns=exclude_url_patterns)
+    config_options = {}
+    if read_request_body:
+        if 'read_request_body' not in inspect.signature(configurator.config).parameters:
+            raise ValueError('read_request_body is only supported for FastAPI')
+        config_options['read_request_body'] = True
+    configurator.config(app, request_response_dto_class, exclude_url_patterns=exclude_url_patterns,
+                        **config_options)
 
     formatter = custom_formatter if custom_formatter else JSONRequestLogFormatter
     request_logger = configurator.request_logger

json_logging/framework/fastapi/implementation.py

@@ -1,4 +1,5 @@
 import logging
+import re
 import sys
 
 import json_logging
@@ -9,6 +10,7 @@
 from json_logging.util import is_not_match_any_pattern
 
 import fastapi
+import starlette
 import starlette.requests
 import starlette.responses
 
@@ -18,6 +20,8 @@
 from starlette.types import ASGIApp
 
 _request_config_class = None
+# Before this Starlette version, reading the body in a BaseHTTPMiddleware leaves the endpoint waiting forever
+_MIN_STARLETTE_FOR_REQUEST_BODY = (0, 28)
 # Only enabled on the Python versions covered by CI (3.9+). On older versions, including 3.6 where at
 # most the thread-local `contextvars` backport exists, the previous call-stack lookup is kept unchanged.
 if sys.version_info >= (3, 9):
@@ -28,8 +32,9 @@
 
 
 class JSONLoggingASGIMiddleware(BaseHTTPMiddleware):
-    def __init__(self, app: ASGIApp, exclude_url_patterns=tuple()) -> None:
+    def __init__(self, app: ASGIApp, exclude_url_patterns=tuple(), read_request_body=False) -> None:
         super().__init__(app)
+        self.read_request_body = read_request_body
         self.request_logger = logging.getLogger('fastapi-request-logger')
         self.exclude_url_patterns = exclude_url_patterns
         logging.getLogger("uvicorn.access").propagate = False
@@ -58,7 +63,13 @@ async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -
             if not log_request:
                 return await call_next(request)
 
+            # Created first so that the response time also covers reading the body
             request_response_data = _request_config_class(request)
+
+            if self.read_request_body:
+                # Starlette caches the body, so the endpoint can still read it afterwards
+                request.state.json_logging_request_body = await request.body()
+
             response = await call_next(request)
             request_response_data.on_request_complete(response)
             self.request_logger.info(
@@ -71,10 +82,16 @@ async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -
 
 
 class FastAPIAppRequestInstrumentationConfigurator(BaseAppRequestInstrumentationConfigurator):
-    def config(self, app, request_response_dto_class, exclude_url_patterns=[]):
+    def config(self, app, request_response_dto_class, exclude_url_patterns=[], read_request_body=False):
         if not isinstance(app, fastapi.FastAPI):
             raise RuntimeError("app is not a valid fastapi.FastAPI instance")
 
+        if read_request_body:
+            starlette_version = tuple(int(part) for part in re.findall(r'\d+', starlette.__version__)[:2])
+            if starlette_version < _MIN_STARLETTE_FOR_REQUEST_BODY:
+                raise RuntimeError("read_request_body needs starlette >= %s, found %s" % (
+                    '.'.join(map(str, _MIN_STARLETTE_FOR_REQUEST_BODY)), starlette.__version__))
+
         global _request_config_class
         _request_config_class = request_response_dto_class
 
@@ -84,7 +101,8 @@ def config(self, app, request_response_dto_class, exclude_url_patterns=[]):
         # noinspection PyAttributeOutsideInit
         self.request_logger = logging.getLogger('fastapi-request-logger')
 
-        app.add_middleware(JSONLoggingASGIMiddleware, exclude_url_patterns=exclude_url_patterns)
+        app.add_middleware(JSONLoggingASGIMiddleware, exclude_url_patterns=exclude_url_patterns,
+                           read_request_body=read_request_body)
 
 
 class FastAPIRequestInfoExtractor(BaseRequestInfoExtractor):

tests/test_fastapi_request_body.py

@@ -0,0 +1,195 @@
+"""Tests for reading the request body from a custom request formatter (FastAPI)"""
+import json
+import logging
+import re
+
+import fastapi
+import fastapi.testclient
+import pytest
+import starlette
+from helpers.handler import FormattedMessageCollectorHandler
+from helpers.imports import undo_imports_from_package
+
+
+def _starlette_can_capture_body():
+    return tuple(int(part) for part in re.findall(r"\d+", starlette.__version__)[:2]) >= (0, 28)
+
+
+needs_body_capture = pytest.mark.skipif(
+    not _starlette_can_capture_body(), reason="read_request_body needs starlette >= 0.28"
+)
+
+
+def _build_client(read_request_body, exclude_url_patterns=("^/excluded",)):
+    import json_logging
+
+    class BodyFormatter(json_logging.JSONRequestLogFormatter):
+        """Adds the body to the request log, the way a user would in a synchronous formatter."""
+
+        def _format_log_object(self, record, request_util):
+            log_object = super()._format_log_object(record, request_util)
+            request = record.request_response_data._request
+            body = getattr(request.state, "json_logging_request_body", None)
+            log_object["request_body"] = body.decode() if body is not None else None
+            return log_object
+
+    app = fastapi.FastAPI()
+    json_logging.init_fastapi(enable_json=True)
+    json_logging.init_request_instrument(
+        app,
+        custom_formatter=BodyFormatter,
+        exclude_url_patterns=list(exclude_url_patterns),
+        read_request_body=read_request_body,
+    )
+
+    @app.post("/echo")
+    async def echo(request: fastapi.Request):
+        return {"body": (await request.body()).decode()}
+
+    @app.post("/echo-sync")
+    def echo_sync(payload: dict):
+        return payload
+
+    @app.post("/excluded")
+    async def excluded(request: fastapi.Request):
+        return {"body": (await request.body()).decode()}
+
+    request_logger = logging.getLogger("fastapi-request-logger")
+    handler = FormattedMessageCollectorHandler()
+    handler.setFormatter(request_logger.handlers[0].formatter)
+    request_logger.addHandler(handler)
+    return fastapi.testclient.TestClient(app), handler, request_logger
+
+
+@pytest.fixture
+def body_client():
+    if not _starlette_can_capture_body():
+        pytest.skip("read_request_body needs starlette >= 0.28")
+    client, handler, request_logger = _build_client(read_request_body=True)
+    yield client, handler
+    request_logger.removeHandler(handler)
+    undo_imports_from_package("json_logging")  # Necessary because of json-logging's global state
+
+
+@pytest.fixture
+def default_client():
+    client, handler, request_logger = _build_client(read_request_body=False)
+    yield client, handler
+    request_logger.removeHandler(handler)
+    undo_imports_from_package("json_logging")
+
+
+def test_formatter_can_read_request_body(body_client):
+    client, handler = body_client
+
+    response = client.post("/echo", content=b"hello body")
+
+    assert response.status_code == 200
+    assert json.loads(handler.messages[0])["request_body"] == "hello body"
+
+
+def test_endpoint_can_still_read_body_after_middleware(body_client):
+    client, _ = body_client
+
+    assert client.post("/echo", content=b"hello body").json() == {"body": "hello body"}
+
+
+def test_sync_endpoint_with_parsed_body_still_works(body_client):
+    client, handler = body_client
+
+    response = client.post("/echo-sync", json={"a": 1})
+
+    assert response.json() == {"a": 1}
+    assert json.loads(json.loads(handler.messages[0])["request_body"]) == {"a": 1}
+
+
+def test_empty_body_is_an_empty_string(body_client):
+    client, handler = body_client
+
+    client.post("/echo")
+
+    assert json.loads(handler.messages[0])["request_body"] == ""
+
+
+def test_body_is_not_read_by_default(default_client):
+    client, handler = default_client
+
+    response = client.post("/echo", content=b"hello body")
+
+    assert response.json() == {"body": "hello body"}
+    assert json.loads(handler.messages[0])["request_body"] is None
+
+
+def test_excluded_url_body_is_not_read_for_logging(body_client):
+    client, handler = body_client
+
+    response = client.post("/excluded", content=b"secret")
+
+    assert response.json() == {"body": "secret"}
+    assert handler.messages == []
+
+
+def test_each_request_gets_its_own_body(body_client):
+    client, handler = body_client
+
+    client.post("/echo", content=b"first")
+    client.post("/echo", content=b"second")
+
+    assert [json.loads(m)["request_body"] for m in handler.messages] == ["first", "second"]
+
+
+def test_read_request_body_is_rejected_for_frameworks_that_do_not_support_it():
+    import flask
+    import json_logging
+
+    try:
+        json_logging.init_flask(enable_json=True)
+        with pytest.raises(ValueError, match="read_request_body"):
+            json_logging.init_request_instrument(flask.Flask(__name__), read_request_body=True)
+    finally:
+        undo_imports_from_package("json_logging")
+
+
+def test_read_request_body_is_rejected_on_starlette_that_would_hang(monkeypatch):
+    import starlette
+    import json_logging
+
+    try:
+        json_logging.init_fastapi(enable_json=True)
+        monkeypatch.setattr(starlette, "__version__", "0.27.0")
+        with pytest.raises(RuntimeError, match="starlette >= 0.28"):
+            json_logging.init_request_instrument(fastapi.FastAPI(), read_request_body=True)
+    finally:
+        undo_imports_from_package("json_logging")
+
+
+@needs_body_capture
+def test_response_time_includes_reading_the_body(monkeypatch):
+    """The request timer starts before the body is read, not after."""
+    import json_logging
+
+    client, handler, request_logger = _build_client(read_request_body=True)
+    try:
+        events = []
+        dto_class = json_logging.framework.fastapi.implementation._request_config_class
+        original_init = dto_class.__init__
+
+        def recording_init(self, request, **kwargs):
+            events.append("timer started")
+            original_init(self, request, **kwargs)
+
+        monkeypatch.setattr(dto_class, "__init__", recording_init)
+        original_body = fastapi.Request.body
+
+        async def recording_body(self):
+            events.append("body read")
+            return await original_body(self)
+
+        monkeypatch.setattr(fastapi.Request, "body", recording_body)
+
+        client.post("/echo", content=b"x")
+
+        assert events.index("timer started") < events.index("body read")
+    finally:
+        request_logger.removeHandler(handler)
+        undo_imports_from_package("json_logging")