feat(logger): add option to clear state per invocation (#467)

Author: heitorlessa

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))

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.

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