feat: add parameter utility

aws_lambda_powertools/utilities/__init__.py

@@ -0,0 +1,3 @@
+# -*- coding: utf-8 -*-
+
+"""General utilities for Powertools"""

aws_lambda_powertools/utilities/parameters/__init__.py

@@ -0,0 +1,247 @@
+# -*- coding: utf-8 -*-
+
+"""
+Parameter retrieval and caching utility
+"""
+
+from typing import Dict, Optional, Union
+
+import boto3
+from boto3.dynamodb.conditions import Key
+
+from .base import BaseProvider, GetParameterError
+
+__all__ = [
+    "BaseProvider",
+    "GetParameterError",
+    "DynamoDBProvider",
+    "SecretsProvider",
+    "SSMProvider",
+    "get_parameter",
+    "get_parameters",
+    "get_secret",
+]
+
+
+class SSMProvider(BaseProvider):
+    """
+    AWS Systems Manager Parameter Store Provider
+    """
+
+    client = None
+
+    def __init__(
+        self, region: Optional[str] = None,
+    ):
+        """
+        Initialize the SSM Parameter Store client
+        """
+
+        client_kwargs = {}
+        if region:
+            client_kwargs["region_name"] = region
+
+        self.client = boto3.client("ssm", **client_kwargs)
+
+        super().__init__()
+
+    def _get(self, name: str, **kwargs) -> str:
+        """
+        Retrieve a parameter value from AWS Systems Manager Parameter Store
+
+        Parameters
+        ----------
+        name: str
+            Parameter name
+        decrypt: bool
+            If the parameter value should be decrypted
+        """
+
+        # Load kwargs
+        decrypt = kwargs.get("decrypt", False)
+
+        return self.client.get_parameter(Name=name, WithDecryption=decrypt)["Parameter"]["Value"]
+
+    def _get_multiple(self, path: str, **kwargs) -> Dict[str, str]:
+        """
+        Retrieve multiple parameter values from AWS Systems Manager Parameter Store
+
+        Parameters
+        ----------
+        path: str
+            Path to retrieve the parameters
+        decrypt: bool
+            If the parameter values should be decrypted
+        recursive: bool
+            If this should retrieve the parameter values recursively or not
+        """
+
+        # Load kwargs
+        decrypt = kwargs.get("decrypt", False)
+        recursive = kwargs.get("recursive", False)
+
+        response = self.client.get_parameters_by_path(Path=path, WithDecryption=decrypt, Recursive=recursive)
+        parameters = response.get("Parameters", [])
+
+        # Keep retrieving parameters
+        while "NextToken" in response:
+            response = self.client.get_parameters_by_path(
+                Path=path, WithDecryption=decrypt, Recursive=recursive, NextToken=response["NextToken"]
+            )
+            parameters.extend(response.get("Parameters", []))
+
+        retval = {}
+        for parameter in parameters:
+
+            # Standardize the parameter name
+            # The parameter name returned by SSM will contained the full path.
+            # However, for readability, we should return only the part after
+            # the path.
+            name = parameter["Name"]
+            if name.startswith(path):
+                name = name[len(path) :]
+            name = name.lstrip("/")
+
+            retval[name] = parameter["Value"]
+
+        return retval
+
+
+class SecretsProvider(BaseProvider):
+    """
+    AWS Secrets Manager Parameter Provider
+    """
+
+    client = None
+
+    def __init__(self, region: Optional[str] = None):
+        """
+        Initialize the Secrets Manager client
+        """
+
+        client_kwargs = {}
+        if region:
+            client_kwargs["region_name"] = region
+
+        self.client = boto3.client("secretsmanager", **client_kwargs)
+
+        super().__init__()
+
+    def _get(self, name: str, **kwargs) -> str:
+        """
+        Retrieve a parameter value from AWS Systems Manager Parameter Store
+        """
+
+        return self.client.get_secret_value(SecretId=name)["SecretString"]
+
+    def _get_multiple(self, path: str, **kwargs) -> Dict[str, str]:
+        """
+        Retrieving multiple parameter values is not supported with AWS Secrets Manager
+        """
+        raise NotImplementedError()
+
+
+class DynamoDBProvider(BaseProvider):
+    """
+    Amazon DynamoDB Parameter Provider
+    """
+
+    table = None
+    key_attr = None
+    value_attr = None
+
+    def __init__(
+        self, table_name: str, key_attr: str = "id", value_attr: str = "value", region: Optional[str] = None,
+    ):
+        """
+        Initialize the DynamoDB client
+        """
+
+        client_kwargs = {}
+        if region:
+            client_kwargs["region_name"] = region
+        self.table = boto3.resource("dynamodb", **client_kwargs).Table(table_name)
+
+        self.key_attr = key_attr
+        self.value_attr = value_attr
+
+        super().__init__()
+
+    def _get(self, name: str, **kwargs) -> str:
+        """
+        Retrieve a parameter value from Amazon DynamoDB
+        """
+
+        return self.table.get_item(Key={self.key_attr: name})["Item"][self.value_attr]
+
+    def _get_multiple(self, path: str, **kwargs) -> Dict[str, str]:
+        """
+        Retrieve multiple parameter values from Amazon DynamoDB
+
+        Parameters
+        ----------
+        path: str
+            Path to retrieve the parameters
+        sort_attr: str
+            Name of the DynamoDB table sort key (defaults to 'sk')
+        """
+
+        sort_attr = kwargs.get("sort_attr", "sk")
+
+        response = self.table.query(KeyConditionExpression=Key(self.key_attr).eq(path))
+        items = response.get("Items", [])
+
+        # Keep querying while there are more items matching the partition key
+        while "LastEvaluatedKey" in response:
+            response = self.table.query(
+                KeyConditionExpression=Key(self.key_attr).eq(path), ExclusiveStartKey=response["LastEvaluatedKey"],
+            )
+            items.extend(response.get("Items", []))
+
+        retval = {}
+        for item in items:
+            retval[item[sort_attr]] = item[self.value_attr]
+
+        return retval
+
+
+# These providers will be dynamically initialized on first use of the helper functions
+_DEFAULT_PROVIDERS = {}
+
+
+def get_parameter(name: str, transform: Optional[str] = None) -> Union[str, list, dict, bytes]:
+    """
+    Retrieve a parameter value from AWS Systems Manager (SSM) Parameter Store
+    """
+
+    # Only create the provider if this function is called at least once
+    if "ssm" not in _DEFAULT_PROVIDERS:
+        _DEFAULT_PROVIDERS["ssm"] = SSMProvider()
+
+    return _DEFAULT_PROVIDERS["ssm"].get(name, transform=transform)
+
+
+def get_parameters(
+    path: str, transform: Optional[str] = None, recursive: bool = False, decrypt: bool = False
+) -> Union[Dict[str, str], Dict[str, dict], Dict[str, bytes]]:
+    """
+    Retrieve multiple parameter values from AWS Systems Manager (SSM) Parameter Store
+    """
+
+    # Only create the provider if this function is called at least once
+    if "ssm" not in _DEFAULT_PROVIDERS:
+        _DEFAULT_PROVIDERS["ssm"] = SSMProvider()
+
+    return _DEFAULT_PROVIDERS["ssm"].get_multiple(path, transform=transform, recursive=recursive, decrypt=decrypt)
+
+
+def get_secret(name: str, transform: Optional[str] = None, decrypt: bool = False) -> Union[str, dict, bytes]:
+    """
+    Retrieve a parameter value from AWS Secrets Manager
+    """
+
+    # Only create the provider if this function is called at least once
+    if "secrets" not in _DEFAULT_PROVIDERS:
+        _DEFAULT_PROVIDERS["secrets"] = SecretsProvider()
+
+    return _DEFAULT_PROVIDERS["secrets"].get(name, transform=transform, decrypt=decrypt)

