diff --git a/aws_lambda_powertools/utilities/idempotency/idempotency.py b/aws_lambda_powertools/utilities/idempotency/idempotency.py index 2b3a4ceef39..b77c3013cbb 100644 --- a/aws_lambda_powertools/utilities/idempotency/idempotency.py +++ b/aws_lambda_powertools/utilities/idempotency/idempotency.py @@ -131,7 +131,7 @@ def handle(self) -> Any: try: # We call save_inprogress first as an optimization for the most common case where no idempotent record # already exists. If it succeeds, there's no need to call get_record. - self.persistence_store.save_inprogress(event=self.event) + self.persistence_store.save_inprogress(event=self.event, context=self.context) except IdempotencyItemAlreadyExistsError: # Now we know the item already exists, we can retrieve it record = self._get_idempotency_record() @@ -151,7 +151,7 @@ def _get_idempotency_record(self) -> DataRecord: """ try: - event_record = self.persistence_store.get_record(self.event) + event_record = self.persistence_store.get_record(event=self.event, context=self.context) except IdempotencyItemNotFoundError: # This code path will only be triggered if the record is removed between save_inprogress and get_record. logger.debug( @@ -219,7 +219,9 @@ def _call_lambda_handler(self) -> Any: # We need these nested blocks to preserve lambda handler exception in case the persistence store operation # also raises an exception try: - self.persistence_store.delete_record(event=self.event, exception=handler_exception) + self.persistence_store.delete_record( + event=self.event, context=self.context, exception=handler_exception + ) except Exception as delete_exception: raise IdempotencyPersistenceLayerError( "Failed to delete record from idempotency store" @@ -228,7 +230,7 @@ def _call_lambda_handler(self) -> Any: else: try: - self.persistence_store.save_success(event=self.event, result=handler_response) + self.persistence_store.save_success(event=self.event, context=self.context, result=handler_response) except Exception as save_exception: raise IdempotencyPersistenceLayerError( "Failed to update record state to success in idempotency store" diff --git a/aws_lambda_powertools/utilities/idempotency/persistence/base.py b/aws_lambda_powertools/utilities/idempotency/persistence/base.py index de726115d95..37c9968b3e0 100644 --- a/aws_lambda_powertools/utilities/idempotency/persistence/base.py +++ b/aws_lambda_powertools/utilities/idempotency/persistence/base.py @@ -23,6 +23,7 @@ IdempotencyKeyError, IdempotencyValidationError, ) +from aws_lambda_powertools.utilities.typing import LambdaContext logger = logging.getLogger(__name__) @@ -152,14 +153,16 @@ def configure(self, config: IdempotencyConfig) -> None: self._cache = LRUDict(max_items=config.local_cache_max_items) self.hash_function = getattr(hashlib, config.hash_function) - def _get_hashed_idempotency_key(self, lambda_event: Dict[str, Any]) -> str: + def _get_hashed_idempotency_key(self, event: Dict[str, Any], context: LambdaContext) -> str: """ Extract data from lambda event using event key jmespath, and return a hashed representation Parameters ---------- - lambda_event: Dict[str, Any] + event: Dict[str, Any] Lambda event + context: LambdaContext + Lambda context Returns ------- @@ -167,19 +170,18 @@ def _get_hashed_idempotency_key(self, lambda_event: Dict[str, Any]) -> str: Hashed representation of the data extracted by the jmespath expression """ - data = lambda_event + data = event if self.event_key_jmespath: - data = self.event_key_compiled_jmespath.search( - lambda_event, options=jmespath.Options(**self.jmespath_options) - ) + data = self.event_key_compiled_jmespath.search(event, options=jmespath.Options(**self.jmespath_options)) if self.is_missing_idempotency_key(data): if self.raise_on_no_idempotency_key: raise IdempotencyKeyError("No data found to create a hashed idempotency_key") warnings.warn(f"No value found for idempotency_key. jmespath: {self.event_key_jmespath}") - return self._generate_hash(data) + generated_hash = self._generate_hash(data) + return f"{context.function_name}#{generated_hash}" @staticmethod def is_missing_idempotency_key(data) -> bool: @@ -298,7 +300,7 @@ def _delete_from_cache(self, idempotency_key: str): if idempotency_key in self._cache: del self._cache[idempotency_key] - def save_success(self, event: Dict[str, Any], result: dict) -> None: + def save_success(self, event: Dict[str, Any], context: LambdaContext, result: dict) -> None: """ Save record of function's execution completing successfully @@ -306,13 +308,15 @@ def save_success(self, event: Dict[str, Any], result: dict) -> None: ---------- event: Dict[str, Any] Lambda event + context: LambdaContext + Lambda context result: dict The response from lambda handler """ response_data = json.dumps(result, cls=Encoder) data_record = DataRecord( - idempotency_key=self._get_hashed_idempotency_key(event), + idempotency_key=self._get_hashed_idempotency_key(event, context), status=STATUS_CONSTANTS["COMPLETED"], expiry_timestamp=self._get_expiry_timestamp(), response_data=response_data, @@ -326,7 +330,7 @@ def save_success(self, event: Dict[str, Any], result: dict) -> None: self._save_to_cache(data_record) - def save_inprogress(self, event: Dict[str, Any]) -> None: + def save_inprogress(self, event: Dict[str, Any], context: LambdaContext) -> None: """ Save record of function's execution being in progress @@ -334,9 +338,11 @@ def save_inprogress(self, event: Dict[str, Any]) -> None: ---------- event: Dict[str, Any] Lambda event + context: LambdaContext + Lambda context """ data_record = DataRecord( - idempotency_key=self._get_hashed_idempotency_key(event), + idempotency_key=self._get_hashed_idempotency_key(event, context), status=STATUS_CONSTANTS["INPROGRESS"], expiry_timestamp=self._get_expiry_timestamp(), payload_hash=self._get_hashed_payload(event), @@ -349,7 +355,7 @@ def save_inprogress(self, event: Dict[str, Any]) -> None: self._put_record(data_record) - def delete_record(self, event: Dict[str, Any], exception: Exception): + def delete_record(self, event: Dict[str, Any], context: LambdaContext, exception: Exception): """ Delete record from the persistence store @@ -357,10 +363,12 @@ def delete_record(self, event: Dict[str, Any], exception: Exception): ---------- event: Dict[str, Any] Lambda event + context: LambdaContext + Lambda context exception The exception raised by the lambda handler """ - data_record = DataRecord(idempotency_key=self._get_hashed_idempotency_key(event)) + data_record = DataRecord(idempotency_key=self._get_hashed_idempotency_key(event, context)) logger.debug( f"Lambda raised an exception ({type(exception).__name__}). Clearing in progress record in persistence " @@ -370,7 +378,7 @@ def delete_record(self, event: Dict[str, Any], exception: Exception): self._delete_from_cache(data_record.idempotency_key) - def get_record(self, event: Dict[str, Any]) -> DataRecord: + def get_record(self, event: Dict[str, Any], context: LambdaContext) -> DataRecord: """ Calculate idempotency key for lambda_event, then retrieve item from persistence store using idempotency key and return it as a DataRecord instance.and return it as a DataRecord instance. @@ -378,6 +386,9 @@ def get_record(self, event: Dict[str, Any]) -> DataRecord: Parameters ---------- event: Dict[str, Any] + Lambda event + context: LambdaContext + Lambda context Returns ------- @@ -392,7 +403,7 @@ def get_record(self, event: Dict[str, Any]) -> DataRecord: Event payload doesn't match the stored record for the given idempotency key """ - idempotency_key = self._get_hashed_idempotency_key(event) + idempotency_key = self._get_hashed_idempotency_key(event, context) cached_record = self._retrieve_from_cache(idempotency_key=idempotency_key) if cached_record: diff --git a/tests/functional/idempotency/conftest.py b/tests/functional/idempotency/conftest.py index 532d551ef40..d34d5da7d12 100644 --- a/tests/functional/idempotency/conftest.py +++ b/tests/functional/idempotency/conftest.py @@ -2,6 +2,7 @@ import hashlib import json import os +from collections import namedtuple from decimal import Decimal from unittest import mock @@ -34,6 +35,18 @@ def lambda_apigw_event(): return event +@pytest.fixture +def lambda_context(): + lambda_context = { + "function_name": "test-func", + "memory_limit_in_mb": 128, + "invoked_function_arn": "arn:aws:lambda:eu-west-1:809313241234:function:test-func", + "aws_request_id": "52fdfc07-2182-154f-163f-5f0f9a621d72", + } + + return namedtuple("LambdaContext", lambda_context.keys())(*lambda_context.values()) + + @pytest.fixture def timestamp_future(): return str(int((datetime.datetime.now() + datetime.timedelta(seconds=3600)).timestamp())) @@ -132,10 +145,10 @@ def expected_params_put_item_with_validation(hashed_idempotency_key, hashed_vali @pytest.fixture -def hashed_idempotency_key(lambda_apigw_event, default_jmespath): +def hashed_idempotency_key(lambda_apigw_event, default_jmespath, lambda_context): compiled_jmespath = jmespath.compile(default_jmespath) data = compiled_jmespath.search(lambda_apigw_event) - return hashlib.md5(json.dumps(data).encode()).hexdigest() + return "test-func#" + hashlib.md5(json.dumps(data).encode()).hexdigest() @pytest.fixture @@ -143,7 +156,7 @@ def hashed_idempotency_key_with_envelope(lambda_apigw_event): event = unwrap_event_from_envelope( data=lambda_apigw_event, envelope=envelopes.API_GATEWAY_HTTP, jmespath_options={} ) - return hashlib.md5(json.dumps(event).encode()).hexdigest() + return "test-func#" + hashlib.md5(json.dumps(event).encode()).hexdigest() @pytest.fixture diff --git a/tests/functional/idempotency/test_idempotency.py b/tests/functional/idempotency/test_idempotency.py index 6f5ba74a7aa..503ec7d6183 100644 --- a/tests/functional/idempotency/test_idempotency.py +++ b/tests/functional/idempotency/test_idempotency.py @@ -34,6 +34,7 @@ def test_idempotent_lambda_already_completed( hashed_idempotency_key, serialized_lambda_response, deserialized_lambda_response, + lambda_context, ): """ Test idempotent decorator where event with matching event key has already been succesfully processed @@ -62,7 +63,7 @@ def test_idempotent_lambda_already_completed( def lambda_handler(event, context): raise Exception - lambda_resp = lambda_handler(lambda_apigw_event, {}) + lambda_resp = lambda_handler(lambda_apigw_event, lambda_context) assert lambda_resp == deserialized_lambda_response stubber.assert_no_pending_responses() @@ -77,6 +78,7 @@ def test_idempotent_lambda_in_progress( lambda_response, timestamp_future, hashed_idempotency_key, + lambda_context, ): """ Test idempotent decorator where lambda_handler is already processing an event with matching event key @@ -106,7 +108,7 @@ def lambda_handler(event, context): return lambda_response with pytest.raises(IdempotencyAlreadyInProgressError) as ex: - lambda_handler(lambda_apigw_event, {}) + lambda_handler(lambda_apigw_event, lambda_context) assert ( ex.value.args[0] == "Execution already in progress with idempotency key: " "body=a3edd699125517bb49d562501179ecbd" @@ -126,6 +128,7 @@ def test_idempotent_lambda_in_progress_with_cache( timestamp_future, hashed_idempotency_key, mocker, + lambda_context, ): """ Test idempotent decorator where lambda_handler is already processing an event with matching event key, cache @@ -165,7 +168,7 @@ def lambda_handler(event, context): loops = 3 for _ in range(loops): with pytest.raises(IdempotencyAlreadyInProgressError) as ex: - lambda_handler(lambda_apigw_event, {}) + lambda_handler(lambda_apigw_event, lambda_context) assert ( ex.value.args[0] == "Execution already in progress with idempotency key: " "body=a3edd699125517bb49d562501179ecbd" @@ -192,6 +195,7 @@ def test_idempotent_lambda_first_execution( serialized_lambda_response, deserialized_lambda_response, hashed_idempotency_key, + lambda_context, ): """ Test idempotent decorator when lambda is executed with an event with a previously unknown event key @@ -208,7 +212,7 @@ def test_idempotent_lambda_first_execution( def lambda_handler(event, context): return lambda_response - lambda_handler(lambda_apigw_event, {}) + lambda_handler(lambda_apigw_event, lambda_context) stubber.assert_no_pending_responses() stubber.deactivate() @@ -225,6 +229,7 @@ def test_idempotent_lambda_first_execution_cached( lambda_response, hashed_idempotency_key, mocker, + lambda_context, ): """ Test idempotent decorator when lambda is executed with an event with a previously unknown event key. Ensure @@ -243,7 +248,7 @@ def test_idempotent_lambda_first_execution_cached( def lambda_handler(event, context): return lambda_response - lambda_handler(lambda_apigw_event, {}) + lambda_handler(lambda_apigw_event, lambda_context) retrieve_from_cache_spy.assert_called_once() save_to_cache_spy.assert_called_once() @@ -251,7 +256,7 @@ def lambda_handler(event, context): assert persistence_store._cache.get(hashed_idempotency_key).status == "COMPLETED" # This lambda call should not call AWS API - lambda_handler(lambda_apigw_event, {}) + lambda_handler(lambda_apigw_event, lambda_context) assert retrieve_from_cache_spy.call_count == 3 retrieve_from_cache_spy.assert_called_with(idempotency_key=hashed_idempotency_key) @@ -270,6 +275,7 @@ def test_idempotent_lambda_expired( expected_params_update_item, expected_params_put_item, hashed_idempotency_key, + lambda_context, ): """ Test idempotent decorator when lambda is called with an event it succesfully handled already, but outside of the @@ -288,7 +294,7 @@ def test_idempotent_lambda_expired( def lambda_handler(event, context): return lambda_response - lambda_handler(lambda_apigw_event, {}) + lambda_handler(lambda_apigw_event, lambda_context) stubber.assert_no_pending_responses() stubber.deactivate() @@ -303,6 +309,7 @@ def test_idempotent_lambda_exception( lambda_response, hashed_idempotency_key, expected_params_put_item, + lambda_context, ): """ Test idempotent decorator when lambda is executed with an event with a previously unknown event key, but @@ -326,7 +333,7 @@ def lambda_handler(event, context): raise Exception("Something went wrong!") with pytest.raises(Exception): - lambda_handler(lambda_apigw_event, {}) + lambda_handler(lambda_apigw_event, lambda_context) stubber.assert_no_pending_responses() stubber.deactivate() @@ -343,6 +350,7 @@ def test_idempotent_lambda_already_completed_with_validation_bad_payload( lambda_response, hashed_idempotency_key, hashed_validation_key, + lambda_context, ): """ Test idempotent decorator where event with matching event key has already been successfully processed @@ -371,7 +379,7 @@ def lambda_handler(event, context): with pytest.raises(IdempotencyValidationError): lambda_apigw_event["requestContext"]["accountId"] += "1" # Alter the request payload - lambda_handler(lambda_apigw_event, {}) + lambda_handler(lambda_apigw_event, lambda_context) stubber.assert_no_pending_responses() stubber.deactivate() @@ -386,6 +394,7 @@ def test_idempotent_lambda_expired_during_request( lambda_response, expected_params_update_item, hashed_idempotency_key, + lambda_context, ): """ Test idempotent decorator when lambda is called with an event it succesfully handled already. Persistence store @@ -427,7 +436,7 @@ def lambda_handler(event, context): # max retries exceeded before get_item and put_item agree on item state, so exception gets raised with pytest.raises(IdempotencyInconsistentStateError): - lambda_handler(lambda_apigw_event, {}) + lambda_handler(lambda_apigw_event, lambda_context) stubber.assert_no_pending_responses() stubber.deactivate() @@ -442,6 +451,7 @@ def test_idempotent_persistence_exception_deleting( lambda_response, hashed_idempotency_key, expected_params_put_item, + lambda_context, ): """ Test idempotent decorator when lambda is executed with an event with a previously unknown event key, but @@ -460,7 +470,7 @@ def lambda_handler(event, context): raise Exception("Something went wrong!") with pytest.raises(IdempotencyPersistenceLayerError) as exc: - lambda_handler(lambda_apigw_event, {}) + lambda_handler(lambda_apigw_event, lambda_context) assert exc.value.args[0] == "Failed to delete record from idempotency store" stubber.assert_no_pending_responses() @@ -476,6 +486,7 @@ def test_idempotent_persistence_exception_updating( lambda_response, hashed_idempotency_key, expected_params_put_item, + lambda_context, ): """ Test idempotent decorator when lambda is executed with an event with a previously unknown event key, but @@ -494,7 +505,7 @@ def lambda_handler(event, context): return {"message": "success!"} with pytest.raises(IdempotencyPersistenceLayerError) as exc: - lambda_handler(lambda_apigw_event, {}) + lambda_handler(lambda_apigw_event, lambda_context) assert exc.value.args[0] == "Failed to update record state to success in idempotency store" stubber.assert_no_pending_responses() @@ -510,6 +521,7 @@ def test_idempotent_persistence_exception_getting( lambda_response, hashed_idempotency_key, expected_params_put_item, + lambda_context, ): """ Test idempotent decorator when lambda is executed with an event with a previously unknown event key, but @@ -526,7 +538,7 @@ def lambda_handler(event, context): return {"message": "success!"} with pytest.raises(IdempotencyPersistenceLayerError) as exc: - lambda_handler(lambda_apigw_event, {}) + lambda_handler(lambda_apigw_event, lambda_context) assert exc.value.args[0] == "Failed to get record from idempotency store" stubber.assert_no_pending_responses() @@ -545,6 +557,7 @@ def test_idempotent_lambda_first_execution_with_validation( lambda_response, hashed_idempotency_key, hashed_validation_key, + lambda_context, ): """ Test idempotent decorator when lambda is executed with an event with a previously unknown event key @@ -560,7 +573,7 @@ def test_idempotent_lambda_first_execution_with_validation( def lambda_handler(event, context): return lambda_response - lambda_handler(lambda_apigw_event, {}) + lambda_handler(lambda_apigw_event, lambda_context) stubber.assert_no_pending_responses() stubber.deactivate() @@ -578,6 +591,7 @@ def test_idempotent_lambda_with_validator_util( deserialized_lambda_response, hashed_idempotency_key_with_envelope, mock_function, + lambda_context, ): """ Test idempotent decorator where event with matching event key has already been succesfully processed, using the @@ -610,7 +624,7 @@ def lambda_handler(event, context): return "shouldn't get here!" mock_function.assert_not_called() - lambda_resp = lambda_handler(lambda_apigw_event, {}) + lambda_resp = lambda_handler(lambda_apigw_event, lambda_context) assert lambda_resp == deserialized_lambda_response stubber.assert_no_pending_responses() @@ -715,7 +729,7 @@ def test_is_missing_idempotency_key(): "idempotency_config", [{"use_local_cache": False, "event_key_jmespath": "body"}], indirect=True ) def test_default_no_raise_on_missing_idempotency_key( - idempotency_config: IdempotencyConfig, persistence_store: DynamoDBPersistenceLayer + idempotency_config: IdempotencyConfig, persistence_store: DynamoDBPersistenceLayer, lambda_context ): # GIVEN a persistence_store with use_local_cache = False and event_key_jmespath = "body" persistence_store.configure(idempotency_config) @@ -723,17 +737,18 @@ def test_default_no_raise_on_missing_idempotency_key( assert "body" in persistence_store.event_key_jmespath # WHEN getting the hashed idempotency key for an event with no `body` key - hashed_key = persistence_store._get_hashed_idempotency_key({}) + hashed_key = persistence_store._get_hashed_idempotency_key({}, lambda_context) # THEN return the hash of None - assert md5(json.dumps(None).encode()).hexdigest() == hashed_key + expected_value = "test-func#" + md5(json.dumps(None).encode()).hexdigest() + assert expected_value == hashed_key @pytest.mark.parametrize( "idempotency_config", [{"use_local_cache": False, "event_key_jmespath": "[body, x]"}], indirect=True ) def test_raise_on_no_idempotency_key( - idempotency_config: IdempotencyConfig, persistence_store: DynamoDBPersistenceLayer + idempotency_config: IdempotencyConfig, persistence_store: DynamoDBPersistenceLayer, lambda_context ): # GIVEN a persistence_store with raise_on_no_idempotency_key and no idempotency key in the request persistence_store.configure(idempotency_config) @@ -743,7 +758,7 @@ def test_raise_on_no_idempotency_key( # WHEN getting the hashed idempotency key for an event with no `body` key with pytest.raises(IdempotencyKeyError) as excinfo: - persistence_store._get_hashed_idempotency_key({}) + persistence_store._get_hashed_idempotency_key({}, lambda_context) # THEN raise IdempotencyKeyError error assert "No data found to create a hashed idempotency_key" in str(excinfo.value) @@ -760,7 +775,7 @@ def test_raise_on_no_idempotency_key( indirect=True, ) def test_jmespath_with_powertools_json( - idempotency_config: IdempotencyConfig, persistence_store: DynamoDBPersistenceLayer + idempotency_config: IdempotencyConfig, persistence_store: DynamoDBPersistenceLayer, lambda_context ): # GIVEN an event_key_jmespath with powertools_json custom function persistence_store.configure(idempotency_config) @@ -773,15 +788,15 @@ def test_jmespath_with_powertools_json( } # WHEN calling _get_hashed_idempotency_key - result = persistence_store._get_hashed_idempotency_key(api_gateway_proxy_event) + result = persistence_store._get_hashed_idempotency_key(api_gateway_proxy_event, lambda_context) # THEN the hashed idempotency key should match the extracted values generated hash - assert result == persistence_store._generate_hash(expected_value) + assert result == "test-func#" + persistence_store._generate_hash(expected_value) @pytest.mark.parametrize("config_with_jmespath_options", ["powertools_json(data).payload"], indirect=True) def test_custom_jmespath_function_overrides_builtin_functions( - config_with_jmespath_options: IdempotencyConfig, persistence_store: DynamoDBPersistenceLayer + config_with_jmespath_options: IdempotencyConfig, persistence_store: DynamoDBPersistenceLayer, lambda_context ): # GIVEN an persistence store with a custom jmespath_options # AND use a builtin powertools custom function @@ -790,10 +805,10 @@ def test_custom_jmespath_function_overrides_builtin_functions( with pytest.raises(jmespath.exceptions.UnknownFunctionError, match="Unknown function: powertools_json()"): # WHEN calling _get_hashed_idempotency_key # THEN raise unknown function - persistence_store._get_hashed_idempotency_key({}) + persistence_store._get_hashed_idempotency_key({}, lambda_context) -def test_idempotent_lambda_save_inprogress_error(persistence_store: DynamoDBPersistenceLayer): +def test_idempotent_lambda_save_inprogress_error(persistence_store: DynamoDBPersistenceLayer, lambda_context): # GIVEN a miss configured persistence layer # like no table was created for the idempotency persistence layer stubber = stub.Stubber(persistence_store.table.meta.client) @@ -807,7 +822,7 @@ def lambda_handler(event, context): # WHEN handling the idempotent call # AND save_inprogress raises a ClientError with pytest.raises(IdempotencyPersistenceLayerError) as e: - lambda_handler({}, {}) + lambda_handler({}, lambda_context) # THEN idempotent should raise an IdempotencyPersistenceLayerError stubber.assert_no_pending_responses()