refactor: strict typing transform_value/method

Author: heitorlessa

aws_lambda_powertools/utilities/feature_flags/appconfig.py

@@ -15,8 +15,6 @@
 from .base import StoreProvider
 from .exceptions import ConfigurationStoreError, StoreClientError
 
-TRANSFORM_TYPE = "json"
-
 
 class AppConfigStore(StoreProvider):
     def __init__(
@@ -74,7 +72,7 @@ def get_raw_configuration(self) -> Dict[str, Any]:
                 dict,
                 self._conf_store.get(
                     name=self.name,
-                    transform=TRANSFORM_TYPE,
+                    transform="json",
                     max_age=self.cache_seconds,
                 ),
             )

aws_lambda_powertools/utilities/parameters/appconfig.py

@@ -9,6 +9,8 @@
 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
 
@@ -132,7 +134,7 @@ def get_app_config(
     name: str,
     environment: str,
     application: Optional[str] = None,
-    transform: Optional[str] = None,
+    transform: TransformOptions = None,
     force_fetch: bool = False,
     max_age: int = DEFAULT_MAX_AGE_SECS,
     **sdk_options

aws_lambda_powertools/utilities/parameters/base.py

@@ -7,11 +7,24 @@
 from abc import ABC, abstractmethod
 from collections import namedtuple
 from datetime import datetime, timedelta
-from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Type, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Optional,
+    Tuple,
+    Type,
+    Union,
+    cast,
+    overload,
+)
 
 import boto3
 from botocore.config import Config
 
+from aws_lambda_powertools.utilities.parameters.types import TransformOptions
+
 from .exceptions import GetParameterError, TransformParameterError
 
 if TYPE_CHECKING:
@@ -30,6 +43,14 @@
 SUPPORTED_TRANSFORM_METHODS = [TRANSFORM_METHOD_JSON, TRANSFORM_METHOD_BINARY]
 ParameterClients = Union["AppConfigDataClient", "SecretsManagerClient", "SSMClient"]
 
+TRANSFORM_METHOD_MAPPING = {
+    TRANSFORM_METHOD_JSON: json.loads,
+    TRANSFORM_METHOD_BINARY: base64.b64decode,
+    ".json": json.loads,
+    ".binary": base64.b64decode,
+    None: lambda x: x,
+}
+
 
 class BaseProvider(ABC):
     """
@@ -52,7 +73,7 @@ def get(
         self,
         name: str,
         max_age: int = DEFAULT_MAX_AGE_SECS,
-        transform: Optional[str] = None,
+        transform: TransformOptions = None,
         force_fetch: bool = False,
         **sdk_options,
     ) -> Optional[Union[str, dict, bytes]]:
@@ -124,7 +145,7 @@ def get_multiple(
         self,
         path: str,
         max_age: int = DEFAULT_MAX_AGE_SECS,
-        transform: Optional[str] = None,
+        transform: TransformOptions = None,
         raise_on_transform_error: bool = False,
         force_fetch: bool = False,
         **sdk_options,
@@ -170,13 +191,7 @@ def get_multiple(
             raise GetParameterError(str(exc))
 
         if transform:
-            transformed_values: dict = {}
-            for (item, value) in values.items():
-                _transform = get_transform_method(item, transform)
-                if not _transform:
-                    continue
-                transformed_values[item] = transform_value(value, _transform, raise_on_transform_error)
-            values.update(transformed_values)
+            values.update(transform_value(values, transform, raise_on_transform_error))
         self.store[key] = ExpirableValue(values, datetime.now() + timedelta(seconds=max_age))
 
         return values
@@ -258,7 +273,7 @@ def _build_boto3_resource_client(
         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]:
+def get_transform_method(key: str, transform: TransformOptions = None) -> Callable[..., Any]:
     """
     Determine the transform method
 
