feat(idempotent): Add support for jmespath_options

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}")

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__)
 

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()

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({})