feat(idempotency): add support for custom Idempotency key prefix (#5898)

* Adding support for custom prefix in Idempotency key

* Refactoring parameter

Author: leandrodamascena

aws_lambda_powertools/utilities/idempotency/base.py

@@ -74,6 +74,7 @@ def __init__(
         config: IdempotencyConfig,
         persistence_store: BasePersistenceLayer,
         output_serializer: BaseIdempotencySerializer | None = None,
+        key_prefix: str | None = None,
         function_args: tuple | None = None,
         function_kwargs: dict | None = None,
     ):
@@ -91,6 +92,8 @@ def __init__(
         output_serializer: BaseIdempotencySerializer | None
             Serializer to transform the data to and from a dictionary.
             If not supplied, no serialization is done via the NoOpSerializer
+        key_prefix: str | Optional
+            Custom prefix for idempotency key: key_prefix#hash
         function_args: tuple | None
             Function arguments
         function_kwargs: dict | None
@@ -102,8 +105,14 @@ def __init__(
         self.fn_args = function_args
         self.fn_kwargs = function_kwargs
         self.config = config
+        self.key_prefix = key_prefix
+
+        persistence_store.configure(
+            config=config,
+            function_name=f"{self.function.__module__}.{self.function.__qualname__}",
+            key_prefix=self.key_prefix,
+        )
 
-        persistence_store.configure(config, f"{self.function.__module__}.{self.function.__qualname__}")
         self.persistence_store = persistence_store
 
     def handle(self) -> Any:

aws_lambda_powertools/utilities/idempotency/idempotency.py

@@ -40,6 +40,7 @@ def idempotent(
     context: LambdaContext,
     persistence_store: BasePersistenceLayer,
     config: IdempotencyConfig | None = None,
+    key_prefix: str | None = None,
     **kwargs,
 ) -> Any:
     """
@@ -57,6 +58,8 @@ def idempotent(
         Instance of BasePersistenceLayer to store data
     config: IdempotencyConfig
         Configuration
+    key_prefix: str | Optional
+        Custom prefix for idempotency key: key_prefix#hash
 
     Examples
     --------
@@ -94,6 +97,7 @@ def idempotent(
         function_payload=event,
         config=config,
         persistence_store=persistence_store,
+        key_prefix=key_prefix,
         function_args=args,
         function_kwargs=kwargs,
     )
@@ -108,6 +112,7 @@ def idempotent_function(
     persistence_store: BasePersistenceLayer,
     config: IdempotencyConfig | None = None,
     output_serializer: BaseIdempotencySerializer | type[BaseIdempotencyModelSerializer] | None = None,
+    key_prefix: str | None = None,
     **kwargs: Any,
 ) -> Any:
     """
@@ -128,6 +133,8 @@ def idempotent_function(
             If not supplied, no serialization is done via the NoOpSerializer.
             In case a serializer of type inheriting BaseIdempotencyModelSerializer is given,
             the serializer is derived from the function return type.
+    key_prefix: str | Optional
+        Custom prefix for idempotency key: key_prefix#hash
 
     Examples
     --------
@@ -154,6 +161,7 @@ def process_order(customer_id: str, order: dict, **kwargs):
                 persistence_store=persistence_store,
                 config=config,
                 output_serializer=output_serializer,
+                key_prefix=key_prefix,
                 **kwargs,
             ),
         )
@@ -191,6 +199,7 @@ def decorate(*args, **kwargs):
             config=config,
             persistence_store=persistence_store,
             output_serializer=output_serializer,
+            key_prefix=key_prefix,
             function_args=args,
             function_kwargs=kwargs,
         )

aws_lambda_powertools/utilities/idempotency/persistence/base.py

@@ -54,7 +54,12 @@ def __init__(self):
         self.use_local_cache = False
         self.hash_function = hashlib.md5
 
-    def configure(self, config: IdempotencyConfig, function_name: str | None = None) -> None:
+    def configure(
+        self,
+        config: IdempotencyConfig,
+        function_name: str | None = None,
+        key_prefix: str | None = None,
+    ) -> None:
         """
         Initialize the base persistence layer from the configuration settings
 
@@ -64,8 +69,12 @@ def configure(self, config: IdempotencyConfig, function_name: str | None = None)
             Idempotency configuration settings
         function_name: str, Optional
             The name of the function being decorated
+        key_prefix: str | Optional
+            Custom prefix for idempotency key: key_prefix#hash
         """
-        self.function_name = f"{os.getenv(constants.LAMBDA_FUNCTION_NAME_ENV, 'test-func')}.{function_name or ''}"
+        self.function_name = (
+            key_prefix or f"{os.getenv(constants.LAMBDA_FUNCTION_NAME_ENV, 'test-func')}.{function_name or ''}"
+        )
 
         if self.configured:
             # Prevent being reconfigured multiple times
@@ -75,9 +84,7 @@ def configure(self, config: IdempotencyConfig, function_name: str | None = None)
         self.event_key_jmespath = config.event_key_jmespath
         if config.event_key_jmespath:
             self.event_key_compiled_jmespath = jmespath.compile(config.event_key_jmespath)
-        self.jmespath_options = config.jmespath_options
-        if not self.jmespath_options:
-            self.jmespath_options = {"custom_functions": PowertoolsFunctions()}
+        self.jmespath_options = config.jmespath_options or {"custom_functions": PowertoolsFunctions()}
         if config.payload_validation_jmespath:
             self.validation_key_jmespath = jmespath.compile(config.payload_validation_jmespath)
             self.payload_validation_enabled = True

docs/utilities/idempotency.md

@@ -18,7 +18,9 @@ The idempotency utility allows you to retry operations within a time window with
 
 The property of idempotency means that an operation does not cause additional side effects if it is called more than once with the same input parameters.
 
-**Idempotency key** is a combination of **(a)** Lambda function name, **(b)** fully qualified name of your function, and **(c)** a hash of the entire payload or part(s) of the payload you specify.
+<!-- markdownlint-disable MD013 -->
+**Idempotency key** By default, this is a combination of **(a)** Lambda function name, **(b)** fully qualified name of your function, and **(c)** a hash of the entire payload or part(s) of the payload you specify. However, you can customize the key generation by using **(a)** a [custom prefix name](#customizing-the-idempotency-key-generation), while still incorporating **(c)** a hash of the entire payload or part(s) of the payload you specify.
+<!-- markdownlint-enable MD013 -->
 
 **Idempotent request** is an operation with the same input previously processed that is not expired in your persistent storage or in-memory cache.
 
@@ -356,6 +358,28 @@ You can change this expiration window with the **`expires_after_seconds`** param
 
     A record might still be valid (`COMPLETE`) when we retrieved, but in some rare cases it might expire a second later. A record could also be [cached in memory](#using-in-memory-cache). You might also want to have idempotent transactions that should expire in seconds.
 
+### Customizing the Idempotency key generation
+
+!!! warning "Warning: Changing the idempotency key generation will invalidate existing idempotency records"
+
+Use **`key_prefix`** parameter in the `@idempotent` or `@idempotent_function` decorators to define a custom prefix for your Idempotency Key. This allows you to decouple idempotency key name from function names. It can be useful during application refactoring, for example.
+
+=== "Using a custom prefix in Lambda Handler"
+
+    ```python hl_lines="25"
+    --8<-- "examples/idempotency/src/working_with_custom_idempotency_key_prefix.py"
+    ```
+
+    1. The Idempotency record will be something like `my_custom_prefix#c4ca4238a0b923820dcc509a6f75849b`
+
+=== "Using a custom prefix in standalone functions"
+
+    ```python hl_lines="32"
+    --8<-- "examples/idempotency/src/working_with_custom_idempotency_key_prefix_standalone.py"
+    ```
+
+    1. The Idempotency record will be something like `my_custom_prefix#c4ca4238a0b923820dcc509a6f75849b`
+
 ### Lambda timeouts
 
 !!! note "You can skip this section if you are using the [`@idempotent` decorator](#idempotent-decorator)"

examples/idempotency/src/working_with_custom_idempotency_key_prefix.py

@@ -0,0 +1,39 @@
+import os
+from dataclasses import dataclass, field
+from uuid import uuid4
+
+from aws_lambda_powertools.utilities.idempotency import (
+    DynamoDBPersistenceLayer,
+    idempotent,
+)
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+table = os.getenv("IDEMPOTENCY_TABLE", "")
+persistence_layer = DynamoDBPersistenceLayer(table_name=table)
+
+
+@dataclass
+class Payment:
+    user_id: str
+    product_id: str
+    payment_id: str = field(default_factory=lambda: f"{uuid4()}")
+
+
+class PaymentError(Exception): ...
+
+
+@idempotent(persistence_store=persistence_layer, key_prefix="my_custom_prefix") # (1)!
+def lambda_handler(event: dict, context: LambdaContext):
+    try:
+        payment: Payment = create_subscription_payment(event)
+        return {
+            "payment_id": payment.payment_id,
+            "message": "success",
+            "statusCode": 200,
+        }
+    except Exception as exc:
+        raise PaymentError(f"Error creating payment {str(exc)}")
+
+
+def create_subscription_payment(event: dict) -> Payment:
+    return Payment(**event)

examples/idempotency/src/working_with_custom_idempotency_key_prefix_standalone.py

@@ -0,0 +1,46 @@
+import os
+from dataclasses import dataclass
+
+from aws_lambda_powertools.utilities.idempotency import (
+    DynamoDBPersistenceLayer,
+    IdempotencyConfig,
+    idempotent_function,
+)
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+table = os.getenv("IDEMPOTENCY_TABLE", "")
+dynamodb = DynamoDBPersistenceLayer(table_name=table)
+config = IdempotencyConfig(event_key_jmespath="order_id")  # see Choosing a payload subset section
+
+
+@dataclass
+class OrderItem:
+    sku: str
+    description: str
+
+
+@dataclass
+class Order:
+    item: OrderItem
+    order_id: int
+
+
+@idempotent_function(
+    data_keyword_argument="order",
+    config=config,
+    persistence_store=dynamodb,
+    key_prefix="my_custom_prefix", # (1)!
+)
+def process_order(order: Order): 
+    return f"processed order {order.order_id}"
+
+
+def lambda_handler(event: dict, context: LambdaContext):
+    # see Lambda timeouts section
+    config.register_lambda_context(context) 
+
+    order_item = OrderItem(sku="fake", description="sample")
+    order = Order(item=order_item, order_id=1)
+
+    # `order` parameter must be called as a keyword argument to work
+    process_order(order=order)

tests/functional/idempotency/_boto3/test_idempotency.py

@@ -1,4 +1,5 @@
 import copy
+import dataclasses
 import datetime
 import warnings
 from typing import Any, Optional
@@ -59,13 +60,6 @@
 TESTS_MODULE_PREFIX = "test-func.tests.functional.idempotency._boto3.test_idempotency"
 
 
-def get_dataclasses_lib():
-    """Python 3.6 doesn't support dataclasses natively"""
-    import dataclasses
-
-    return dataclasses
-
-
 # Using parametrize to run test twice, with two separate instances of persistence store. One instance with caching
 # enabled, and one without.
 @pytest.mark.parametrize("idempotency_config", [{"use_local_cache": False}, {"use_local_cache": True}], indirect=True)
@@ -1313,7 +1307,6 @@ def record_handler(record):
 @pytest.mark.parametrize("output_serializer_type", ["explicit", "deduced"])
 def test_idempotent_function_serialization_dataclass(output_serializer_type: str):
     # GIVEN
-    dataclasses = get_dataclasses_lib()
     config = IdempotencyConfig(use_local_cache=True)
     mock_event = {"customer_id": "fake", "transaction_id": "fake-id"}
     idempotency_key = f"{TESTS_MODULE_PREFIX}.test_idempotent_function_serialization_dataclass.<locals>.collect_payment#{hash_idempotency_key(mock_event)}"  # noqa E501
@@ -1359,7 +1352,6 @@ def collect_payment(payment: PaymentInput) -> PaymentOutput:
 
 def test_idempotent_function_serialization_dataclass_failure_no_return_type():
     # GIVEN
-    dataclasses = get_dataclasses_lib()
     config = IdempotencyConfig(use_local_cache=True)
     mock_event = {"customer_id": "fake", "transaction_id": "fake-id"}
     idempotency_key = f"{TESTS_MODULE_PREFIX}.test_idempotent_function_serialization_pydantic_failure_no_return_type.<locals>.collect_payment#{hash_idempotency_key(mock_event)}"  # noqa E501
@@ -1655,7 +1647,6 @@ def test_invalid_dynamodb_persistence_layer():
 
 def test_idempotent_function_dataclasses():
     # Scenario _prepare_data should convert a python dataclasses to a dict
-    dataclasses = get_dataclasses_lib()
 
     @dataclasses.dataclass
     class Foo:
@@ -1670,7 +1661,6 @@ class Foo:
 
 def test_idempotent_function_dataclass_with_jmespath():
     # GIVEN
-    dataclasses = get_dataclasses_lib()
     config = IdempotencyConfig(event_key_jmespath="transaction_id", use_local_cache=True)
     mock_event = {"customer_id": "fake", "transaction_id": "fake-id"}
     idempotency_key = f"{TESTS_MODULE_PREFIX}.test_idempotent_function_dataclass_with_jmespath.<locals>.collect_payment#{hash_idempotency_key(mock_event['transaction_id'])}"  # noqa E501
@@ -2019,7 +2009,6 @@ def lambda_handler(event, context):
 @pytest.mark.parametrize("output_serializer_type", ["explicit", "deduced"])
 def test_idempotent_function_serialization_dataclass_with_optional_return(output_serializer_type: str):
     # GIVEN
-    dataclasses = get_dataclasses_lib()
     config = IdempotencyConfig(use_local_cache=True)
     mock_event = {"customer_id": "fake", "transaction_id": "fake-id"}
     idempotency_key = f"{TESTS_MODULE_PREFIX}.test_idempotent_function_serialization_dataclass_with_optional_return.<locals>.collect_payment#{hash_idempotency_key(mock_event)}"  # noqa E501
@@ -2061,3 +2050,41 @@ def collect_payment(payment: PaymentInput) -> Optional[PaymentOutput]:
     assert isinstance(second_call, PaymentOutput)
     assert second_call.customer_id == payment.customer_id
     assert second_call.transaction_id == payment.transaction_id
+
+
+def test_idempotent_function_with_custom_prefix_standalone_function():
+    # Scenario to validate we can use idempotent_function with any function
+    mock_event = {"data": "value"}
+    idempotency_key = f"my-custom-prefix#{hash_idempotency_key(mock_event)}"
+    persistence_layer = MockPersistenceLayer(expected_idempotency_key=idempotency_key)
+    expected_result = {"message": "Foo"}
+
+    @idempotent_function(
+        persistence_store=persistence_layer,
+        data_keyword_argument="record",
+        key_prefix="my-custom-prefix",
+    )
+    def record_handler(record):
+        return expected_result
+
+    # WHEN calling the function
+    result = record_handler(record=mock_event)
+    # THEN we expect the function to execute successfully
+    assert result == expected_result
+
+
+def test_idempotent_function_with_custom_prefix_lambda_handler(lambda_context):
+    # Scenario to validate we can use idempotent_function with any function
+    mock_event = {"data": "value"}
+    idempotency_key = f"my-custom-prefix#{hash_idempotency_key(mock_event)}"
+    persistence_layer = MockPersistenceLayer(expected_idempotency_key=idempotency_key)
+    expected_result = {"message": "Foo"}
+
+    @idempotent(persistence_store=persistence_layer, key_prefix="my-custom-prefix")
+    def lambda_handler(record, context):
+        return expected_result
+
+    # WHEN calling the function
+    result = lambda_handler(mock_event, lambda_context)
+    # THEN we expect the function to execute successfully
+    assert result == expected_result