feat(idempotent): Add support for jmespath_options

aws_lambda_powertools/utilities/idempotency/persistence/base.py

@@ -19,6 +19,7 @@
     IdempotencyItemAlreadyExistsError,
     IdempotencyValidationError,
 )
+from aws_lambda_powertools.utilities.validation.jmespath_functions import PowertoolsFunctions
 
 logger = logging.getLogger(__name__)
 
@@ -112,6 +113,7 @@ def __init__(
         use_local_cache: bool = False,
         local_cache_max_items: int = 256,
         hash_function: str = "md5",
+        jmespath_options: Dict = None,
     ) -> None:
         """
         Initialize the base persistence layer
@@ -130,6 +132,8 @@ def __init__(
             Max number of items to store in local cache, by default 1024
         hash_function: str, optional
             Function to use for calculating hashes, by default md5.
+        jmespath_options : Dict
+            Alternative JMESPath options to be included when filtering expr
         """
         self.event_key_jmespath = event_key_jmespath
         if self.event_key_jmespath:
@@ -143,6 +147,9 @@ def __init__(
             self.validation_key_jmespath = jmespath.compile(payload_validation_jmespath)
             self.payload_validation_enabled = True
         self.hash_function = getattr(hashlib, hash_function)
+        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:
         """
@@ -160,8 +167,12 @@ 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)
+            )
+
         return self._generate_hash(data)
 
     def _get_hashed_payload(self, lambda_event: Dict[str, Any]) -> str:

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

@@ -1,6 +1,8 @@
 import copy
+import json
 import sys
 
+import jmespath
 import pytest
 from botocore import stub
 
@@ -638,3 +640,33 @@ def test_delete_from_cache_when_empty(persistence_store):
     except KeyError:
         # THEN we should not get a KeyError
         pytest.fail("KeyError should not happen")
+
+
+@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({})