Skip to content

refactor(typing): enable boto3 implicit type annotations #4692

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Aug 6, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
Loading