feat(logging): Add correlation_id_path option

Author: Michael Brewer

aws_lambda_powertools/logging/logger.py

@@ -6,6 +6,8 @@
 import sys
 from typing import Any, Callable, Dict, Union
 
+import jmespath
+
 from ..shared import constants
 from ..shared.functions import resolve_env_var_choice, resolve_truthy_env_var_choice
 from .exceptions import InvalidLoggerSamplingRateError
@@ -204,7 +206,9 @@ def _configure_sampling(self):
                 f"Please review POWERTOOLS_LOGGER_SAMPLE_RATE environment variable."
             )
 
-    def inject_lambda_context(self, lambda_handler: Callable[[Dict, Any], Any] = None, log_event: bool = None):
+    def inject_lambda_context(
+        self, lambda_handler: Callable[[Dict, Any], Any] = None, log_event: bool = None, correlation_id_path: str = None
+    ):
         """Decorator to capture Lambda contextual info and inject into logger
 
         Parameters
@@ -213,6 +217,8 @@ def inject_lambda_context(self, lambda_handler: Callable[[Dict, Any], Any] = Non
             Method to inject the lambda context
         log_event : bool, optional
             Instructs logger to log Lambda Event, by default False
+        correlation_id_path: str, optional
+            Optional JMESPath for the correlation_id
 
         Environment variables
         ---------------------
@@ -251,7 +257,9 @@ def handler(event, context):
         # Return a partial function with args filled
         if lambda_handler is None:
             logger.debug("Decorator called with parameters")
-            return functools.partial(self.inject_lambda_context, log_event=log_event)
+            return functools.partial(
+                self.inject_lambda_context, log_event=log_event, correlation_id_path=correlation_id_path
+            )
 
         log_event = resolve_truthy_env_var_choice(
             choice=log_event, env=os.getenv(constants.LOGGER_LOG_EVENT_ENV, "false")
@@ -263,6 +271,9 @@ def decorate(event, context):
             cold_start = _is_cold_start()
             self.structure_logs(append=True, cold_start=cold_start, **lambda_context.__dict__)
 
+            if correlation_id_path:
+                self.set_correlation_id(jmespath.search(correlation_id_path, event))
+
             if log_event:
                 logger.debug("Event received")
                 self.info(event)

tests/functional/test_logger.py

@@ -439,7 +439,7 @@ def test_logger_exception_extract_exception_name(stdout, service_name):
     assert "ValueError" == log["exception_name"]
 
 
-def test_logger_correlation_id(lambda_context, stdout, service_name):
+def test_logger_set_correlation_id(lambda_context, stdout, service_name):
     # GIVEN
     logger = Logger(service=service_name, stream=stdout)
     request_id = "xxx-111-222"
@@ -455,3 +455,21 @@ def handler(event, _):
     # THEN
     log = capture_logging_output(stdout)
     assert request_id == log["correlation_id"]
+
+
+def test_logger_set_correlation_id_path(lambda_context, stdout, service_name):
+    # GIVEN
+    logger = Logger(service=service_name, stream=stdout)
+    request_id = "xxx-111-222"
+    mock_event = {"requestContext": {"requestId": request_id}}
+
+    @logger.inject_lambda_context(correlation_id_path="requestContext.requestId")
+    def handler(event, context):
+        logger.info("Foo")
+
+    # WHEN
+    handler(mock_event, lambda_context)
+
+    # THEN
+    log = capture_logging_output(stdout)
+    assert request_id == log["correlation_id"]