feat(logger): add option to remove correlation_id

Call remove_keys when the correlation_id is set to None

Author: Michael Brewer

aws_lambda_powertools/logging/logger.py

@@ -387,15 +387,18 @@ def structure_logs(self, append: bool = False, **keys):
             formatter = self.logger_formatter or LambdaPowertoolsFormatter(**log_keys)
             self.registered_handler.setFormatter(formatter)
 
-    def set_correlation_id(self, value: str):
+    def set_correlation_id(self, value: Optional[str]):
         """Sets the correlation_id in the logging json
 
         Parameters
         ----------
-        value : str
-            Value for the correlation id
+        value : str, optional
+            Value for the correlation id. None will remove the correlation_id
         """
-        self.append_keys(correlation_id=value)
+        if value is None:
+            self.remove_keys(["correlation_id"])
+        else:
+            self.append_keys(correlation_id=value)
 
     @staticmethod
     def _get_log_level(level: Union[str, int, None]) -> Union[str, int]:

tests/functional/test_logger.py

@@ -460,6 +460,20 @@ def handler(event, _):
     assert request_id == log["correlation_id"]
 
 
+def test_logger_set_correlation_id_to_none(lambda_context, stdout, service_name):
+    # GIVEN a logger with a correlation_id set
+    logger = Logger(service=service_name, stream=stdout)
+    logger.set_correlation_id("foo")
+
+    # WHEN calling set_correlation_id with None
+    logger.set_correlation_id(None)
+
+    # THEN there should be no correlation_id
+    logger.info("Foo")
+    log = capture_logging_output(stdout)
+    assert "correlation_id" not in log
+
+
 def test_logger_set_correlation_id_path(lambda_context, stdout, service_name):
     # GIVEN
     logger = Logger(service=service_name, stream=stdout)