Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 3 commits into from
Jan 23, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion aws_lambda_powertools/utilities/idempotency/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
):
Expand All @@ -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
Expand All @@ -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:
Expand Down
9 changes: 9 additions & 0 deletions aws_lambda_powertools/utilities/idempotency/idempotency.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ def idempotent(
context: LambdaContext,
persistence_store: BasePersistenceLayer,
config: IdempotencyConfig | None = None,
key_prefix: str | None = None,
**kwargs,
) -> Any:
"""
Expand All @@ -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
--------
Expand Down Expand Up @@ -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,
)
Expand All @@ -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:
"""
Expand All @@ -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
--------
Expand All @@ -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,
),
)
Expand Down Expand Up @@ -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,
)
Expand Down
17 changes: 12 additions & 5 deletions aws_lambda_powertools/utilities/idempotency/persistence/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
Expand All @@ -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
Expand Down
26 changes: 25 additions & 1 deletion docs/utilities/idempotency.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down Expand Up @@ -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)"
Expand Down
Original file line number Diff line number Diff line change
@@ -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)
Original file line number Diff line number Diff line change
@@ -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)
51 changes: 39 additions & 12 deletions tests/functional/idempotency/_boto3/test_idempotency.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import copy
import dataclasses
import datetime
import warnings
from typing import Any, Optional
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Loading