feat(parameters): accept boto3_client to support private endpoints and ease testing (#1096)

Co-authored-by: Ran Isenberg <ran.isenberg@cyberark.com>
Co-authored-by: heitorlessa <lessa@amazon.co.uk>

Author: 3 people

aws_lambda_powertools/utilities/parameters/appconfig.py

@@ -4,12 +4,15 @@
 
 
 import os
-from typing import Any, Dict, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 from uuid import uuid4
 
 import boto3
 from botocore.config import Config
 
+if TYPE_CHECKING:
+    from mypy_boto3_appconfig import AppConfigClient
+
 from ...shared import constants
 from ...shared.functions import resolve_env_var_choice
 from .base import DEFAULT_MAX_AGE_SECS, DEFAULT_PROVIDERS, BaseProvider
@@ -30,7 +33,9 @@ class AppConfigProvider(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
+            Boto3 session to create a boto3_client from
+    boto3_client: AppConfigClient, optional
+            Boto3 AppConfig Client to use, boto3_session will be ignored if both are provided
 
     Example
     -------
@@ -68,22 +73,24 @@ def __init__(
         application: Optional[str] = None,
         config: Optional[Config] = None,
         boto3_session: Optional[boto3.session.Session] = None,
+        boto3_client: Optional["AppConfigClient"] = None,
     ):
         """
         Initialize the App Config client
         """
 
-        config = config or Config()
-        session = boto3_session or boto3.session.Session()
-        self.client = session.client("appconfig", config=config)
+        super().__init__()
+
+        self.client: "AppConfigClient" = self._build_boto3_client(
+            service_name="appconfig", client=boto3_client, session=boto3_session, config=config
+        )
+
         self.application = resolve_env_var_choice(
             choice=application, env=os.getenv(constants.SERVICE_NAME_ENV, "service_undefined")
         )
         self.environment = environment
         self.current_version = ""
 
-        super().__init__()
-
     def _get(self, name: str, **sdk_options) -> str:
         """
         Retrieve a parameter value from AWS App config.

aws_lambda_powertools/utilities/parameters/base.py

@@ -7,17 +7,28 @@
 from abc import ABC, abstractmethod
 from collections import namedtuple
 from datetime import datetime, timedelta
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Type, Union
+
+import boto3
+from botocore.config import Config
 
 from .exceptions import GetParameterError, TransformParameterError
 
+if TYPE_CHECKING:
+    from mypy_boto3_appconfig import AppConfigClient
+    from mypy_boto3_dynamodb import DynamoDBServiceResource
+    from mypy_boto3_secretsmanager import SecretsManagerClient
+    from mypy_boto3_ssm import SSMClient
+
+
 DEFAULT_MAX_AGE_SECS = 5
 ExpirableValue = namedtuple("ExpirableValue", ["value", "ttl"])
 # 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["AppConfigClient", "SecretsManagerClient", "SSMClient"]
 
 
 class BaseProvider(ABC):
@@ -180,6 +191,72 @@ def _get_multiple(self, path: str, **sdk_options) -> Dict[str, str]:
     def clear_cache(self):
         self.store.clear()
 
+    @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:
+            return client
+
+        session = session or boto3.Session()
+        config = config or Config()
+        return session.client(service_name=service_name, config=config)
+
+    # 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:
+            return client
+
+        session = session or boto3.Session()
+        config = config or Config()
+        return session.resource(service_name=service_name, config=config, endpoint_url=endpoint_url)
+
 
 def get_transform_method(key: str, transform: Optional[str] = None) -> Optional[str]:
     """

aws_lambda_powertools/utilities/parameters/dynamodb.py

@@ -3,14 +3,18 @@
 """
 
 
-from typing import Dict, Optional
+from typing import TYPE_CHECKING, Dict, Optional
 
 import boto3
 from boto3.dynamodb.conditions import Key
 from botocore.config import Config
 
 from .base import BaseProvider
 
+if TYPE_CHECKING:
+    from mypy_boto3_dynamodb import DynamoDBServiceResource
+    from mypy_boto3_dynamodb.service_resource import Table
+
 
 class DynamoDBProvider(BaseProvider):
     """
@@ -31,7 +35,9 @@ class DynamoDBProvider(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
+            Boto3 session to create a boto3_client from
+    boto3_client: DynamoDBServiceResource, optional
+            Boto3 DynamoDB Resource Client to use; boto3_session will be ignored if both are provided
 
     Example
     -------
@@ -152,15 +158,18 @@ def __init__(
         endpoint_url: Optional[str] = None,
         config: Optional[Config] = None,
         boto3_session: Optional[boto3.session.Session] = None,
+        boto3_client: Optional["DynamoDBServiceResource"] = None,
     ):
         """
         Initialize the DynamoDB client
         """
-
-        config = config or Config()
-        session = boto3_session or boto3.session.Session()
-
-        self.table = session.resource("dynamodb", endpoint_url=endpoint_url, config=config).Table(table_name)
+        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)
 
         self.key_attr = key_attr
         self.sort_attr = sort_attr
@@ -183,7 +192,9 @@ def _get(self, name: str, **sdk_options) -> str:
         # Explicit arguments will take precedence over keyword arguments
         sdk_options["Key"] = {self.key_attr: name}
 
-        return self.table.get_item(**sdk_options)["Item"][self.value_attr]
+        # maintenance: look for better ways to correctly type DynamoDB multiple return types
+        # without a breaking change within ABC return type
+        return self.table.get_item(**sdk_options)["Item"][self.value_attr]  # type: ignore[return-value]
 
     def _get_multiple(self, path: str, **sdk_options) -> Dict[str, str]:
         """
@@ -209,4 +220,6 @@ def _get_multiple(self, path: str, **sdk_options) -> Dict[str, str]:
             response = self.table.query(**sdk_options)
             items.extend(response.get("Items", []))
 
-        return {item[self.sort_attr]: item[self.value_attr] for item in items}
+        # maintenance: look for better ways to correctly type DynamoDB multiple return types
+        # without a breaking change within ABC return type
+        return {item[self.sort_attr]: item[self.value_attr] for item in items}  # type: ignore[misc]

aws_lambda_powertools/utilities/parameters/secrets.py

@@ -3,11 +3,14 @@
 """
 
 
-from typing import Any, Dict, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
 import boto3
 from botocore.config import Config
 
+if TYPE_CHECKING:
+    from mypy_boto3_secretsmanager import SecretsManagerClient
+
 from .base import DEFAULT_MAX_AGE_SECS, DEFAULT_PROVIDERS, BaseProvider
 
 
@@ -20,7 +23,9 @@ 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
+            Boto3 session to create a boto3_client from
+    boto3_client: SecretsManagerClient, optional
+            Boto3 SecretsManager Client to use, boto3_session will be ignored if both are provided
 
     Example
     -------
@@ -60,17 +65,22 @@ class SecretsProvider(BaseProvider):
 
     client: Any = None
 
-    def __init__(self, config: Optional[Config] = None, boto3_session: Optional[boto3.session.Session] = None):
+    def __init__(
+        self,
+        config: Optional[Config] = None,
+        boto3_session: Optional[boto3.session.Session] = None,
+        boto3_client: Optional["SecretsManagerClient"] = None,
+    ):
         """
         Initialize the Secrets Manager client
         """
 
-        config = config or Config()
-        session = boto3_session or boto3.session.Session()
-        self.client = session.client("secretsmanager", config=config)
-
         super().__init__()
 
+        self.client: "SecretsManagerClient" = self._build_boto3_client(
+            service_name="secretsmanager", client=boto3_client, session=boto3_session, config=config
+        )
+
     def _get(self, name: str, **sdk_options) -> str:
         """
         Retrieve a parameter value from AWS Systems Manager Parameter Store

aws_lambda_powertools/utilities/parameters/ssm.py

@@ -3,13 +3,16 @@
 """
 
 
-from typing import Any, Dict, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
 import boto3
 from botocore.config import Config
 
 from .base import DEFAULT_MAX_AGE_SECS, DEFAULT_PROVIDERS, BaseProvider
 
+if TYPE_CHECKING:
+    from mypy_boto3_ssm import SSMClient
+
 
 class SSMProvider(BaseProvider):
     """
@@ -20,7 +23,9 @@ 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
+            Boto3 session to create a boto3_client from
+    boto3_client: SSMClient, optional
+            Boto3 SSM Client to use, boto3_session will be ignored if both are provided
 
     Example
     -------
@@ -76,17 +81,22 @@ class SSMProvider(BaseProvider):
 
     client: Any = None
 
-    def __init__(self, config: Optional[Config] = None, boto3_session: Optional[boto3.session.Session] = None):
+    def __init__(
+        self,
+        config: Optional[Config] = None,
+        boto3_session: Optional[boto3.session.Session] = None,
+        boto3_client: Optional["SSMClient"] = None,
+    ):
         """
         Initialize the SSM Parameter Store client
         """
 
-        config = config or Config()
-        session = boto3_session or boto3.session.Session()
-        self.client = session.client("ssm", config=config)
-
         super().__init__()
 
+        self.client: "SSMClient" = self._build_boto3_client(
+            service_name="ssm", client=boto3_client, session=boto3_session, config=config
+        )
+
     # We break Liskov substitution principle due to differences in signatures of this method and superclass get method
     # We ignore mypy error, as changes to the signature here or in a superclass is a breaking change to users
     def get(  # type: ignore[override]