diff --git a/aws_lambda_powertools/logging/logger.py b/aws_lambda_powertools/logging/logger.py index ab29b6af649..61bfb82867b 100644 --- a/aws_lambda_powertools/logging/logger.py +++ b/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: diff --git a/docs/core/logger.md b/docs/core/logger.md index 9915f7cc4b4..27fb532ad00 100644 --- a/docs/core/logger.md +++ b/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. diff --git a/examples/logger/src/after_clear_state.json b/examples/logger/src/after_clear_state.json new file mode 100644 index 00000000000..54dd72ed41e --- /dev/null +++ b/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" +} \ No newline at end of file diff --git a/examples/logger/src/before_clear_state.json b/examples/logger/src/before_clear_state.json new file mode 100644 index 00000000000..a710dbde0d6 --- /dev/null +++ b/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" + } + ] +} \ No newline at end of file diff --git a/examples/logger/src/clear_state_method.py b/examples/logger/src/clear_state_method.py new file mode 100644 index 00000000000..1564b4c2c39 --- /dev/null +++ b/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" diff --git a/tests/functional/logger/required_dependencies/test_logger.py b/tests/functional/logger/required_dependencies/test_logger.py index 5c85677d73d..92c8e27ff4b 100644 --- a/tests/functional/logger/required_dependencies/test_logger.py +++ b/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