fix(logger): allow custom JMESPath functions to extract correlation ID (#3382)

Author: heitorlessa

aws_lambda_powertools/logging/logger.py

@@ -22,15 +22,14 @@
     overload,
 )
 
-import jmespath
-
 from aws_lambda_powertools.logging import compat
 from aws_lambda_powertools.shared import constants
 from aws_lambda_powertools.shared.functions import (
     extract_event_from_common_models,
     resolve_env_var_choice,
     resolve_truthy_env_var_choice,
 )
+from aws_lambda_powertools.utilities import jmespath_utils
 
 from ..shared.types import AnyCallableT
 from .exceptions import InvalidLoggerSamplingRateError
@@ -443,7 +442,9 @@ def decorate(event, context, *args, **kwargs):
                 self.append_keys(cold_start=cold_start, **lambda_context.__dict__)
 
             if correlation_id_path:
-                self.set_correlation_id(jmespath.search(correlation_id_path, event))
+                self.set_correlation_id(
+                    jmespath_utils.extract_data_from_envelope(envelope=correlation_id_path, data=event),
+                )
 
             if log_event:
                 logger.debug("Event received")

docs/core/logger.md

@@ -85,10 +85,10 @@ When debugging in non-production environments, you can instruct Logger to log th
 
 ### Setting a Correlation ID
 
-You can set a Correlation ID using `correlation_id_path` param by passing a [JMESPath expression](https://jmespath.org/tutorial.html){target="_blank" rel="nofollow"}.
+You can set a Correlation ID using `correlation_id_path` param by passing a [JMESPath expression](https://jmespath.org/tutorial.html){target="_blank" rel="nofollow"}, including [our custom JMESPath Functions](../utilities/jmespath_functions.md#powertools_json-function).
 
 ???+ tip
-	You can retrieve correlation IDs via `get_correlation_id` method
+	You can retrieve correlation IDs via `get_correlation_id` method.
 
 === "set_correlation_id.py"
 

tests/functional/test_logger.py

@@ -569,6 +569,24 @@ def handler(event, context):
     assert request_id == log["correlation_id"]
 
 
+def test_logger_set_correlation_id_path_custom_functions(lambda_context, stdout, service_name):
+    # GIVEN an initialized Logger
+    # AND a Lambda handler decorated with a JMESPath expression using Powertools custom functions
+    logger = Logger(service=service_name, stream=stdout)
+
+    @logger.inject_lambda_context(correlation_id_path="Records[*].powertools_json(body).id")
+    def handler(event, context):
+        ...
+
+    # WHEN handler is called
+    request_id = "xxx-111-222"
+    mock_event = {"Records": [{"body": json.dumps({"id": request_id})}]}
+    handler(mock_event, lambda_context)
+
+    # THEN there should be no exception and correlation ID should match
+    assert logger.get_correlation_id() == [request_id]
+
+
 def test_logger_append_remove_keys(stdout, service_name):
     # GIVEN a Logger is initialized
     logger = Logger(service=service_name, stream=stdout)