fix(logger): ensure state is cleared for custom formatters (aws-powertools#1072)

Author: heitorlessa

aws_lambda_powertools/logging/formatter.py

@@ -45,6 +45,11 @@ def append_keys(self, **additional_keys):
     def remove_keys(self, keys: Iterable[str]):
         raise NotImplementedError()
 
+    @abstractmethod
+    def clear_state(self):
+        """Removes any previously added logging keys"""
+        raise NotImplementedError()
+
 
 class LambdaPowertoolsFormatter(BasePowertoolsFormatter):
     """AWS Lambda Powertools Logging formatter.
@@ -180,6 +185,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,26 @@ 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. Register a Powertools Formatter for the first time
+        ## 2. Append new keys to the current logger formatter; deprecated in favour of append_keys
+        ## 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}
+        # Mode 1
+        log_keys = {**self._default_log_keys, **keys}
+        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)
+
+        # Mode 2 (legacy)
+        if append:
+            # Maintenance: Add deprecation warning for major version
+            return self.append_keys(**keys)
+
+        # Mode 3
+        self.registered_formatter.clear_state()
+        self.registered_formatter.append_keys(**log_keys)
 
     def set_correlation_id(self, value: Optional[str]):
         """Sets the correlation_id in the logging json

docs/core/logger.md

@@ -862,7 +862,10 @@ logger.info("Collecting payment")
 
 By default, Logger uses [LambdaPowertoolsFormatter](#lambdapowertoolsformatter) that persists its custom structure between non-cold start invocations. There could be scenarios where the existing feature set isn't sufficient to your formatting needs.
 
-For **minor changes like remapping keys** after all log record processing has completed, you can override `serialize` method from [LambdaPowertoolsFormatter](#lambdapowertoolsformatter):
+???+ info
+    The most common use cases are remapping keys by bringing your existing schema, and redacting sensitive information you know upfront.
+
+For these, you can override the `serialize` method from [LambdaPowertoolsFormatter](#lambdapowertoolsformatter).
 
 === "custom_formatter.py"
 
@@ -892,28 +895,39 @@ For **minor changes like remapping keys** after all log record processing has co
 	}
 	```
 
-For **replacing the formatter entirely**, you can subclass `BasePowertoolsFormatter`, implement `append_keys` method, and override `format` standard logging method. This ensures the current feature set of Logger like [injecting Lambda context](#capturing-lambda-context-info) and [sampling](#sampling-debug-logs) will continue to work.
+The `log` argument is the final log record containing [our standard keys](#standard-structured-keys), optionally [Lambda context keys](#capturing-lambda-context-info), and any custom key you might have added via [append_keys](#append_keys-method) or the [extra parameter](#extra-parameter).
+
+For exceptional cases where you want to completely replace our formatter logic, you can subclass `BasePowertoolsFormatter`.
+
+???+ warning
+    You will need to implement `append_keys`, `clear_state`, override `format`, and optionally `remove_keys` to keep the same feature set Powertools Logger provides. This also means keeping state of logging keys added.
 
-???+ info
-	You might need to implement `remove_keys` method if you make use of the feature too.
 
 === "collect.py"
 
-    ```python hl_lines="2 4 7 12 16 27"
+    ```python hl_lines="5 7 9-10 13 17 21 24 35"
+    import logging
+    from typing import Iterable, List, Optional
+
     from aws_lambda_powertools import Logger
     from aws_lambda_powertools.logging.formatter import BasePowertoolsFormatter
 
     class CustomFormatter(BasePowertoolsFormatter):
-        custom_format = {}  # arbitrary dict to hold our structured keys
+        def __init__(self, log_record_order: Optional[List[str]], *args, **kwargs):
+            self.log_record_order = log_record_order or ["level", "location", "message", "timestamp"]
+            self.log_format = dict.fromkeys(self.log_record_order)
+            super().__init__(*args, **kwargs)
 
         def append_keys(self, **additional_keys):
             # also used by `inject_lambda_context` decorator
-            self.custom_format.update(additional_keys)
+            self.log_format.update(additional_keys)
 
-        # Optional unless you make use of this Logger feature
         def remove_keys(self, keys: Iterable[str]):
             for key in keys:
-                self.custom_format.pop(key, None)
+                self.log_format.pop(key, None)
+
+        def clear_state(self):
+            self.log_format = dict.fromkeys(self.log_record_order)
 
         def format(self, record: logging.LogRecord) -> str:  # noqa: A003
             """Format logging record as structured JSON str"""
@@ -922,7 +936,7 @@ For **replacing the formatter entirely**, you can subclass `BasePowertoolsFormat
                     "event": super().format(record),
                     "timestamp": self.formatTime(record),
                     "my_default_key": "test",
-                    **self.custom_format,
+                    **self.log_format,
                 }
             )
 

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
@@ -524,6 +525,9 @@ def remove_keys(self, keys: Iterable[str]):
             for key in keys:
                 self.custom_format.pop(key, None)
 
+        def clear_state(self):
+            self.custom_format.clear()
+
         def format(self, record: logging.LogRecord) -> str:  # noqa: A003
             return json.dumps(
                 {
@@ -564,6 +568,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"