aws_lambda_powertools/utilities/parameters/base.py

@@ -0,0 +1,124 @@
+"""
+Base for Parameter providers
+"""
+
+import base64
+import json
+from abc import ABC, abstractmethod
+from collections import namedtuple
+from datetime import datetime, timedelta
+from typing import Dict, Optional, Union
+
+DEFAULT_MAX_AGE_SECS = 5
+ExpirableValue = namedtuple("ExpirableValue", ["value", "ttl"])
+
+
+class GetParameterError(Exception):
+    """When a provider raises an exception on parameter retrieval"""
+
+
+class BaseProvider(ABC):
+    """
+    Abstract Base Class for Parameter providers
+    """
+
+    store = None
+
+    def __init__(self):
+        """
+        Initialize the base provider
+        """
+
+        self.store = {}
+
+    def get(
+        self, name: str, max_age: int = DEFAULT_MAX_AGE_SECS, transform: Optional[str] = None, **kwargs
+    ) -> Union[str, list, dict, bytes]:
+        """
+        Retrieve a parameter value or return the cached value
+
+        Parameters
+        ----------
+
+        name: str
+            Parameter name
+        max_age: int
+            Maximum age of the cached value
+        transform: str
+            Optional transformation of the parameter value. Supported values
+            are "json" for JSON strings and "binary" for base 64 encoded
+            values.
+
+        Raises
+        ------
+
+        GetParameterError
+            When the parameter provider fails to retrieve a parameter value for
+            a given name.
+        """
+
+        # If there are multiple calls to the same parameter but in a different
+        # transform, they will be stored multiple times. This allows us to
+        # optimize by transforming the data only once per retrieval, thus there
+        # is no need to transform cached values multiple times. However, this
+        # means that we need to make multiple calls to the underlying parameter
+        # store if we need to return it in different transforms. Since the number
+        # of supported transform is small and the probability that a given
+        # parameter will always be used in a specific transform, this should be
+        # an acceptable tradeoff.
+        key = (name, transform)
+
+        if key not in self.store or self.store[key].ttl < datetime.now():
+            try:
+                value = self._get(name, **kwargs)
+            # Encapsulate all errors into a generic GetParameterError
+            except Exception as exc:
+                raise GetParameterError(str(exc))
+
+            if transform == "json":
+                value = json.loads(value)
+            elif transform == "binary":
+                value = base64.b64decode(value)
+
+            self.store[key] = ExpirableValue(value, datetime.now() + timedelta(seconds=max_age),)
+
+        return self.store[key].value
+
+    @abstractmethod
+    def _get(self, name: str, **kwargs) -> str:
+        """
+        Retrieve paramater value from the underlying parameter store
+        """
+        raise NotImplementedError()
+
+    def get_multiple(
+        self, path: str, max_age: int = DEFAULT_MAX_AGE_SECS, transform: Optional[str] = None, **kwargs
+    ) -> Union[Dict[str, str], Dict[str, dict], Dict[str, bytes]]:
+        """
+        Retrieve multiple parameters based on a path prefix
+        """
+
+        key = (path, transform)
+
+        if key not in self.store or self.store[key].ttl < datetime.now():
+            try:
+                values = self._get_multiple(path, **kwargs)
+            # Encapsulate all errors into a generic GetParameterError
+            except Exception as exc:
+                raise GetParameterError(str(exc))
+
+            if transform == "json":
+                values = {k: json.loads(v) for k, v in values.items()}
+            elif transform == "binary":
+                values = {k: base64.b64decode(v) for k, v in values.items()}
+
+            self.store[key] = ExpirableValue(values, datetime.now() + timedelta(seconds=max_age),)
+
+        return self.store[key].value
+
+    @abstractmethod
+    def _get_multiple(self, path: str, **kwargs) -> Dict[str, str]:
+        """
+        Retrieve multiple parameter values from the underlying parameter store
+        """
+        raise NotImplementedError()