diff --git a/aws_lambda_powertools/logging/logger.py b/aws_lambda_powertools/logging/logger.py index 3231f30eccd..689409d9813 100644 --- a/aws_lambda_powertools/logging/logger.py +++ b/aws_lambda_powertools/logging/logger.py @@ -260,12 +260,18 @@ def _configure_sampling(self): ) def inject_lambda_context( - self, lambda_handler: Callable[[Dict, Any], Any] = None, log_event: bool = None, correlation_id_path: str = None + self, + lambda_handler: Callable[[Dict, Any], Any] = None, + log_event: bool = None, + correlation_id_path: str = None, + clear_state: bool = False, ): """Decorator to capture Lambda contextual info and inject into logger Parameters ---------- + clear_state : bool, optional + Instructs logger to remove any custom keys previously added lambda_handler : Callable Method to inject the lambda context log_event : bool, optional @@ -311,7 +317,10 @@ def handler(event, context): if lambda_handler is None: logger.debug("Decorator called with parameters") return functools.partial( - self.inject_lambda_context, log_event=log_event, correlation_id_path=correlation_id_path + self.inject_lambda_context, + log_event=log_event, + correlation_id_path=correlation_id_path, + clear_state=clear_state, ) log_event = resolve_truthy_env_var_choice( @@ -322,7 +331,11 @@ def handler(event, context): def decorate(event, context): lambda_context = build_lambda_context_model(context) cold_start = _is_cold_start() - self.append_keys(cold_start=cold_start, **lambda_context.__dict__) + + if clear_state: + self.structure_logs(cold_start=cold_start, **lambda_context.__dict__) + else: + self.append_keys(cold_start=cold_start, **lambda_context.__dict__) if correlation_id_path: self.set_correlation_id(jmespath.search(correlation_id_path, event)) diff --git a/docs/core/logger.md b/docs/core/logger.md index a544bf91e4b..45119ca51d6 100644 --- a/docs/core/logger.md +++ b/docs/core/logger.md @@ -231,8 +231,9 @@ We provide [built-in JMESPath expressions](#built-in-correlation-id-expressions) ### Appending additional keys -!!! info "Keys might be persisted across invocations" - Always set additional keys as part of your handler to ensure they have the latest value. Additional keys are kept in memory as part of a Logger instance and might be reused in non-cold start scenarios. +!!! info "Custom keys are persisted across warm invocations" + Always set additional keys as part of your handler to ensure they have the latest value, or explicitly clear them with [`clear_state=True`](#clearing-all-state). + You can append additional keys using either mechanism: @@ -426,6 +427,73 @@ You can remove any additional key from Logger state using `remove_keys`. } ``` +#### Clearing all 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), 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. + +!!! info + This 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 "This can have unintended side effects if you use Layers" + Lambda Layers code is imported before the Lambda handler. + + This means that `clear_state=True` will instruct Logger to remove any keys previously added before Lambda handler execution proceeds. + + You can either avoid running any code as part of Lambda Layers global scope, or override keys with their latest value as part of handler's execution. + +=== "collect.py" + + ```python hl_lines="5 8" + from aws_lambda_powertools import Logger + + logger = Logger(service="payment") + + @logger.inject_lambda_context(clear_state=True) + def handler(event, context): + if event.get("special_key"): + # Should only be available in the first request log + # as the second request doesn't contain `special_key` + logger.append_keys(debugging_key="value") + + logger.info("Collecting payment") + ``` + +=== "#1 request" + + ```json hl_lines="7" + { + "level": "INFO", + "location": "collect.handler:10", + "message": "Collecting payment", + "timestamp": "2021-05-03 11:47:12,494+0200", + "service": "payment", + "special_key": "debug_key", + "cold_start": true, + "lambda_function_name": "test", + "lambda_function_memory_size": 128, + "lambda_function_arn": "arn:aws:lambda:eu-west-1:12345678910:function:test", + "lambda_request_id": "52fdfc07-2182-154f-163f-5f0f9a621d72" + } + ``` + +=== "#2 request" + + ```json hl_lines="7" + { + "level": "INFO", + "location": "collect.handler:10", + "message": "Collecting payment", + "timestamp": "2021-05-03 11:47:12,494+0200", + "service": "payment", + "cold_start": false, + "lambda_function_name": "test", + "lambda_function_memory_size": 128, + "lambda_function_arn": "arn:aws:lambda:eu-west-1:12345678910:function:test", + "lambda_request_id": "52fdfc07-2182-154f-163f-5f0f9a621d72" + } + ``` + + ### Logging exceptions Use `logger.exception` method to log contextual information about exceptions. Logger will include `exception_name` and `exception` keys to aid troubleshooting and error enumeration. diff --git a/tests/functional/test_logger.py b/tests/functional/test_logger.py index ba6e82b72af..44249af6250 100644 --- a/tests/functional/test_logger.py +++ b/tests/functional/test_logger.py @@ -562,3 +562,23 @@ def handler(event, context): # THEN we should output to a file not stdout log = log_file.read_text() assert "custom handler" in log + + +def test_clear_state_on_inject_lambda_context(lambda_context, stdout, service_name): + # GIVEN + logger = Logger(service=service_name, stream=stdout) + + # WHEN clear_state is set and a key was conditionally added in the first 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("Foo") + + # THEN custom key should only exist in the first log + handler({"add_key": True}, lambda_context) + handler({}, lambda_context) + + first_log, second_log = capture_multiple_logging_statements_output(stdout) + assert "my_key" in first_log + assert "my_key" not in second_log