fix(logger): ensure state is cleared for all formatters

Author: heitorlessa

aws_lambda_powertools/logging/formatter.py

@@ -45,6 +45,10 @@ def append_keys(self, **additional_keys):
     def remove_keys(self, keys: Iterable[str]):
         raise NotImplementedError()
 
+    def clear_state(self):
+        """Removes any previously added logging keys"""
+        raise NotImplementedError()
+
 
 class LambdaPowertoolsFormatter(BasePowertoolsFormatter):
     """AWS Lambda Powertools Logging formatter.
@@ -180,6 +184,9 @@ def remove_keys(self, keys: Iterable[str]):
         for key in keys:
             self.log_format.pop(key, None)
 
+    def clear_state(self):
+        self.log_format = dict.fromkeys(self.log_record_order)
+
     @staticmethod
     def _build_default_keys():
         return {

aws_lambda_powertools/logging/logger.py

@@ -385,14 +385,29 @@ def structure_logs(self, append: bool = False, **keys):
         append : bool, optional
             append keys provided to logger formatter, by default False
         """
+        # There are 3 operational modes for this method
+        ## 1. Append new keys to the current logger formatter; deprecated in favour of append_keys
+        ## 2. Register a Powertools Formatter for the first time
+        ## 3. Add new keys and discard existing to the registered formatter
 
         if append:
-            # Maintenance: Add deprecation warning for major version. Refer to append_keys() when docs are updated
-            self.append_keys(**keys)
-        else:
-            log_keys = {**self._default_log_keys, **keys}
+            # Maintenance: Add deprecation warning for major version
+            return self.append_keys(**keys)
+
+        log_keys = {**self._default_log_keys, **keys}
+
+        # Behaviour 2
+        is_logger_preconfigured = getattr(self._logger, "init", False)
+        if not is_logger_preconfigured:
             formatter = self.logger_formatter or LambdaPowertoolsFormatter(**log_keys)  # type: ignore
-            self.registered_handler.setFormatter(formatter)
+            return self.registered_handler.setFormatter(formatter)
+
+        # Behaviour 3
+        try:
+            self.registered_formatter.clear_state()
+            self.registered_formatter.append_keys(**log_keys)
+        except (AttributeError, NotImplementedError):
+            logger.warning(f"Formatter {self.registered_formatter} doesn't implement clear_state method; ignoring...")
 
     def set_correlation_id(self, value: Optional[str]):
         """Sets the correlation_id in the logging json
@@ -438,6 +453,9 @@ def _get_caller_filename():
         caller_frame = frame.f_back.f_back.f_back
         return caller_frame.f_globals["__name__"]
 
+    def _reset_logger_state(self):
+        self.registered_formatter.clear_state()
+
 
 def set_package_logger(
     level: Union[str, int] = logging.DEBUG,

tests/functional/test_logger.py

@@ -5,16 +5,17 @@
 import random
 import re
 import string
+from ast import Dict
 from collections import namedtuple
 from datetime import datetime, timezone
-from typing import Iterable
+from typing import Any, Callable, Iterable, List, Optional, Union
 
 import pytest
 
 from aws_lambda_powertools import Logger, Tracer
 from aws_lambda_powertools.logging import correlation_paths
 from aws_lambda_powertools.logging.exceptions import InvalidLoggerSamplingRateError
-from aws_lambda_powertools.logging.formatter import BasePowertoolsFormatter
+from aws_lambda_powertools.logging.formatter import BasePowertoolsFormatter, LambdaPowertoolsFormatter
 from aws_lambda_powertools.logging.logger import set_package_logger
 from aws_lambda_powertools.shared import constants
 from aws_lambda_powertools.utilities.data_classes import S3Event, event_source
@@ -564,6 +565,63 @@ def handler(event, context):
     assert logger.get_correlation_id() is None
 
 
+def test_logger_custom_powertools_formatter_clear_state(stdout, service_name, lambda_context):
+    class CustomFormatter(LambdaPowertoolsFormatter):
+        def __init__(
+            self,
+            json_serializer: Optional[Callable[[Dict], str]] = None,
+            json_deserializer: Optional[Callable[[Union[Dict, str, bool, int, float]], str]] = None,
+            json_default: Optional[Callable[[Any], Any]] = None,
+            datefmt: Optional[str] = None,
+            use_datetime_directive: bool = False,
+            log_record_order: Optional[List[str]] = None,
+            utc: bool = False,
+            **kwargs,
+        ):
+            super().__init__(
+                json_serializer,
+                json_deserializer,
+                json_default,
+                datefmt,
+                use_datetime_directive,
+                log_record_order,
+                utc,
+                **kwargs,
+            )
+
+    custom_formatter = CustomFormatter()
+
+    # GIVEN a Logger is initialized with a custom formatter
+    logger = Logger(service=service_name, stream=stdout, logger_formatter=custom_formatter)
+
+    # WHEN a lambda function is decorated with logger
+    # and state is to be cleared in the next invocation
+    @logger.inject_lambda_context(clear_state=True)
+    def handler(event, context):
+        if event.get("add_key"):
+            logger.append_keys(my_key="value")
+        logger.info("Hello")
+
+    handler({"add_key": True}, lambda_context)
+    handler({}, lambda_context)
+
+    lambda_context_keys = (
+        "function_name",
+        "function_memory_size",
+        "function_arn",
+        "function_request_id",
+    )
+
+    first_log, second_log = capture_multiple_logging_statements_output(stdout)
+
+    # THEN my_key should only present once
+    # and lambda contextual info should also be in both logs
+    assert "my_key" in first_log
+    assert "my_key" not in second_log
+    assert all(k in first_log for k in lambda_context_keys)
+    assert all(k in second_log for k in lambda_context_keys)
+
+
 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"