feat: boto3 sessions in batch, parameters & idempotency (#717)

aws_lambda_powertools/utilities/batch/sqs.py

@@ -31,6 +31,8 @@ class PartialSQSProcessor(BasePartialProcessor):
         botocore config object
     suppress_exception: bool, optional
         Supress exception raised if any messages fail processing, by default False
+    boto3_session : boto3.session.Session, optional
+            Boto3 session to use for AWS API communication
 
 
     Example
@@ -56,12 +58,18 @@ class PartialSQSProcessor(BasePartialProcessor):
 
     """
 
-    def __init__(self, config: Optional[Config] = None, suppress_exception: bool = False):
+    def __init__(
+        self,
+        config: Optional[Config] = None,
+        suppress_exception: bool = False,
+        boto3_session: Optional[boto3.session.Session] = None,
+    ):
         """
         Initializes sqs client.
         """
         config = config or Config()
-        self.client = boto3.client("sqs", config=config)
+        session = boto3_session or boto3.session.Session()
+        self.client = session.client("sqs", config=config)
         self.suppress_exception = suppress_exception
 
         super().__init__()
@@ -142,6 +150,7 @@ def sqs_batch_processor(
     record_handler: Callable,
     config: Optional[Config] = None,
     suppress_exception: bool = False,
+    boto3_session: Optional[boto3.session.Session] = None,
 ):
     """
     Middleware to handle SQS batch event processing
@@ -160,6 +169,8 @@ def sqs_batch_processor(
             botocore config object
     suppress_exception: bool, optional
         Supress exception raised if any messages fail processing, by default False
+    boto3_session : boto3.session.Session, optional
+            Boto3 session to use for AWS API communication
 
     Examples
     --------
@@ -180,7 +191,9 @@ def sqs_batch_processor(
 
     """
     config = config or Config()
-    processor = PartialSQSProcessor(config=config, suppress_exception=suppress_exception)
+    session = boto3_session or boto3.session.Session()
+
+    processor = PartialSQSProcessor(config=config, suppress_exception=suppress_exception, boto3_session=session)
 
     records = event["Records"]
 

aws_lambda_powertools/utilities/parameters/appconfig.py

@@ -29,6 +29,8 @@ class AppConfigProvider(BaseProvider):
         Application of the configuration to pass during client initialization
     config: botocore.config.Config, optional
         Botocore configuration to pass during client initialization
+    boto3_session : boto3.session.Session, optional
+            Boto3 session to use for AWS API communication
 
     Example
     -------
@@ -60,13 +62,20 @@ class AppConfigProvider(BaseProvider):
 
     client: Any = None
 
-    def __init__(self, environment: str, application: Optional[str] = None, config: Optional[Config] = None):
+    def __init__(
+        self,
+        environment: str,
+        application: Optional[str] = None,
+        config: Optional[Config] = None,
+        boto3_session: Optional[boto3.session.Session] = None,
+    ):
         """
         Initialize the App Config client
         """
 
         config = config or Config()
-        self.client = boto3.client("appconfig", config=config)
+        session = boto3_session or boto3.session.Session()
+        self.client = session.client("appconfig", config=config)
         self.application = resolve_env_var_choice(
             choice=application, env=os.getenv(constants.SERVICE_NAME_ENV, "service_undefined")
         )

aws_lambda_powertools/utilities/parameters/dynamodb.py

@@ -30,6 +30,8 @@ class DynamoDBProvider(BaseProvider):
         Complete url to reference local DynamoDB instance, e.g. http://localhost:8080
     config: botocore.config.Config, optional
         Botocore configuration to pass during client initialization
+    boto3_session : boto3.session.Session, optional
+            Boto3 session to use for AWS API communication
 
     Example
     -------
@@ -149,13 +151,16 @@ def __init__(
         value_attr: str = "value",
         endpoint_url: Optional[str] = None,
         config: Optional[Config] = None,
+        boto3_session: Optional[boto3.session.Session] = None,
     ):
         """
         Initialize the DynamoDB client
         """
 
         config = config or Config()
-        self.table = boto3.resource("dynamodb", endpoint_url=endpoint_url, config=config).Table(table_name)
+        session = boto3_session or boto3.session.Session()
+
+        self.table = session.resource("dynamodb", endpoint_url=endpoint_url, config=config).Table(table_name)
 
         self.key_attr = key_attr
         self.sort_attr = sort_attr

aws_lambda_powertools/utilities/parameters/secrets.py

@@ -19,6 +19,8 @@ class SecretsProvider(BaseProvider):
     ----------
     config: botocore.config.Config, optional
         Botocore configuration to pass during client initialization
+    boto3_session : boto3.session.Session, optional
+            Boto3 session to use for AWS API communication
 
     Example
     -------
@@ -58,14 +60,14 @@ class SecretsProvider(BaseProvider):
 
     client: Any = None
 
-    def __init__(self, config: Optional[Config] = None):
+    def __init__(self, config: Optional[Config] = None, boto3_session: Optional[boto3.session.Session] = None):
         """
         Initialize the Secrets Manager client
         """
 
         config = config or Config()
-
-        self.client = boto3.client("secretsmanager", config=config)
+        session = boto3_session or boto3.session.Session()
+        self.client = session.client("secretsmanager", config=config)
 
         super().__init__()
 

aws_lambda_powertools/utilities/parameters/ssm.py

@@ -19,6 +19,8 @@ class SSMProvider(BaseProvider):
     ----------
     config: botocore.config.Config, optional
         Botocore configuration to pass during client initialization
+    boto3_session : boto3.session.Session, optional
+            Boto3 session to use for AWS API communication
 
     Example
     -------
@@ -74,13 +76,14 @@ class SSMProvider(BaseProvider):
 
     client: Any = None
 
-    def __init__(self, config: Optional[Config] = None):
+    def __init__(self, config: Optional[Config] = None, boto3_session: Optional[boto3.session.Session] = None):
         """
         Initialize the SSM Parameter Store client
         """
 
         config = config or Config()
-        self.client = boto3.client("ssm", config=config)
+        session = boto3_session or boto3.session.Session()
+        self.client = session.client("ssm", config=config)
 
         super().__init__()
 

docs/utilities/batch.md

@@ -143,10 +143,13 @@ Use `PartialSQSProcessor` context manager to access a list of all return values
         return result
     ```
 
-### Passing custom boto3 config
+### Customizing boto configuration
 
-If you need to pass custom configuration such as region to the SDK, you can pass your own [botocore config object](https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html) to
-the `sqs_batch_processor` decorator:
+The **`config`** and **`boto3_session`** parameters enable you to pass in a custom [botocore config object](https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html)
+or a custom [boto3 session](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html) when using the `sqs_batch_processor`
+decorator or `PartialSQSProcessor` class.
+
+> Custom config example
 
 === "Decorator"
 
@@ -193,6 +196,53 @@ the `sqs_batch_processor` decorator:
         return result
     ```
 
+> Custom boto3 session example
+
+=== "Decorator"
+
+    ```python  hl_lines="4 12"
+    from aws_lambda_powertools.utilities.batch import sqs_batch_processor
+    from botocore.config import Config
+
+    session = boto3.session.Session()
+
+    def record_handler(record):
+        # This will be called for each individual message from a batch
+        # It should raise an exception if the message was not processed successfully
+        return_value = do_something_with(record["body"])
+        return return_value
+
+    @sqs_batch_processor(record_handler=record_handler, boto3_session=session)
+    def lambda_handler(event, context):
+        return {"statusCode": 200}
+    ```
+
+=== "Context manager"
+
+    ```python  hl_lines="4 16"
+    from aws_lambda_powertools.utilities.batch import PartialSQSProcessor
+    import boto3
+
+    session = boto3.session.Session()
+
+    def record_handler(record):
+        # This will be called for each individual message from a batch
+        # It should raise an exception if the message was not processed successfully
+        return_value = do_something_with(record["body"])
+        return return_value
+
+
+    def lambda_handler(event, context):
+        records = event["Records"]
+
+        processor = PartialSQSProcessor(boto3_session=session)
+
+        with processor(records, record_handler):
+            result = processor.process()
+
+        return result
+    ```
+
 ### Suppressing exceptions
 
 If you want to disable the default behavior where `SQSBatchProcessingError` is raised if there are any errors, you can pass the `suppress_exception` boolean argument.

docs/utilities/idempotency.md

@@ -548,7 +548,7 @@ This means that we will raise **`IdempotencyKeyError`** if the evaluation of **`
 
 ### Customizing boto configuration
 
-You can provide a custom boto configuration via **`boto_config`**, or an existing boto session via **`boto3_session`** parameters, when constructing the persistence store.
+The **`boto_config`** and **`boto3_session`** parameters enable you to pass in a custom [botocore config object](https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html) or a custom [boto3 session](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html) when constructing the persistence store.
 
 === "Custom session"
 

docs/utilities/parameters.md

@@ -501,3 +501,40 @@ Here is the mapping between this utility's functions and methods and the underly
 | DynamoDB            | `DynamoDBProvider.get`          | `dynamodb`       | ([Table resource](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/dynamodb.html#table)) | [get_item](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/dynamodb.html#DynamoDB.Table.get_item)
 | DynamoDB            | `DynamoDBProvider.get_multiple` | `dynamodb`       | ([Table resource](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/dynamodb.html#table)) | [query](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/dynamodb.html#DynamoDB.Table.query)
 | App Config          | `get_app_config`                | `appconfig`      | [get_configuration](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/appconfig.html#AppConfig.Client.get_configuration) |
+
+
+### Customizing boto configuration
+
+The **`config`** and **`boto3_session`** parameters enable you to pass in a custom [botocore config object](https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html) or a custom [boto3 session](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html) when constructing any of the built-in provider classes.
+
+> **Example**
+
+
+=== "Custom session"
+
+    ```python hl_lines="2 4 5"
+    from aws_lambda_powertools.utilities import parameters
+    import boto3
+
+    boto3_session = boto3.session.Session()
+    ssm_provider = parameters.SSMProvider(boto3_session=boto3_session)
+
+    def handler(event, context):
+        # Retrieve a single parameter
+        value = ssm_provider.get("/my/parameter")
+        ...
+    ```
+=== "Custom config"
+
+    ```python hl_lines="2 4 5"
+    from aws_lambda_powertools.utilities import parameters
+    from botocore.config import Config
+
+    boto_config = Config()
+    ssm_provider = parameters.SSMProvider(config=boto_config)
+
+    def handler(event, context):
+        # Retrieve a single parameter
+        value = ssm_provider.get("/my/parameter")
+        ...
+    ```