feat(feature_flags): allow customers to bring their own boto3 client and session (#4717)

* feat: add `sdk_client` argument to AppConfigStore `__init__`

Signed-off-by: Adrián Tomás <22175056+adriantomas@users.noreply.github.com>

* style: fix typing

Signed-off-by: Adrián Tomás <22175056+adriantomas@users.noreply.github.com>

* test: modify tests including sdk_client

Signed-off-by: Adrián Tomás <22175056+adriantomas@users.noreply.github.com>

* docs: add sdk_client to feature flags docs

Signed-off-by: Adrián Tomás <22175056+adriantomas@users.noreply.github.com>

* Adding some small changes

* Addressing Adrian's feedback

---------

Signed-off-by: Adrián Tomás <22175056+adriantomas@users.noreply.github.com>
Co-authored-by: Leandro Damascena <lcdama@amazon.pt>

Author: adriantomas

aws_lambda_powertools/utilities/feature_flags/appconfig.py

@@ -1,7 +1,8 @@
 import logging
 import traceback
-from typing import Any, Dict, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union, cast
 
+import boto3
 from botocore.config import Config
 
 from aws_lambda_powertools.utilities import jmespath_utils
@@ -15,6 +16,9 @@
 from .base import StoreProvider
 from .exceptions import ConfigurationStoreError, StoreClientError
 
+if TYPE_CHECKING:
+    from mypy_boto3_appconfigdata import AppConfigDataClient
+
 
 class AppConfigStore(StoreProvider):
     def __init__(
@@ -27,6 +31,9 @@ def __init__(
         envelope: Optional[str] = "",
         jmespath_options: Optional[Dict] = None,
         logger: Optional[Union[logging.Logger, Logger]] = None,
+        boto_config: Optional[Config] = None,
+        boto3_session: Optional[boto3.session.Session] = None,
+        boto3_client: Optional["AppConfigDataClient"] = None,
     ):
         """This class fetches JSON schemas from AWS AppConfig
 
@@ -48,17 +55,29 @@ def __init__(
             Alternative JMESPath options to be included when filtering expr
         logger: A logging object
             Used to log messages. If None is supplied, one will be created.
+        boto_config: botocore.config.Config, optional
+            Botocore configuration to pass during client initialization
+        boto3_session : boto3.Session, optional
+            Boto3 session to use for AWS API communication
+        boto3_client : AppConfigDataClient, optional
+            Boto3 AppConfigDataClient Client to use, boto3_session and boto_config will be ignored if both are provided
         """
         super().__init__()
         self.logger = logger or logging.getLogger(__name__)
         self.environment = environment
         self.application = application
         self.name = name
         self.cache_seconds = max_age
-        self.config = sdk_config
+        self.config = sdk_config or boto_config
         self.envelope = envelope
         self.jmespath_options = jmespath_options
-        self._conf_store = AppConfigProvider(environment=environment, application=application, config=sdk_config)
+        self._conf_store = AppConfigProvider(
+            environment=environment,
+            application=application,
+            config=sdk_config or boto_config,
+            boto3_client=boto3_client,
+            boto3_session=boto3_session,
+        )
 
     @property
     def get_raw_configuration(self) -> Dict[str, Any]:

docs/utilities/feature_flags.md

@@ -496,16 +496,18 @@ AppConfig store provider fetches any JSON document from AWS AppConfig.
 
 These are the available options for further customization.
 
-| Parameter            | Default          | Description                                                                                                                                            |
-| -------------------- | ---------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------ |
-| **environment**      | `""`             | AWS AppConfig Environment, e.g. `dev`                                                                                                                  |
-| **application**      | `""`             | AWS AppConfig Application, e.g. `product-catalogue`                                                                                                    |
-| **name**             | `""`             | AWS AppConfig Configuration name, e.g `features`                                                                                                       |
-| **envelope**         | `None`           | JMESPath expression to use to extract feature flags configuration from AWS AppConfig configuration                                                     |
-| **max_age**          | `5`              | Number of seconds to cache feature flags configuration fetched from AWS AppConfig                                                                      |
-| **sdk_config**       | `None`           | [Botocore Config object](https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html){target="_blank"}                            |
+| Parameter            | Default          | Description                                                                                                                                                           |
+| -------------------- | ---------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| **environment**      | `""`             | AWS AppConfig Environment, e.g. `dev`                                                                                                                                 |
+| **application**      | `""`             | AWS AppConfig Application, e.g. `product-catalogue`                                                                                                                   |
+| **name**             | `""`             | AWS AppConfig Configuration name, e.g `features`                                                                                                                      |
+| **envelope**         | `None`           | JMESPath expression to use to extract feature flags configuration from AWS AppConfig configuration                                                                    |
+| **max_age**          | `5`              | Number of seconds to cache feature flags configuration fetched from AWS AppConfig                                                                                     |
 | **jmespath_options** | `None`           | For advanced use cases when you want to bring your own [JMESPath functions](https://github.com/jmespath/jmespath.py#custom-functions){target="_blank" rel="nofollow"} |
-| **logger**           | `logging.Logger` | Logger to use for debug.  You can optionally supply an instance of Powertools for AWS Lambda (Python) Logger.                                          |
+| **logger**           | `logging.Logger` | Logger to use for debug.  You can optionally supply an instance of Powertools for AWS Lambda (Python) Logger.                                                         |
+| **boto3_client**     | `None`           | [AppConfigData boto3 client](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/appconfigdata.html#AppConfigData.Client){target="_blank"}     |
+| **boto3_session**    | `None`           | [Boto3 session](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html){target="_blank"}     |
+| **boto_config**     | `None`            | [Botocore config](https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html){target="_blank"}     |
 
 === "appconfig_provider_options.py"
 
@@ -525,6 +527,27 @@ These are the available options for further customization.
     --8<-- "examples/feature_flags/src/appconfig_provider_options_features.json"
     ```
 
+#### Customizing boto configuration
+
+<!-- markdownlint-disable MD013 -->
+The **`boto_config`** , **`boto3_session`**, and **`boto3_client`**  parameters enable you to pass in a custom [botocore config object](https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html){target="_blank"}, [boto3 session](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html){target="_blank"}, or  a [boto3 client](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/boto3.html){target="_blank"} when constructing the AppConfig store provider.
+<!-- markdownlint-enable MD013 -->
+
+=== "custom_boto_session_feature_flags.py"
+    ```python hl_lines="8 14"
+    --8<-- "examples/feature_flags/src/custom_boto_session_feature_flags.py"
+    ```
+
+=== "custom_boto_config_feature_flags.py"
+    ```python hl_lines="8 14"
+    --8<-- "examples/feature_flags/src/custom_boto_config_feature_flags.py"
+    ```
+
+=== "custom_boto_client_feature_flags.py"
+    ```python hl_lines="8 14"
+    --8<-- "examples/feature_flags/src/custom_boto_client_feature_flags.py"
+    ```
+
 ### Create your own store provider
 
 You can create your own custom FeatureFlags store provider by inheriting the `StoreProvider` class, and implementing both `get_raw_configuration()` and `get_configuration()` methods to retrieve the configuration from your custom store.

examples/feature_flags/src/appconfig_provider_options.py

@@ -26,7 +26,7 @@ def _func_special_decoder(self, features):
     name="features",
     max_age=120,
     envelope="special_decoder(features)",  # using a custom function defined in CustomFunctions Class
-    sdk_config=boto_config,
+    boto_config=boto_config,
     jmespath_options=custom_jmespath_options,
 )
 

examples/feature_flags/src/custom_boto_client_feature_flags.py

@@ -0,0 +1,29 @@
+from typing import Any
+
+import boto3
+
+from aws_lambda_powertools.utilities.feature_flags import AppConfigStore, FeatureFlags
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+boto3_client = boto3.client("appconfigdata")
+
+app_config = AppConfigStore(
+    environment="dev",
+    application="product-catalogue",
+    name="features",
+    boto3_client=boto3_client,
+)
+
+feature_flags = FeatureFlags(store=app_config)
+
+
+def lambda_handler(event: dict, context: LambdaContext):
+    apply_discount: Any = feature_flags.evaluate(name="ten_percent_off_campaign", default=False)
+
+    price: Any = event.get("price")
+
+    if apply_discount:
+        # apply 10% discount to product
+        price = price * 0.9
+
+    return {"price": price}

examples/feature_flags/src/custom_boto_config_feature_flags.py

@@ -0,0 +1,29 @@
+from typing import Any
+
+from botocore.config import Config
+
+from aws_lambda_powertools.utilities.feature_flags import AppConfigStore, FeatureFlags
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+boto_config = Config(read_timeout=10, retries={"total_max_attempts": 2})
+
+app_config = AppConfigStore(
+    environment="dev",
+    application="product-catalogue",
+    name="features",
+    boto_config=boto_config,
+)
+
+feature_flags = FeatureFlags(store=app_config)
+
+
+def lambda_handler(event: dict, context: LambdaContext):
+    apply_discount: Any = feature_flags.evaluate(name="ten_percent_off_campaign", default=False)
+
+    price: Any = event.get("price")
+
+    if apply_discount:
+        # apply 10% discount to product
+        price = price * 0.9
+
+    return {"price": price}

examples/feature_flags/src/custom_boto_session_feature_flags.py

@@ -0,0 +1,29 @@
+from typing import Any
+
+import boto3
+
+from aws_lambda_powertools.utilities.feature_flags import AppConfigStore, FeatureFlags
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+boto3_session = boto3.session.Session()
+
+app_config = AppConfigStore(
+    environment="dev",
+    application="product-catalogue",
+    name="features",
+    boto3_session=boto3_session,
+)
+
+feature_flags = FeatureFlags(store=app_config)
+
+
+def lambda_handler(event: dict, context: LambdaContext):
+    apply_discount: Any = feature_flags.evaluate(name="ten_percent_off_campaign", default=False)
+
+    price: Any = event.get("price")
+
+    if apply_discount:
+        # apply 10% discount to product
+        price = price * 0.9
+
+    return {"price": price}

tests/functional/feature_flags/_boto3/test_feature_flags.py

@@ -1,7 +1,12 @@
+from io import BytesIO
+from json import dumps
 from typing import Dict, List, Optional
 
+import boto3
 import pytest
 from botocore.config import Config
+from botocore.response import StreamingBody
+from botocore.stub import Stubber
 
 from aws_lambda_powertools.utilities.feature_flags import (
     ConfigurationStoreError,
@@ -37,17 +42,46 @@ def init_feature_flags(
     envelope: str = "",
     jmespath_options: Optional[Dict] = None,
 ) -> FeatureFlags:
-    mocked_get_conf = mocker.patch("aws_lambda_powertools.utilities.parameters.AppConfigProvider.get")
-    mocked_get_conf.return_value = mock_schema
+    environment = "test_env"
+    application = "test_app"
+    name = "test_conf_name"
+    configuration_token = "foo"
+    mock_schema_to_bytes = dumps(mock_schema).encode()
+
+    client = boto3.client("appconfigdata", config=config)
+    stubber = Stubber(client)
+
+    stubber.add_response(
+        method="start_configuration_session",
+        expected_params={
+            "ConfigurationProfileIdentifier": name,
+            "ApplicationIdentifier": application,
+            "EnvironmentIdentifier": environment,
+        },
+        service_response={"InitialConfigurationToken": configuration_token},
+    )
+    stubber.add_response(
+        method="get_latest_configuration",
+        expected_params={"ConfigurationToken": configuration_token},
+        service_response={
+            "Configuration": StreamingBody(
+                raw_stream=BytesIO(mock_schema_to_bytes),
+                content_length=len(mock_schema_to_bytes),
+            ),
+            "NextPollConfigurationToken": configuration_token,
+        },
+    )
+    stubber.activate()
 
     app_conf_fetcher = AppConfigStore(
-        environment="test_env",
-        application="test_app",
-        name="test_conf_name",
+        environment=environment,
+        application=application,
+        name=name,
         max_age=600,
-        sdk_config=config,
         envelope=envelope,
         jmespath_options=jmespath_options,
+        boto_config=config,
+        boto3_client=client,
     )
     feature_flags: FeatureFlags = FeatureFlags(store=app_conf_fetcher)
     return feature_flags