Skip to content

Commit

Permalink
refactor(typing): enable boto3 implicit type annotations (aws-powerto…
Browse files Browse the repository at this point in the history
…ols#4692)

* chore(deps-dev): enable mypy-boto3 implicit type annotations

Add boto3-stubs and botocore-stubs to enable implicit type annotations
with mypy_boto3_* packages. See
https://youtype.github.io/boto3_stubs_docs/mypy_boto3_s3/usage/#implicit-type-annotations

This enables removing very few unneeded explicit type annotations and
helps detect some new errors automatically.

No need to do

    boto3_config = boto3_config or Config()

before passing the boto3_config to the boto3_session client or resource
methods, as None is an accepted value.

Refactor the aws_lambda_powertools.utilities.parameters.base.BaseProvider
class: move the instantiation of clients in _build_boto3_client and
_build_boto3_resource_client to the subclasses and the user_agent method
calls to the constructor. Keeping correct typing information with the
original implementation was going to be too complex and hopefully the
refactor even made the code tidier.

Keep imports consistent: use boto3.session.Session, mypy_boto3_*.client,
mypy_boto3_*.service_resource and mypy_boto3_*.type_defs
as documented at
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
and
https://youtype.github.io/boto3_stubs_docs/mypy_boto3_s3/usage/#explicit-type-annotations

* Ignore mypy arg-type error with ConditionalCheckFailedException

* Reassign to boto3_session instead of inlined expression
  • Loading branch information
ericbn authored Aug 6, 2024
1 parent 0f610dc commit 26cfe7f
Show file tree
Hide file tree
Showing 24 changed files with 552 additions and 251 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -220,7 +220,7 @@ class S3ObjectLambdaEvent(DictWrapper):
import requests
from aws_lambda_powertools.utilities.data_classes.s3_object_event import S3ObjectLambdaEvent
session = boto3.Session()
session = boto3.session.Session()
s3 = session.client("s3")
def lambda_handler(event, context):
Expand Down
34 changes: 17 additions & 17 deletions aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
)

if TYPE_CHECKING:
from mypy_boto3_dynamodb import DynamoDBClient
from mypy_boto3_dynamodb.client import DynamoDBClient
from mypy_boto3_dynamodb.type_defs import AttributeValueTypeDef

logger = logging.getLogger(__name__)
Expand All @@ -43,7 +43,7 @@ def __init__(
validation_key_attr: str = "validation",
boto_config: Optional[Config] = None,
boto3_session: Optional[boto3.session.Session] = None,
boto3_client: "DynamoDBClient" | None = None,
boto3_client: Optional[DynamoDBClient] = None,
):
"""
Initialize the DynamoDB client
Expand Down Expand Up @@ -71,7 +71,7 @@ def __init__(
DynamoDB attribute name for hashed representation of the parts of the event used for validation
boto_config: botocore.config.Config, optional
Botocore configuration to pass during client initialization
boto3_session : boto3.Session, optional
boto3_session : boto3.session.Session, optional
Boto3 session to use for AWS API communication
boto3_client : DynamoDBClient, optional
Boto3 DynamoDB Client to use, boto3_session and boto_config will be ignored if both are provided
Expand All @@ -91,11 +91,9 @@ def __init__(
>>> return {"StatusCode": 200}
"""
if boto3_client is None:
self._boto_config = boto_config or Config()
self._boto3_session: boto3.Session = boto3_session or boto3.session.Session()
self.client: "DynamoDBClient" = self._boto3_session.client("dynamodb", config=self._boto_config)
else:
self.client = boto3_client
boto3_session = boto3_session or boto3.session.Session()
boto3_client = boto3_session.client("dynamodb", config=boto_config)
self.client = boto3_client

user_agent.register_feature_to_client(client=self.client, feature="idempotency")

