feat: add support for custom boto3 client

Author: leandrodamascena

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

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()

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"}'