diff --git a/aws_lambda_powertools/utilities/validation/jmespath_functions.py b/aws_lambda_powertools/shared/jmespath_functions.py similarity index 100% rename from aws_lambda_powertools/utilities/validation/jmespath_functions.py rename to aws_lambda_powertools/shared/jmespath_functions.py diff --git a/aws_lambda_powertools/utilities/idempotency/persistence/base.py b/aws_lambda_powertools/utilities/idempotency/persistence/base.py index c866d75d98d..352ba40b5f6 100644 --- a/aws_lambda_powertools/utilities/idempotency/persistence/base.py +++ b/aws_lambda_powertools/utilities/idempotency/persistence/base.py @@ -14,6 +14,7 @@ import jmespath from aws_lambda_powertools.shared.cache_dict import LRUDict +from aws_lambda_powertools.shared.jmespath_functions import PowertoolsFunctions from aws_lambda_powertools.shared.json_encoder import Encoder from aws_lambda_powertools.utilities.idempotency.exceptions import ( IdempotencyInvalidStatusError, @@ -115,6 +116,7 @@ def __init__( local_cache_max_items: int = 256, hash_function: str = "md5", raise_on_no_idempotency_key: bool = False, + jmespath_options: Dict = None, ) -> None: """ Initialize the base persistence layer @@ -135,6 +137,8 @@ def __init__( Function to use for calculating hashes, by default md5. raise_on_no_idempotency_key: bool, optional Raise exception if no idempotency key was found in the request, by default False + jmespath_options : Dict + Alternative JMESPath options to be included when filtering expr """ self.event_key_jmespath = event_key_jmespath if self.event_key_jmespath: @@ -149,6 +153,9 @@ def __init__( self.payload_validation_enabled = True self.hash_function = getattr(hashlib, hash_function) self.raise_on_no_idempotency_key = raise_on_no_idempotency_key + if not jmespath_options: + jmespath_options = {"custom_functions": PowertoolsFunctions()} + self.jmespath_options = jmespath_options def _get_hashed_idempotency_key(self, lambda_event: Dict[str, Any]) -> str: """ @@ -166,8 +173,11 @@ def _get_hashed_idempotency_key(self, lambda_event: Dict[str, Any]) -> str: """ data = lambda_event + if self.event_key_jmespath: - data = self.event_key_compiled_jmespath.search(lambda_event) + data = self.event_key_compiled_jmespath.search( + lambda_event, options=jmespath.Options(**self.jmespath_options) + ) if self.is_missing_idempotency_key(data): warnings.warn(f"No value found for idempotency_key. jmespath: {self.event_key_jmespath}") diff --git a/aws_lambda_powertools/utilities/validation/base.py b/aws_lambda_powertools/utilities/validation/base.py index bacd25a4efa..a5c82503735 100644 --- a/aws_lambda_powertools/utilities/validation/base.py +++ b/aws_lambda_powertools/utilities/validation/base.py @@ -5,8 +5,9 @@ import jmespath from jmespath.exceptions import LexerError +from aws_lambda_powertools.shared.jmespath_functions import PowertoolsFunctions + from .exceptions import InvalidEnvelopeExpressionError, InvalidSchemaFormatError, SchemaValidationError -from .jmespath_functions import PowertoolsFunctions logger = logging.getLogger(__name__) diff --git a/tests/functional/idempotency/conftest.py b/tests/functional/idempotency/conftest.py index 68492648337..9ae030f02d1 100644 --- a/tests/functional/idempotency/conftest.py +++ b/tests/functional/idempotency/conftest.py @@ -9,6 +9,7 @@ import pytest from botocore import stub from botocore.config import Config +from jmespath import functions from aws_lambda_powertools.shared.json_encoder import Encoder from aws_lambda_powertools.utilities.idempotency import DynamoDBPersistenceLayer @@ -180,6 +181,23 @@ def persistence_store_with_validation(config, request, default_jmespath): return persistence_store +@pytest.fixture +def persistence_store_with_jmespath_options(config, request): + class CustomFunctions(functions.Functions): + @functions.signature({"types": ["string"]}) + def _func_echo_decoder(self, value): + return value + + persistence_store = DynamoDBPersistenceLayer( + table_name=TABLE_NAME, + boto_config=config, + use_local_cache=False, + event_key_jmespath=request.param, + jmespath_options={"custom_functions": CustomFunctions()}, + ) + return persistence_store + + @pytest.fixture def mock_function(): return mock.MagicMock() diff --git a/tests/functional/idempotency/test_idempotency.py b/tests/functional/idempotency/test_idempotency.py index 675d8d1c5c8..872d9f39365 100644 --- a/tests/functional/idempotency/test_idempotency.py +++ b/tests/functional/idempotency/test_idempotency.py @@ -3,6 +3,7 @@ import sys from hashlib import md5 +import jmespath import pytest from botocore import stub @@ -690,3 +691,33 @@ def test_raise_on_no_idempotency_key(persistence_store): # THEN raise IdempotencyKeyError error assert "No data found to create a hashed idempotency_key" in str(excinfo.value) + + +@pytest.mark.parametrize("persistence_store", [{"use_local_cache": True}], indirect=True) +def test_jmespath_with_powertools_json(persistence_store): + # GIVEN an event_key_jmespath with powertools_json custom function + persistence_store.event_key_jmespath = "[requestContext.authorizer.claims.sub, powertools_json(body).id]" + persistence_store.event_key_compiled_jmespath = jmespath.compile(persistence_store.event_key_jmespath) + sub_attr_value = "cognito_user" + key_attr_value = "some_key" + expected_value = [sub_attr_value, key_attr_value] + api_gateway_proxy_event = { + "requestContext": {"authorizer": {"claims": {"sub": sub_attr_value}}}, + "body": json.dumps({"id": key_attr_value}), + } + + # WHEN calling _get_hashed_idempotency_key + result = persistence_store._get_hashed_idempotency_key(api_gateway_proxy_event) + + # THEN the hashed idempotency key should match the extracted values generated hash + assert result == persistence_store._generate_hash(expected_value) + + +@pytest.mark.parametrize("persistence_store_with_jmespath_options", ["powertools_json(data).payload"], indirect=True) +def test_custom_jmespath_function_overrides_builtin_functions(persistence_store_with_jmespath_options): + # GIVEN an persistence store with a custom jmespath_options + # AND use a builtin powertools custom function + with pytest.raises(jmespath.exceptions.UnknownFunctionError, match="Unknown function: powertools_json()"): + # WHEN calling _get_hashed_idempotency_key + # THEN raise unknown function + persistence_store_with_jmespath_options._get_hashed_idempotency_key({})