Expand Down Expand Up @@ -246,13 +244,20 @@ def _put_record(self, data_record: DataRecord) -> None:
":now_in_millis": {"N": str(int(now.timestamp() * 1000))},
":inprogress": {"S": STATUS_CONSTANTS["INPROGRESS"]},
},
**self.return_value_on_condition, # type: ignore
**self.return_value_on_condition, # type: ignore[arg-type]
)
except ClientError as exc:
error_code = exc.response.get("Error", {}).get("Code")
if error_code == "ConditionalCheckFailedException":
old_data_record = self._item_to_data_record(exc.response["Item"]) if "Item" in exc.response else None
if old_data_record is not None:
try:
item = exc.response["Item"] # type: ignore[typeddict-item]
except KeyError:
logger.debug(
f"Failed to put record for already existing idempotency key: {data_record.idempotency_key}",
)
raise IdempotencyItemAlreadyExistsError() from exc
else:
old_data_record = self._item_to_data_record(item)
logger.debug(
f"Failed to put record for already existing idempotency key: "
f"{data_record.idempotency_key} with status: {old_data_record.status}, "
Expand All @@ -268,11 +273,6 @@ def _put_record(self, data_record: DataRecord) -> None:

raise IdempotencyItemAlreadyExistsError(old_data_record=old_data_record) from exc

logger.debug(
f"Failed to put record for already existing idempotency key: {data_record.idempotency_key}",
)
raise IdempotencyItemAlreadyExistsError() from exc

raise

@staticmethod
Expand All @@ -297,7 +297,7 @@ def boto3_supports_condition_check_failure(boto3_version: str) -> bool:
def _update_record(self, data_record: DataRecord):
logger.debug(f"Updating record for idempotency key: {data_record.idempotency_key}")
update_expression = "SET #response_data = :response_data, #expiry = :expiry, #status = :status"
expression_attr_values: Dict[str, "AttributeValueTypeDef"] = {
expression_attr_values: Dict[str, AttributeValueTypeDef] = {
":expiry": {"N": str(data_record.expiry_timestamp)},
":response_data": {"S": data_record.response_data},
":status": {"S": data_record.status},
Expand Down
25 changes: 10 additions & 15 deletions aws_lambda_powertools/utilities/parameters/appconfig.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,15 +3,15 @@
"""

import os
from typing import TYPE_CHECKING, Any, Dict, Optional, Union
from typing import TYPE_CHECKING, Dict, Optional, Union

import boto3
from botocore.config import Config

from aws_lambda_powertools.utilities.parameters.types import TransformOptions

if TYPE_CHECKING:
from mypy_boto3_appconfigdata import AppConfigDataClient
from mypy_boto3_appconfigdata.client import AppConfigDataClient

from aws_lambda_powertools.shared import constants
from aws_lambda_powertools.shared.functions import (
Expand Down Expand Up @@ -67,8 +67,6 @@ class AppConfigProvider(BaseProvider):
"""

client: Any = None

def __init__(
self,
environment: str,
Expand All @@ -80,15 +78,10 @@ def __init__(
"""
Initialize the App Config client
"""

super().__init__()

self.client: "AppConfigDataClient" = self._build_boto3_client(
service_name="appconfigdata",
client=boto3_client,
session=boto3_session,
config=config,
)
if boto3_client is None:
boto3_session = boto3_session or boto3.session.Session()
boto3_client = boto3_session.client("appconfigdata", config=config)
self.client = boto3_client

self.application = resolve_env_var_choice(
choice=application,
Expand All @@ -99,9 +92,11 @@ def __init__(

self._next_token: Dict[str, str] = {} # nosec - token for get_latest_configuration executions
# Dict to store the recently retrieved value for a specific configuration.
self.last_returned_value: Dict[str, str] = {}
self.last_returned_value: Dict[str, bytes] = {}

super().__init__(client=self.client)

def _get(self, name: str, **sdk_options) -> str:
def _get(self, name: str, **sdk_options) -> bytes:
"""
Retrieve a parameter value from AWS App config.
Expand Down
91 changes: 5 additions & 86 deletions aws_lambda_powertools/utilities/parameters/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,43 +10,30 @@
from abc import ABC, abstractmethod
from datetime import datetime, timedelta
from typing import (
TYPE_CHECKING,
Any,
Callable,
Dict,
NamedTuple,
Optional,
Tuple,
Type,
Union,
cast,
overload,
)

import boto3
from botocore.config import Config

from aws_lambda_powertools.shared import constants, user_agent
from aws_lambda_powertools.shared.functions import resolve_max_age
from aws_lambda_powertools.utilities.parameters.types import TransformOptions

from .exceptions import GetParameterError, TransformParameterError

if TYPE_CHECKING:
from mypy_boto3_appconfigdata import AppConfigDataClient
from mypy_boto3_dynamodb import DynamoDBServiceResource
from mypy_boto3_secretsmanager import SecretsManagerClient
from mypy_boto3_ssm import SSMClient


DEFAULT_MAX_AGE_SECS = "300"

# These providers will be dynamically initialized on first use of the helper functions
DEFAULT_PROVIDERS: Dict[str, Any] = {}
TRANSFORM_METHOD_JSON = "json"
TRANSFORM_METHOD_BINARY = "binary"
SUPPORTED_TRANSFORM_METHODS = [TRANSFORM_METHOD_JSON, TRANSFORM_METHOD_BINARY]
ParameterClients = Union["AppConfigDataClient", "SecretsManagerClient", "SSMClient"]

TRANSFORM_METHOD_MAPPING = {
TRANSFORM_METHOD_JSON: json.loads,
Expand All @@ -69,10 +56,14 @@ class BaseProvider(ABC):

store: Dict[Tuple, ExpirableValue]

def __init__(self):
def __init__(self, *, client=None, resource=None):
"""
Initialize the base provider
"""
if client is not None:
user_agent.register_feature_to_client(client=client, feature="parameters")
if resource is not None:
user_agent.register_feature_to_resource(resource=resource, feature="parameters")

self.store: Dict[Tuple, ExpirableValue] = {}

Expand Down Expand Up @@ -262,78 +253,6 @@ def _build_cache_key(
"""
return (name, transform, is_nested)

@staticmethod
def _build_boto3_client(
service_name: str,
client: Optional[ParameterClients] = None,
session: Optional[Type[boto3.Session]] = None,
config: Optional[Type[Config]] = None,
) -> Type[ParameterClients]:
"""Builds a low level boto3 client with session and config provided
Parameters
----------
service_name : str
AWS service name to instantiate a boto3 client, e.g. ssm
client : Optional[ParameterClients], optional
boto3 client instance, by default None
session : Optional[Type[boto3.Session]], optional
boto3 session instance, by default None
config : Optional[Type[Config]], optional
botocore config instance to configure client with, by default None
Returns
-------
Type[ParameterClients]
Instance of a boto3 client for Parameters feature (e.g., ssm, appconfig, secretsmanager, etc.)
"""
if client is not None:
user_agent.register_feature_to_client(client=client, feature="parameters")
return client

session = session or boto3.Session()
config = config or Config()
client = session.client(service_name=service_name, config=config)
user_agent.register_feature_to_client(client=client, feature="parameters")
return client

# maintenance: change DynamoDBServiceResource type to ParameterResourceClients when we expand
@staticmethod
def _build_boto3_resource_client(
service_name: str,
client: Optional["DynamoDBServiceResource"] = None,
session: Optional[Type[boto3.Session]] = None,
config: Optional[Type[Config]] = None,
endpoint_url: Optional[str] = None,
) -> "DynamoDBServiceResource":
"""Builds a high level boto3 resource client with session, config and endpoint_url provided
Parameters
----------
service_name : str
AWS service name to instantiate a boto3 client, e.g. ssm
client : Optional[DynamoDBServiceResource], optional
boto3 client instance, by default None
session : Optional[Type[boto3.Session]], optional
boto3 session instance, by default None
config : Optional[Type[Config]], optional
botocore config instance to configure client, by default None
Returns
-------
Type[DynamoDBServiceResource]
Instance of a boto3 resource client for Parameters feature (e.g., dynamodb, etc.)
"""
if client is not None:
user_agent.register_feature_to_resource(resource=client, feature="parameters")
return client

session = session or boto3.Session()
config = config or Config()
client = session.resource(service_name=service_name, config=config, endpoint_url=endpoint_url)
user_agent.register_feature_to_resource(resource=client, feature="parameters")
return client


def get_transform_method(value: str, transform: TransformOptions = None) -> Callable[..., Any]:
"""
Expand Down
17 changes: 6 additions & 11 deletions aws_lambda_powertools/utilities/parameters/dynamodb.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,7 @@
from .base import BaseProvider

if TYPE_CHECKING:
from mypy_boto3_dynamodb import DynamoDBServiceResource
from mypy_boto3_dynamodb.service_resource import Table
from mypy_boto3_dynamodb.service_resource import DynamoDBServiceResource


class DynamoDBProvider(BaseProvider):
Expand Down Expand Up @@ -162,19 +161,15 @@ def __init__(
"""
Initialize the DynamoDB client
"""
self.table: "Table" = self._build_boto3_resource_client(
service_name="dynamodb",
client=boto3_client,
session=boto3_session,
config=config,
endpoint_url=endpoint_url,
).Table(table_name)

if boto3_client is None:
boto3_session = boto3_session or boto3.session.Session()
boto3_client = boto3_session.resource("dynamodb", config=config, endpoint_url=endpoint_url)
self.table = boto3_client.Table(table_name)
self.key_attr = key_attr
self.sort_attr = sort_attr
self.value_attr = value_attr

super().__init__()
super().__init__(resource=boto3_client)

def _get(self, name: str, **sdk_options) -> str:
"""
Expand Down
Loading

0 comments on commit 26cfe7f

Please sign in to comment.