@@ -278,37 +293,50 @@ def get_transform_method(key: str, transform: Optional[str] = None) -> Optional[
     Parameters
     ---------
     key: str
-        Only used when the tranform is "auto".
+        Only used when the transform is "auto".
     transform: str, optional
         Original transform method, only "auto" will try to detect the transform method by the key
 
     Returns
     ------
-    Optional[str]:
-        The transform method either when transform is "auto" then None, "json" or "binary" is returned
-        or the original transform method
+    Callable:
+        Transform function could be json.loads, base64.b64decode, or a lambda that echo the str value
     """
-    if transform != "auto":
-        return transform
+    transform_method = TRANSFORM_METHOD_MAPPING.get(transform)
+
+    if transform == "auto":
+        key_suffix = key.rsplit(".")[-1]
+        transform_method = TRANSFORM_METHOD_MAPPING.get(key_suffix, TRANSFORM_METHOD_MAPPING[None])
+
+    return cast(Callable, transform_method)  # https://github.com/python/mypy/issues/10740
+
+
+@overload
+def transform_value(
+    value: Dict[str, Any], transform: TransformOptions, raise_on_transform_error: bool = False
+) -> Dict[str, Any]:
+    ...
 
-    for transform_method in SUPPORTED_TRANSFORM_METHODS:
-        if key.endswith("." + transform_method):
-            return transform_method
-    return None
+
+@overload
+def transform_value(
+    value: Union[str, bytes, Dict[str, Any]], transform: TransformOptions, raise_on_transform_error: bool = False
+) -> Optional[Union[str, bytes, Dict[str, Any]]]:
+    ...
 
 
 def transform_value(
-    value: str, transform: str, raise_on_transform_error: Optional[bool] = True
-) -> Optional[Union[dict, bytes]]:
+    value: Union[str, bytes, Dict[str, Any]], transform: TransformOptions, raise_on_transform_error: bool = False
+) -> Optional[Union[str, bytes, Dict[str, Any]]]:
     """
-    Apply a transform to a value
+    Transform a value using one of the available options.
 
     Parameters
     ---------
     value: str
         Parameter value to transform
     transform: str
-        Type of transform, supported values are "json" and "binary"
+        Type of transform, supported values are "json", "binary", and "auto" based on suffix (.json, .binary)
     raise_on_transform_error: bool, optional
         Raises an exception if any transform fails, otherwise this will
         return a None value for each transform that failed
@@ -318,18 +346,35 @@ def transform_value(
     TransformParameterError:
         When the parameter value could not be transformed
     """
+    # Maintenance: For v3, we should consider returning the original value for soft transform failures.
 
-    try:
-        if transform == TRANSFORM_METHOD_JSON:
-            return json.loads(value)
-        elif transform == TRANSFORM_METHOD_BINARY:
-            return base64.b64decode(value)
-        else:
-            raise ValueError(f"Invalid transform type '{transform}'")
+    err_msg = "Unable to transform value using '{transform}' transform: {exc}"
+
+    if isinstance(value, bytes):
+        value = value.decode("utf-8")
 
+    if isinstance(value, dict):
+        # NOTE: We must handle partial failures when receiving multiple values
+        # where one of the keys might fail during transform, e.g. `{"a": "valid", "b": "{"}`
+        # expected: `{"a": "valid", "b": None}`
+
+        transformed_values: Dict[str, Any] = {}
+        for dict_key, dict_value in value.items():
+            transform_method = get_transform_method(key=dict_key, transform=transform)
+            try:
+                transformed_values[dict_key] = transform_method(dict_value)
+            except Exception as exc:
+                if raise_on_transform_error:
+                    raise TransformParameterError(err_msg.format(transform=transform, exc=exc)) from exc
+                transformed_values[dict_key] = None
+        return transformed_values
+
+    try:
+        transform_method = get_transform_method(key=value, transform=transform)
+        return transform_method(value)
     except Exception as exc:
         if raise_on_transform_error:
-            raise TransformParameterError(str(exc))
+            raise TransformParameterError(err_msg.format(transform=transform, exc=exc)) from exc
         return None
 
 

aws_lambda_powertools/utilities/parameters/ssm.py

@@ -103,7 +103,7 @@ def get(  # type: ignore[override]
         self,
         name: str,
         max_age: int = DEFAULT_MAX_AGE_SECS,
-        transform: Optional[str] = None,
+        transform: TransformOptions = None,
         decrypt: bool = False,
         force_fetch: bool = False,
         **sdk_options,

tests/functional/test_utilities_parameters.py

@@ -14,7 +14,11 @@
 from botocore.response import StreamingBody
 
 from aws_lambda_powertools.utilities import parameters
-from aws_lambda_powertools.utilities.parameters.base import BaseProvider, ExpirableValue
+from aws_lambda_powertools.utilities.parameters.base import (
+    TRANSFORM_METHOD_MAPPING,
+    BaseProvider,
+    ExpirableValue,
+)
 
 
 @pytest.fixture(scope="function")
@@ -1863,17 +1867,6 @@ def test_transform_value_binary_exception():
     assert "Incorrect padding" in str(excinfo)
 
 
-def test_transform_value_wrong(mock_value):
-    """
-    Test transform_value() with an incorrect transform
-    """
-
-    with pytest.raises(parameters.TransformParameterError) as excinfo:
-        parameters.base.transform_value(mock_value, "INCORRECT")
-
-    assert "Invalid transform type" in str(excinfo)
-
-
 def test_transform_value_ignore_error(mock_value):
     """
     Test transform_value() does not raise errors when raise_on_transform_error is False
@@ -1884,35 +1877,14 @@ def test_transform_value_ignore_error(mock_value):
     assert value is None
 
 
-@pytest.mark.parametrize("original_transform", ["json", "binary", "other", "Auto", None])
-def test_get_transform_method_preserve_original(original_transform):
-    """
-    Check if original transform method is returned for anything other than "auto"
-    """
-    transform = parameters.base.get_transform_method("key", original_transform)
-
-    assert transform == original_transform
-
-
 @pytest.mark.parametrize("extension", ["json", "binary"])
 def test_get_transform_method_preserve_auto(extension, mock_name):
     """
     Check if we can auto detect the transform method by the support extensions json / binary
     """
     transform = parameters.base.get_transform_method(f"{mock_name}.{extension}", "auto")
 
-    assert transform == extension
-
-
-@pytest.mark.parametrize("key", ["json", "binary", "example", "example.jsonp"])
-def test_get_transform_method_preserve_auto_unhandled(key):
-    """
-    Check if any key that does not end with a supported extension returns None when
-    using the transform="auto"
-    """
-    transform = parameters.base.get_transform_method(key, "auto")
-
-    assert transform is None
+    assert transform == TRANSFORM_METHOD_MAPPING[extension]
 
 
 def test_base_provider_get_multiple_force_update(mock_name, mock_value):