From d2428e90b917223bf42b279067ef468d6a4d9d2a Mon Sep 17 00:00:00 2001 From: Leandro Damascena Date: Sat, 15 Nov 2025 00:10:23 +0000 Subject: [PATCH 1/2] feat: add support for custom boto3 client --- docs/getting-started.md | 27 ++++++++++ .../execution.py | 26 +++++++--- tests/execution_test.py | 51 +++++++++++++++++++ 3 files changed, 98 insertions(+), 6 deletions(-) diff --git a/docs/getting-started.md b/docs/getting-started.md index 2ed0940..d6b050d 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -228,6 +228,33 @@ def handler(event: dict, context: DurableContext) -> str: Deploy this to Lambda and you have a durable function. The `greet_user` step is checkpointed automatically. +### Using a custom boto3 Lambda client + +If you need to customize the boto3 Lambda client used for durable execution operations (for example, to configure custom endpoints, retry settings, or credentials), you can pass a `boto3_client` parameter to the decorator. The client must be a boto3 Lambda client: + +```python +import boto3 +from botocore.config import Config +from aws_durable_execution_sdk_python import durable_execution, DurableContext + +# Create a custom boto3 Lambda client with specific configuration +custom_lambda_client = boto3.client( + 'lambda', + config=Config( + retries={'max_attempts': 5, 'mode': 'adaptive'}, + connect_timeout=10, + read_timeout=60, + ) +) + +@durable_execution(boto3_client=custom_lambda_client) +def handler(event: dict, context: DurableContext) -> dict: + # Your durable function logic + return {"status": "success"} +``` + +The custom Lambda client is used for all checkpoint and state management operations. If you don't provide a `boto3_client`, the SDK initializes a default Lambda client from your environment. + [↑ Back to top](#table-of-contents) ## Next steps diff --git a/src/aws_durable_execution_sdk_python/execution.py b/src/aws_durable_execution_sdk_python/execution.py index f9e3d74..a7e20ca 100644 --- a/src/aws_durable_execution_sdk_python/execution.py +++ b/src/aws_durable_execution_sdk_python/execution.py @@ -1,6 +1,7 @@ from __future__ import annotations import contextlib +import functools import json import logging from concurrent.futures import ThreadPoolExecutor @@ -30,6 +31,8 @@ if TYPE_CHECKING: from collections.abc import Callable, MutableMapping + import boto3 # type: ignore + from aws_durable_execution_sdk_python.types import LambdaContext @@ -193,8 +196,15 @@ def create_succeeded(cls, result: str) -> DurableExecutionInvocationOutput: def durable_execution( - func: Callable[[Any, DurableContext], Any], + func: Callable[[Any, DurableContext], Any] | None = None, + *, + boto3_client: boto3.client | None = None, # type: ignore ) -> Callable[[Any, LambdaContext], Any]: + # Decorator called with parameters + if func is None: + logger.debug("Decorator called with parameters") + return functools.partial(durable_execution, boto3_client=boto3_client) + logger.debug("Starting durable execution handler...") def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: @@ -210,11 +220,15 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: logger.debug("durableExecutionArn: %s", event.get("DurableExecutionArn")) invocation_input = DurableExecutionInvocationInput.from_dict(event) - service_client = ( - LambdaClient.initialize_local_runner_client() - if invocation_input.is_local_runner - else LambdaClient.initialize_from_env() - ) + # Local runner always uses its own client, otherwise use custom or default + if invocation_input.is_local_runner: + service_client = LambdaClient.initialize_local_runner_client() + else: + service_client = ( + LambdaClient(client=boto3_client) + if boto3_client is not None + else LambdaClient.initialize_from_env() + ) raw_input_payload: str | None = ( invocation_input.initial_execution_state.get_input_payload() diff --git a/tests/execution_test.py b/tests/execution_test.py index bc49913..68a06d1 100644 --- a/tests/execution_test.py +++ b/tests/execution_test.py @@ -1961,3 +1961,54 @@ def test_handler(event: Any, context: DurableContext) -> dict: assert call_args[0][0] == "Checkpoint system failed" assert call_args[1]["extra"]["Error"] == error_obj assert call_args[1]["extra"]["ResponseMetadata"] == metadata_obj + + +def test_durable_execution_with_boto3_client_parameter(): + """Test durable_execution decorator accepts boto3_client parameter.""" + # GIVEN a custom boto3 Lambda client + mock_boto3_client = Mock() + mock_boto3_client.checkpoint_durable_execution.return_value = { + "CheckpointToken": "new_token", + "NewExecutionState": {"Operations": [], "NextMarker": ""}, + } + mock_boto3_client.get_durable_execution_state.return_value = { + "Operations": [], + "NextMarker": "", + } + + # GIVEN a durable function decorated with the custom client + @durable_execution(boto3_client=mock_boto3_client) + def test_handler(event: Any, context: DurableContext) -> dict: + return {"result": "success"} + + event = { + "DurableExecutionArn": "arn:test:execution", + "CheckpointToken": "token123", + "InitialExecutionState": { + "Operations": [ + { + "Id": "exec1", + "Type": "EXECUTION", + "Status": "STARTED", + "ExecutionDetails": {"InputPayload": '{"input": "test"}'}, + } + ], + "NextMarker": "", + }, + "LocalRunner": False, + } + + lambda_context = Mock() + lambda_context.aws_request_id = "test-request" + lambda_context.client_context = None + lambda_context.identity = None + lambda_context._epoch_deadline_time_in_ms = 1000000 # noqa: SLF001 + lambda_context.invoked_function_arn = None + lambda_context.tenant_id = None + + # WHEN the handler is invoked + result = test_handler(event, lambda_context) + + # THEN the execution succeeds using the custom client + assert result["Status"] == InvocationStatus.SUCCEEDED.value + assert result["Result"] == '{"result": "success"}' From 27c9f493fe50c154b3489d65b20fde5907710608 Mon Sep 17 00:00:00 2001 From: Leandro Damascena Date: Sat, 15 Nov 2025 00:14:42 +0000 Subject: [PATCH 2/2] feat: remove unused ignore --- src/aws_durable_execution_sdk_python/execution.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aws_durable_execution_sdk_python/execution.py b/src/aws_durable_execution_sdk_python/execution.py index a7e20ca..e3e7be5 100644 --- a/src/aws_durable_execution_sdk_python/execution.py +++ b/src/aws_durable_execution_sdk_python/execution.py @@ -198,7 +198,7 @@ def create_succeeded(cls, result: str) -> DurableExecutionInvocationOutput: def durable_execution( func: Callable[[Any, DurableContext], Any] | None = None, *, - boto3_client: boto3.client | None = None, # type: ignore + boto3_client: boto3.client | None = None, ) -> Callable[[Any, LambdaContext], Any]: # Decorator called with parameters if func is None: