feat(logger): add clear_state method (#5956)

* feat(logger): add clear_state method to Logger

* changes after feedback

* change reset default keys

---------

Co-authored-by: Leandro Damascena <[email protected]>

aws_lambda_powertools/logging/logger.py

@@ -222,6 +222,7 @@ def __init__(
             choice=sampling_rate,
             env=os.getenv(constants.LOGGER_LOG_SAMPLING_RATE),
         )
+        self._default_log_keys: dict[str, Any] = {"service": self.service, "sampling_rate": self.sampling_rate}
         self.child = child
         self.logger_formatter = logger_formatter
         self._stream = stream or sys.stdout
@@ -231,7 +232,6 @@ def __init__(
         self._is_deduplication_disabled = resolve_truthy_env_var_choice(
             env=os.getenv(constants.LOGGER_LOG_DEDUPLICATION_ENV, "false"),
         )
-        self._default_log_keys = {"service": self.service, "sampling_rate": self.sampling_rate}
         self._logger = self._get_logger()
 
         # NOTE: This is primarily to improve UX, so IDEs can autocomplete LambdaPowertoolsFormatter options
@@ -605,6 +605,14 @@ def append_context_keys(self, **additional_keys: Any) -> Generator[None, None, N
         with self.registered_formatter.append_context_keys(**additional_keys):
             yield
 
+    def clear_state(self) -> None:
+        """Removes all custom keys that were appended to the Logger."""
+        # Clear all custom keys from the formatter
+        self.registered_formatter.clear_state()
+
+        # Reset to default keys
+        self.structure_logs(**self._default_log_keys)
+
     # These specific thread-safe methods are necessary to manage shared context in concurrent environments.
     # They prevent race conditions and ensure data consistency across multiple threads.
     def thread_safe_append_keys(self, **additional_keys: object) -> None:

docs/core/logger.md

@@ -274,13 +274,15 @@ You can remove any additional key from Logger state using `remove_keys`.
 
 #### Clearing all state
 
+##### Decorator with clear_state
+
 Logger is commonly initialized in the global scope. Due to [Lambda Execution Context reuse](https://docs.aws.amazon.com/lambda/latest/dg/runtimes-context.html){target="_blank"}, this means that custom keys can be persisted across invocations. If you want all custom keys to be deleted, you can use `clear_state=True` param in `inject_lambda_context` decorator.
 
 ???+ tip "Tip: When is this useful?"
     It is useful when you add multiple custom keys conditionally, instead of setting a default `None` value if not present. Any key with `None` value is automatically removed by Logger.
 
 ???+ danger "Danger: This can have unintended side effects if you use Layers"
-    Lambda Layers code is imported before the Lambda handler.
+    Lambda Layers code is imported before the Lambda handler. When a Lambda function starts, it first imports and executes all code in the Layers (including any global scope code) before proceeding to the function's own code.
 
     This means that `clear_state=True` will instruct Logger to remove any keys previously added before Lambda handler execution proceeds.
 
@@ -304,6 +306,27 @@ Logger is commonly initialized in the global scope. Due to [Lambda Execution Con
     --8<-- "examples/logger/src/clear_state_event_two.json"
     ```
 
+##### clear_state method
+
+You can call `clear_state()` as a method explicitly within your code to clear appended keys at any point during the execution of your Lambda invocation.
+
+=== "clear_state_method.py"
+
+    ```python hl_lines="12"
+    --8<-- "examples/logger/src/clear_state_method.py"
+    ```
+=== "Output before clear_state()"
+
+    ```json hl_lines="9 17"
+    --8<-- "examples/logger/src/before_clear_state.json"
+    ```
+
+=== "Output after clear_state()"
+
+    ```json hl_lines="4"
+    --8<-- "examples/logger/src/after_clear_state.json"
+    ```
+
 ### Accessing currently configured keys
 
 You can view all currently configured keys from the Logger state using the `get_current_keys()` method. This method is useful when you need to avoid overwriting keys that are already configured.

examples/logger/src/after_clear_state.json

@@ -0,0 +1,7 @@
+{
+    "level": "INFO",
+    "location": "lambda_handler:126",
+    "message": "State after clearing - only show default keys",
+    "timestamp": "2025-01-30 13:56:03,158-0300",
+    "service": "payment"
+}

examples/logger/src/before_clear_state.json

@@ -0,0 +1,20 @@
+{
+    "logs": [
+        {
+            "level": "INFO",
+            "location": "lambda_handler:122",
+            "message": "Starting order processing",
+            "timestamp": "2025-01-30 13:56:03,157-0300",
+            "service": "payment",
+            "order_id": "12345"
+        },
+        {
+            "level": "INFO",
+            "location": "lambda_handler:124",
+            "message": "Final state before clearing",
+            "timestamp": "2025-01-30 13:56:03,157-0300",
+            "service": "payment",
+            "order_id": "12345"
+        }
+    ]
+}

examples/logger/src/clear_state_method.py

@@ -0,0 +1,15 @@
+from aws_lambda_powertools import Logger
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+logger = Logger(service="payment", level="DEBUG")
+
+
+def lambda_handler(event: dict, context: LambdaContext) -> str:
+    try:
+        logger.append_keys(order_id="12345")
+        logger.info("Starting order processing")
+    finally:
+        logger.info("Final state before clearing")
+        logger.clear_state()
+        logger.info("State after clearing - only show default keys")
+    return "Completed"

tests/functional/logger/required_dependencies/test_logger.py

@@ -1232,3 +1232,57 @@ def test_logger_change_level_child_logger(stdout, service_name):
     logs = list(stdout.getvalue().strip().split("\n"))
     assert len(logs) == 1
     assert "service" in logs[0]
+
+
+def test_clear_state_with_append_keys():
+    # GIVEN a Logger is initialized
+    logger = Logger(service="service_name", stream=stdout)
+
+    # WHEN append keys are added
+    logger.append_keys(custom_key="custom_key")
+    logger.info("message with appended keys")
+    logger.clear_state()
+
+    # THEN context keys should be cleared
+    assert "custom_key" not in logger.get_current_keys()
+
+
+def test_clear_state(stdout, service_name):
+    # GIVEN a Logger is initialized
+    logger = Logger(service=service_name, stream=stdout)
+    logger.info("message for the user")
+
+    # WHEN the clear_state method is called
+    logger.clear_state()
+
+    # THEN the logger's current keys should be reset to their default values
+    expected_keys = {
+        "level": "%(levelname)s",
+        "location": "%(funcName)s:%(lineno)d",
+        "message": None,
+        "timestamp": "%(asctime)s",
+        "service": service_name,
+        "sampling_rate": None,
+    }
+    assert logger.get_current_keys() == expected_keys
+
+
+def test_clear_state_log_output(stdout, service_name):
+    # GIVEN a Logger is initialized
+    logger = Logger(service=service_name, stream=stdout)
+
+    # WHEN we append a custom key and log
+    logger.append_keys(custom_key="test_value")
+    logger.info("first message")
+
+    # AND we clear the state and log again
+    logger.clear_state()
+    logger.info("second message")
+
+    # THEN the first log should contain the custom key
+    # AND the second log should not contain the custom key
+    first_log, second_log = capture_multiple_logging_statements_output(stdout)
+
+    assert "custom_key" in first_log
+    assert first_log["custom_key"] == "test_value"
+    assert "custom_key" not in second_log