feat(idempotent): Add support for jmespath_options

Like for the validator, idempotent utility should allow for extracting an idempotent key using custom functions

Author: Michael Brewer

Diff for: 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:

Diff for: 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,23 @@ 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)