From f088ddd320bc4c410a97dfe4b4ac7424d98aac02 Mon Sep 17 00:00:00 2001
From: Abraham Chavez <achave11@ucsc.edu>
Date: Fri, 30 Aug 2024 18:06:12 -0700
Subject: [PATCH] Fix: Log request body (#6404)

---
 src/azul/chalice.py              | 24 ++++++++++++++++++------
 src/azul/logging.py              | 18 ++++++++++++------
 test/service/test_app_logging.py |  7 ++++++-
 3 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/src/azul/chalice.py b/src/azul/chalice.py
index f7ad22b12c..db5e758d13 100644
--- a/src/azul/chalice.py
+++ b/src/azul/chalice.py
@@ -61,6 +61,9 @@
     copy_json,
     json_head,
 )
+from azul.logging import (
+    http_body_log_message,
+)
 from azul.types import (
     JSON,
     LambdaContext,
@@ -376,24 +379,33 @@ def __authenticate(self):
             log.info('Authenticated request as %r', auth)
 
     def _log_request(self):
+        log.info(logging.INFO)
         if log.isEnabledFor(logging.INFO):
             context = self.current_request.context
             request_info = {
                 'query': self.current_request.query_params,
                 'headers': self.current_request.headers
             }
-            log.info('Received %s request for %r, with %s.',
+            log.info('Received %s request for %r, with %s',
                      context['httpMethod'],
                      context['path'],
                      json.dumps(request_info, cls=self._LogJSONEncoder))
 
+            body, n = self._body_logging_prep(self.current_request.raw_body)
+            log.info(http_body_log_message('request', body, max_len=n))
+
+    def _body_logging_prep(self, body: str | JSON) -> tuple[str, int]:
+        if config.debug == 2:
+            body = json.dumps(body, cls=self._LogJSONEncoder)
+            n = len(body)
+        else:
+            n = 1024
+            body = body[:n] if isinstance(body, str) else json_head(n, body)
+        return body, n
+
     def _log_response(self, response):
         if log.isEnabledFor(logging.DEBUG):
-            n = 1024
-            if isinstance(response.body, str):
-                body = response.body[:n]
-            else:
-                body = json_head(n, response.body)
+            body, n = self._body_logging_prep(response.body)
             log.debug('Returning %i response with headers %s. '
                       'See next line for the first %i characters of the body.\n%s',
                       response.status_code,
diff --git a/src/azul/logging.py b/src/azul/logging.py
index f5a24bb37a..7649f099f2 100644
--- a/src/azul/logging.py
+++ b/src/azul/logging.py
@@ -4,6 +4,7 @@
 import logging
 from typing import (
     Optional,
+    TYPE_CHECKING,
 )
 
 import attr
@@ -15,13 +16,18 @@
 from azul import (
     config,
 )
-from azul.chalice import (
-    AzulChaliceApp,
-)
 from azul.strings import (
     trunc_ellipses,
 )
 
+if TYPE_CHECKING:
+    from azul.chalice import (
+        AzulChaliceApp,
+    )
+else:
+    class AzulChaliceApp:
+        pass
+
 
 @attr.s(frozen=False, kw_only=False, auto_attribs=True)
 class LambdaLogFilter(logging.Filter):
@@ -169,15 +175,15 @@ def http_body_log_message(body_type: str,
                           body: bytes | bytearray | str | None,
                           *,
                           verbatim: bool = False,
-                          ) -> str:
+                          max_len: int = 128) -> str:
     if body is None:
         return f'… without {body_type} body'
     elif isinstance(body, (bytes, bytearray, str)):
-        if verbatim:
+        if verbatim or max_len > 1028:
             if isinstance(body, (bytes, bytearray)):
                 body = body.decode(errors='ignore')
         else:
-            body = trunc_ellipses(body, max_len=128)
+            body = trunc_ellipses(body, max_len=max_len)
         return f'… with {body_type} body {body!r}'
     else:
         return f'… with nonprintable body ({type(body)!r})'
diff --git a/test/service/test_app_logging.py b/test/service/test_app_logging.py
index 2d4251f561..c0b62a7b9b 100644
--- a/test/service/test_app_logging.py
+++ b/test/service/test_app_logging.py
@@ -52,7 +52,12 @@ def test_request_logs(self):
                         (
                             INFO,
                             f"Received GET request for '/health/basic', "
-                            f"with {json.dumps({'query': None, 'headers': headers})}."),
+                            f"with {json.dumps({'query': None, 'headers': headers})}"
+                        ),
+                        (
+                            INFO,
+                            "… with request body ''"
+                        ),
                         (
                             INFO,
                             "Authenticated request as OAuth2(access_token='foo_token')"