fix(logger): preserve std keys when using custom formatters (#1264)

Author: heitorlessa

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:]

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:

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"