From 7df82ff2f8c045e546b6f18aefc24075b0c0cabc Mon Sep 17 00:00:00 2001 From: Michael Brewer Date: Tue, 2 Mar 2021 20:12:09 -0800 Subject: [PATCH 1/4] feat(idempotent): Add support for jmespath_options Like for the validator, idempotent utility should allow for extracting an idempotent key using custom functions --- .../utilities/idempotency/persistence/base.py | 13 ++++++++++- .../idempotency/test_idempotency.py | 22 +++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/aws_lambda_powertools/utilities/idempotency/persistence/base.py b/aws_lambda_powertools/utilities/idempotency/persistence/base.py index 7e31c7d394b..d33b46a86b9 100644 --- a/aws_lambda_powertools/utilities/idempotency/persistence/base.py +++ b/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 --git a/tests/functional/idempotency/test_idempotency.py b/tests/functional/idempotency/test_idempotency.py index 6a85d69d957..903c92331c2 100644 --- a/tests/functional/idempotency/test_idempotency.py +++ b/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) From b5b740b4a86a0e43ba6c511a6dd526c8615fa0cf Mon Sep 17 00:00:00 2001 From: Michael Brewer Date: Tue, 2 Mar 2021 20:29:33 -0800 Subject: [PATCH 2/4] test(idempotent): Add test for support jmepath custom function --- .../functional/idempotency/test_idempotency.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tests/functional/idempotency/test_idempotency.py b/tests/functional/idempotency/test_idempotency.py index 903c92331c2..f10f6fae6aa 100644 --- a/tests/functional/idempotency/test_idempotency.py +++ b/tests/functional/idempotency/test_idempotency.py @@ -5,7 +5,9 @@ import jmespath import pytest from botocore import stub +from jmespath import functions +from aws_lambda_powertools.utilities.idempotency import DynamoDBPersistenceLayer from aws_lambda_powertools.utilities.idempotency.exceptions import ( IdempotencyAlreadyInProgressError, IdempotencyInconsistentStateError, @@ -660,3 +662,18 @@ def test_jmespath_with_powertools_json(persistence_store): # THEN the hashed idempotency key should match the extracted values generated hash assert result == persistence_store._generate_hash(expected_value) + + +def test_custom_jmespath_function_overrides_builtin_functions(): + class CustomFunctions(functions.Functions): + @functions.signature({"types": ["string"]}) + def _func_echo_decoder(self, value): + return value + + persistence_store = DynamoDBPersistenceLayer( + table_name="foo", + event_key_jmespath="powertools_json(data).payload", + jmespath_options={"custom_functions": CustomFunctions()}, + ) + with pytest.raises(jmespath.exceptions.UnknownFunctionError, match="Unknown function: powertools_json()"): + persistence_store._get_hashed_idempotency_key({}) From 79ef461679e91dd93caa72715edd5254a190c4c2 Mon Sep 17 00:00:00 2001 From: Michael Brewer Date: Tue, 2 Mar 2021 20:45:30 -0800 Subject: [PATCH 3/4] chore: Fix test by using fixtures --- tests/functional/idempotency/conftest.py | 18 ++++++++++++++++ .../idempotency/test_idempotency.py | 21 +++++++------------ 2 files changed, 25 insertions(+), 14 deletions(-) diff --git a/tests/functional/idempotency/conftest.py b/tests/functional/idempotency/conftest.py index 918eac9a507..84827b111c7 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 f10f6fae6aa..0aab9274f88 100644 --- a/tests/functional/idempotency/test_idempotency.py +++ b/tests/functional/idempotency/test_idempotency.py @@ -5,9 +5,7 @@ import jmespath import pytest from botocore import stub -from jmespath import functions -from aws_lambda_powertools.utilities.idempotency import DynamoDBPersistenceLayer from aws_lambda_powertools.utilities.idempotency.exceptions import ( IdempotencyAlreadyInProgressError, IdempotencyInconsistentStateError, @@ -664,16 +662,11 @@ def test_jmespath_with_powertools_json(persistence_store): assert result == persistence_store._generate_hash(expected_value) -def test_custom_jmespath_function_overrides_builtin_functions(): - class CustomFunctions(functions.Functions): - @functions.signature({"types": ["string"]}) - def _func_echo_decoder(self, value): - return value - - persistence_store = DynamoDBPersistenceLayer( - table_name="foo", - event_key_jmespath="powertools_json(data).payload", - jmespath_options={"custom_functions": CustomFunctions()}, - ) +@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()"): - persistence_store._get_hashed_idempotency_key({}) + # WHEN calling _get_hashed_idempotency_key + # THEN raise unknown function + persistence_store_with_jmespath_options._get_hashed_idempotency_key({}) From 922f7f0e588805fa31b47e0cfea41c415a82d2d0 Mon Sep 17 00:00:00 2001 From: Michael Brewer Date: Wed, 3 Mar 2021 07:36:17 -0800 Subject: [PATCH 4/4] refactor(jmespath_functions): Move into shared --- .../{utilities/validation => shared}/jmespath_functions.py | 0 .../utilities/idempotency/persistence/base.py | 2 +- aws_lambda_powertools/utilities/validation/base.py | 3 ++- 3 files changed, 3 insertions(+), 2 deletions(-) rename aws_lambda_powertools/{utilities/validation => shared}/jmespath_functions.py (100%) 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 426a2564dee..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, @@ -21,7 +22,6 @@ IdempotencyKeyError, IdempotencyValidationError, ) -from aws_lambda_powertools.utilities.validation.jmespath_functions import PowertoolsFunctions logger = logging.getLogger(__name__) 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__)