Skip to content

Commit 3493789

Browse files
committed
* 'develop' of https://github.com/awslabs/aws-lambda-powertools-python: feat: boto3 sessions in batch, parameters & idempotency (aws-powertools#717)
2 parents 3d4305b + c21ba45 commit 3493789

File tree

8 files changed

+134
-15
lines changed

8 files changed

+134
-15
lines changed

aws_lambda_powertools/utilities/batch/sqs.py

+16-3
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,8 @@ class PartialSQSProcessor(BasePartialProcessor):
3131
botocore config object
3232
suppress_exception: bool, optional
3333
Supress exception raised if any messages fail processing, by default False
34+
boto3_session : boto3.session.Session, optional
35+
Boto3 session to use for AWS API communication
3436
3537
3638
Example
@@ -56,12 +58,18 @@ class PartialSQSProcessor(BasePartialProcessor):
5658
5759
"""
5860

59-
def __init__(self, config: Optional[Config] = None, suppress_exception: bool = False):
61+
def __init__(
62+
self,
63+
config: Optional[Config] = None,
64+
suppress_exception: bool = False,
65+
boto3_session: Optional[boto3.session.Session] = None,
66+
):
6067
"""
6168
Initializes sqs client.
6269
"""
6370
config = config or Config()
64-
self.client = boto3.client("sqs", config=config)
71+
session = boto3_session or boto3.session.Session()
72+
self.client = session.client("sqs", config=config)
6573
self.suppress_exception = suppress_exception
6674

6775
super().__init__()
@@ -142,6 +150,7 @@ def sqs_batch_processor(
142150
record_handler: Callable,
143151
config: Optional[Config] = None,
144152
suppress_exception: bool = False,
153+
boto3_session: Optional[boto3.session.Session] = None,
145154
):
146155
"""
147156
Middleware to handle SQS batch event processing
@@ -160,6 +169,8 @@ def sqs_batch_processor(
160169
botocore config object
161170
suppress_exception: bool, optional
162171
Supress exception raised if any messages fail processing, by default False
172+
boto3_session : boto3.session.Session, optional
173+
Boto3 session to use for AWS API communication
163174
164175
Examples
165176
--------
@@ -180,7 +191,9 @@ def sqs_batch_processor(
180191
181192
"""
182193
config = config or Config()
183-
processor = PartialSQSProcessor(config=config, suppress_exception=suppress_exception)
194+
session = boto3_session or boto3.session.Session()
195+
196+
processor = PartialSQSProcessor(config=config, suppress_exception=suppress_exception, boto3_session=session)
184197

185198
records = event["Records"]
186199

aws_lambda_powertools/utilities/parameters/appconfig.py

+11-2
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@ class AppConfigProvider(BaseProvider):
2929
Application of the configuration to pass during client initialization
3030
config: botocore.config.Config, optional
3131
Botocore configuration to pass during client initialization
32+
boto3_session : boto3.session.Session, optional
33+
Boto3 session to use for AWS API communication
3234
3335
Example
3436
-------
@@ -60,13 +62,20 @@ class AppConfigProvider(BaseProvider):
6062

6163
client: Any = None
6264

63-
def __init__(self, environment: str, application: Optional[str] = None, config: Optional[Config] = None):
65+
def __init__(
66+
self,
67+
environment: str,
68+
application: Optional[str] = None,
69+
config: Optional[Config] = None,
70+
boto3_session: Optional[boto3.session.Session] = None,
71+
):
6472
"""
6573
Initialize the App Config client
6674
"""
6775

6876
config = config or Config()
69-
self.client = boto3.client("appconfig", config=config)
77+
session = boto3_session or boto3.session.Session()
78+
self.client = session.client("appconfig", config=config)
7079
self.application = resolve_env_var_choice(
7180
choice=application, env=os.getenv(constants.SERVICE_NAME_ENV, "service_undefined")
7281
)

aws_lambda_powertools/utilities/parameters/dynamodb.py

+6-1
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@ class DynamoDBProvider(BaseProvider):
3030
Complete url to reference local DynamoDB instance, e.g. http://localhost:8080
3131
config: botocore.config.Config, optional
3232
Botocore configuration to pass during client initialization
33+
boto3_session : boto3.session.Session, optional
34+
Boto3 session to use for AWS API communication
3335
3436
Example
3537
-------
@@ -149,13 +151,16 @@ def __init__(
149151
value_attr: str = "value",
150152
endpoint_url: Optional[str] = None,
151153
config: Optional[Config] = None,
154+
boto3_session: Optional[boto3.session.Session] = None,
152155
):
153156
"""
154157
Initialize the DynamoDB client
155158
"""
156159

157160
config = config or Config()
158-
self.table = boto3.resource("dynamodb", endpoint_url=endpoint_url, config=config).Table(table_name)
161+
session = boto3_session or boto3.session.Session()
162+
163+
self.table = session.resource("dynamodb", endpoint_url=endpoint_url, config=config).Table(table_name)
159164

160165
self.key_attr = key_attr
161166
self.sort_attr = sort_attr

aws_lambda_powertools/utilities/parameters/secrets.py

+5-3
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@ class SecretsProvider(BaseProvider):
1919
----------
2020
config: botocore.config.Config, optional
2121
Botocore configuration to pass during client initialization
22+
boto3_session : boto3.session.Session, optional
23+
Boto3 session to use for AWS API communication
2224
2325
Example
2426
-------
@@ -58,14 +60,14 @@ class SecretsProvider(BaseProvider):
5860

5961
client: Any = None
6062

61-
def __init__(self, config: Optional[Config] = None):
63+
def __init__(self, config: Optional[Config] = None, boto3_session: Optional[boto3.session.Session] = None):
6264
"""
6365
Initialize the Secrets Manager client
6466
"""
6567

6668
config = config or Config()
67-
68-
self.client = boto3.client("secretsmanager", config=config)
69+
session = boto3_session or boto3.session.Session()
70+
self.client = session.client("secretsmanager", config=config)
6971

7072
super().__init__()
7173

aws_lambda_powertools/utilities/parameters/ssm.py

+5-2
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@ class SSMProvider(BaseProvider):
1919
----------
2020
config: botocore.config.Config, optional
2121
Botocore configuration to pass during client initialization
22+
boto3_session : boto3.session.Session, optional
23+
Boto3 session to use for AWS API communication
2224
2325
Example
2426
-------
@@ -74,13 +76,14 @@ class SSMProvider(BaseProvider):
7476

7577
client: Any = None
7678

77-
def __init__(self, config: Optional[Config] = None):
79+
def __init__(self, config: Optional[Config] = None, boto3_session: Optional[boto3.session.Session] = None):
7880
"""
7981
Initialize the SSM Parameter Store client
8082
"""
8183

8284
config = config or Config()
83-
self.client = boto3.client("ssm", config=config)
85+
session = boto3_session or boto3.session.Session()
86+
self.client = session.client("ssm", config=config)
8487

8588
super().__init__()
8689

docs/utilities/batch.md

+53-3
Original file line numberDiff line numberDiff line change
@@ -143,10 +143,13 @@ Use `PartialSQSProcessor` context manager to access a list of all return values
143143
return result
144144
```
145145

146-
### Passing custom boto3 config
146+
### Customizing boto configuration
147147

148-
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
149-
the `sqs_batch_processor` decorator:
148+
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)
149+
or a custom [boto3 session](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html) when using the `sqs_batch_processor`
150+
decorator or `PartialSQSProcessor` class.
151+
152+
> Custom config example
150153
151154
=== "Decorator"
152155

@@ -193,6 +196,53 @@ the `sqs_batch_processor` decorator:
193196
return result
194197
```
195198

199+
> Custom boto3 session example
200+
201+
=== "Decorator"
202+
203+
```python hl_lines="4 12"
204+
from aws_lambda_powertools.utilities.batch import sqs_batch_processor
205+
from botocore.config import Config
206+
207+
session = boto3.session.Session()
208+
209+
def record_handler(record):
210+
# This will be called for each individual message from a batch
211+
# It should raise an exception if the message was not processed successfully
212+
return_value = do_something_with(record["body"])
213+
return return_value
214+
215+
@sqs_batch_processor(record_handler=record_handler, boto3_session=session)
216+
def lambda_handler(event, context):
217+
return {"statusCode": 200}
218+
```
219+
220+
=== "Context manager"
221+
222+
```python hl_lines="4 16"
223+
from aws_lambda_powertools.utilities.batch import PartialSQSProcessor
224+
import boto3
225+
226+
session = boto3.session.Session()
227+
228+
def record_handler(record):
229+
# This will be called for each individual message from a batch
230+
# It should raise an exception if the message was not processed successfully
231+
return_value = do_something_with(record["body"])
232+
return return_value
233+
234+
235+
def lambda_handler(event, context):
236+
records = event["Records"]
237+
238+
processor = PartialSQSProcessor(boto3_session=session)
239+
240+
with processor(records, record_handler):
241+
result = processor.process()
242+
243+
return result
244+
```
245+
196246
### Suppressing exceptions
197247

198248
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

+1-1
Original file line numberDiff line numberDiff line change
@@ -548,7 +548,7 @@ This means that we will raise **`IdempotencyKeyError`** if the evaluation of **`
548548

549549
### Customizing boto configuration
550550

551-
You can provide a custom boto configuration via **`boto_config`**, or an existing boto session via **`boto3_session`** parameters, when constructing the persistence store.
551+
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.
552552

553553
=== "Custom session"
554554

docs/utilities/parameters.md

+37
Original file line numberDiff line numberDiff line change
@@ -501,3 +501,40 @@ Here is the mapping between this utility's functions and methods and the underly
501501
| 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)
502502
| 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)
503503
| App Config | `get_app_config` | `appconfig` | [get_configuration](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/appconfig.html#AppConfig.Client.get_configuration) |
504+
505+
506+
### Customizing boto configuration
507+
508+
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.
509+
510+
> **Example**
511+
512+
513+
=== "Custom session"
514+
515+
```python hl_lines="2 4 5"
516+
from aws_lambda_powertools.utilities import parameters
517+
import boto3
518+
519+
boto3_session = boto3.session.Session()
520+
ssm_provider = parameters.SSMProvider(boto3_session=boto3_session)
521+
522+
def handler(event, context):
523+
# Retrieve a single parameter
524+
value = ssm_provider.get("/my/parameter")
525+
...
526+
```
527+
=== "Custom config"
528+
529+
```python hl_lines="2 4 5"
530+
from aws_lambda_powertools.utilities import parameters
531+
from botocore.config import Config
532+
533+
boto_config = Config()
534+
ssm_provider = parameters.SSMProvider(config=boto_config)
535+
536+
def handler(event, context):
537+
# Retrieve a single parameter
538+
value = ssm_provider.get("/my/parameter")
539+
...
540+
```

0 commit comments

Comments
 (0)