From 6c1a4d2c9a23cd1c56d58ae5b12a62b3164be8ac Mon Sep 17 00:00:00 2001
From: Heitor Lessa <lessa@amazon.com>
Date: Mon, 4 Jul 2022 13:44:34 +0200
Subject: [PATCH] fix(logger): preserve std keys when using custom formatters
 (#1264)

---
 aws_lambda_powertools/logging/formatter.py |  5 ++++
 aws_lambda_powertools/logging/logger.py    | 16 +++++++---
 tests/functional/test_logger.py            | 34 ++++++++++++++++++++++
 3 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/aws_lambda_powertools/logging/formatter.py b/aws_lambda_powertools/logging/formatter.py
index becfc9de85..1f01015051 100644
--- a/aws_lambda_powertools/logging/formatter.py
+++ b/aws_lambda_powertools/logging/formatter.py
@@ -1,3 +1,4 @@
+import inspect
 import json
 import logging
 import os
@@ -286,3 +287,7 @@ def _strip_none_records(records: Dict[str, Any]) -> Dict[str, Any]:
 
 
 JsonFormatter = LambdaPowertoolsFormatter  # alias to previous formatter
+
+
+# Fetch current and future parameters from PowertoolsFormatter that should be reserved
+RESERVED_FORMATTER_CUSTOM_KEYS: List[str] = inspect.getfullargspec(LambdaPowertoolsFormatter).args[1:]
diff --git a/aws_lambda_powertools/logging/logger.py b/aws_lambda_powertools/logging/logger.py
index 157d53adf7..f70224caba 100644
--- a/aws_lambda_powertools/logging/logger.py
+++ b/aws_lambda_powertools/logging/logger.py
@@ -12,7 +12,7 @@
 from ..shared.functions import resolve_env_var_choice, resolve_truthy_env_var_choice
 from .exceptions import InvalidLoggerSamplingRateError
 from .filters import SuppressFilter
-from .formatter import BasePowertoolsFormatter, LambdaPowertoolsFormatter
+from .formatter import RESERVED_FORMATTER_CUSTOM_KEYS, BasePowertoolsFormatter, LambdaPowertoolsFormatter
 from .lambda_context import build_lambda_context_model
 
 logger = logging.getLogger(__name__)
@@ -82,7 +82,7 @@ class Logger(logging.Logger):  # lgtm [py/missing-call-to-init]
     datefmt: str, optional
         String directives (strftime) to format log timestamp using `time`, by default it uses RFC
         3339.
-    use_datetime_directive: str, optional
+    use_datetime_directive: bool, optional
         Interpret `datefmt` as a format string for `datetime.datetime.strftime`, rather than
         `time.strftime`.
 
@@ -368,7 +368,7 @@ def registered_handler(self) -> logging.Handler:
         return handlers[0]
 
     @property
-    def registered_formatter(self) -> PowertoolsFormatter:
+    def registered_formatter(self) -> BasePowertoolsFormatter:
         """Convenience property to access logger formatter"""
         return self.registered_handler.formatter  # type: ignore
 
@@ -395,7 +395,15 @@ def structure_logs(self, append: bool = False, **keys):
         is_logger_preconfigured = getattr(self._logger, "init", False)
         if not is_logger_preconfigured:
             formatter = self.logger_formatter or LambdaPowertoolsFormatter(**log_keys)  # type: ignore
-            return self.registered_handler.setFormatter(formatter)
+            self.registered_handler.setFormatter(formatter)
+
+            # when using a custom Lambda Powertools Formatter
+            # standard and custom keys that are not Powertools Formatter parameters should be appended
+            # and custom keys that might happen to be Powertools Formatter parameters should be discarded
+            # this prevents adding them as custom keys, for example, `json_default=<callable>`
+            # see https://github.com/awslabs/aws-lambda-powertools-python/issues/1263
+            custom_keys = {k: v for k, v in log_keys.items() if k not in RESERVED_FORMATTER_CUSTOM_KEYS}
+            return self.registered_formatter.append_keys(**custom_keys)
 
         # Mode 2 (legacy)
         if append:
diff --git a/tests/functional/test_logger.py b/tests/functional/test_logger.py
index c76faa7cde..c8b3dc6175 100644
--- a/tests/functional/test_logger.py
+++ b/tests/functional/test_logger.py
@@ -625,6 +625,40 @@ def handler(event, context):
     assert all(k in second_log for k in lambda_context_keys)
 
 
+def test_logger_custom_formatter_has_standard_and_custom_keys(stdout, service_name, lambda_context):
+    class CustomFormatter(LambdaPowertoolsFormatter):
+        ...
+
+    # GIVEN a Logger is initialized with a custom formatter
+    logger = Logger(service=service_name, stream=stdout, logger_formatter=CustomFormatter(), my_key="value")
+
+    # WHEN a lambda function is decorated with logger
+    @logger.inject_lambda_context
+    def handler(event, context):
+        logger.info("Hello")
+
+    handler({}, lambda_context)
+
+    standard_keys = (
+        "level",
+        "location",
+        "message",
+        "timestamp",
+        "service",
+        "cold_start",
+        "function_name",
+        "function_memory_size",
+        "function_arn",
+        "function_request_id",
+    )
+
+    log = capture_logging_output(stdout)
+
+    # THEN all standard keys should be available
+    assert all(k in log for k in standard_keys)
+    assert "my_key" in log
+
+
 def test_logger_custom_handler(lambda_context, service_name, tmp_path):
     # GIVEN a Logger is initialized with a FileHandler
     log_file = tmp_path / "log.json"