diff --git a/aws_lambda_powertools/utilities/batch/sqs.py b/aws_lambda_powertools/utilities/batch/sqs.py index e37fdbd3fb5..38773a399dd 100644 --- a/aws_lambda_powertools/utilities/batch/sqs.py +++ b/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"] diff --git a/aws_lambda_powertools/utilities/parameters/appconfig.py b/aws_lambda_powertools/utilities/parameters/appconfig.py index 4a400aa7789..d1613c14513 100644 --- a/aws_lambda_powertools/utilities/parameters/appconfig.py +++ b/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") ) diff --git a/aws_lambda_powertools/utilities/parameters/dynamodb.py b/aws_lambda_powertools/utilities/parameters/dynamodb.py index 39bd1a8d6b7..9220edf3b05 100644 --- a/aws_lambda_powertools/utilities/parameters/dynamodb.py +++ b/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 diff --git a/aws_lambda_powertools/utilities/parameters/secrets.py b/aws_lambda_powertools/utilities/parameters/secrets.py index 5699876d90e..b64e70ae184 100644 --- a/aws_lambda_powertools/utilities/parameters/secrets.py +++ b/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__() diff --git a/aws_lambda_powertools/utilities/parameters/ssm.py b/aws_lambda_powertools/utilities/parameters/ssm.py index 2a16ad91f08..4cbb16354c7 100644 --- a/aws_lambda_powertools/utilities/parameters/ssm.py +++ b/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__() diff --git a/docs/utilities/batch.md b/docs/utilities/batch.md index 96770fb1849..56ab160e9f9 100644 --- a/docs/utilities/batch.md +++ b/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. diff --git a/docs/utilities/idempotency.md b/docs/utilities/idempotency.md index a9a5a129e63..3508e2190e3 100644 --- a/docs/utilities/idempotency.md +++ b/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" diff --git a/docs/utilities/parameters.md b/docs/utilities/parameters.md index 081d22817ab..51fd0196abd 100644 --- a/docs/utilities/parameters.md +++ b/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") + ... + ```