diff --git a/aws_lambda_powertools/logging/logger.py b/aws_lambda_powertools/logging/logger.py index b147824e2cf..88c903b7cb6 100644 --- a/aws_lambda_powertools/logging/logger.py +++ b/aws_lambda_powertools/logging/logger.py @@ -22,8 +22,6 @@ overload, ) -import jmespath - from aws_lambda_powertools.logging import compat from aws_lambda_powertools.shared import constants from aws_lambda_powertools.shared.functions import ( @@ -31,6 +29,7 @@ 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") diff --git a/docs/core/logger.md b/docs/core/logger.md index de968a4ddc7..8c915fcd589 100644 --- a/docs/core/logger.md +++ b/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" diff --git a/tests/functional/test_logger.py b/tests/functional/test_logger.py index b7eeed779ce..dbe2ed1917f 100644 --- a/tests/functional/test_logger.py +++ b/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)