feat: adding expiration time for secret cache in secret manager plugin

# Conflicts:
#	docs/using-the-python-driver/using-plugins/UsingTheAwsSecretsManagerPlugin.md

Author: JuanLeee

aws_advanced_python_wrapper/aws_secrets_manager_plugin.py

@@ -17,9 +17,10 @@
 from json import JSONDecodeError, loads
 from re import search
 from types import SimpleNamespace
-from typing import TYPE_CHECKING, Callable, Dict, Optional, Set, Tuple
+from typing import TYPE_CHECKING, Callable, Optional, Set, Tuple
 
 import boto3
+from aws_advanced_python_wrapper.utils.cache_map import CacheMap
 from botocore.exceptions import ClientError, EndpointConnectionError
 
 if TYPE_CHECKING:
@@ -46,8 +47,10 @@ class AwsSecretsManagerPlugin(Plugin):
     _SUBSCRIBED_METHODS: Set[str] = {"connect", "force_connect"}
 
     _SECRETS_ARN_PATTERN = r"^arn:aws:secretsmanager:(?P<region>[^:\n]*):[^:\n]*:([^:/\n]*[:/])?(.*)$"
+    _ONE_YEAR_IN_SECONDS = 60 * 60 * 24 * 365
 
-    _secrets_cache: Dict[Tuple, SimpleNamespace] = {}
+    _secret: Optional[SimpleNamespace] = None
+    _secrets_cache: CacheMap[Tuple, SimpleNamespace] = CacheMap()
     _secret_key: Tuple = ()
 
     @property
@@ -94,7 +97,13 @@ def force_connect(
         return self._connect(props, force_connect_func)
 
     def _connect(self, props: Properties, connect_func: Callable) -> Connection:
-        secret_fetched: bool = self._update_secret()
+        token_expiration_sec: int = WrapperProperties.SECRETS_MANAGER_EXPIRATION.get_int(props)
+        # if value is less than 0, default to one year
+        if token_expiration_sec < 0:
+            token_expiration_sec = AwsSecretsManagerPlugin._ONE_YEAR_IN_SECONDS
+        token_expiration_ns = token_expiration_sec * 1000
+
+        secret_fetched: bool = self._update_secret(token_expiration_ns=token_expiration_ns)
 
         try:
             self._apply_secret_to_properties(props)
@@ -105,7 +114,7 @@ def _connect(self, props: Properties, connect_func: Callable) -> Connection:
                 raise AwsWrapperError(
                     Messages.get_formatted("AwsSecretsManagerPlugin.ConnectException", e)) from e
 
-            secret_fetched = self._update_secret(True)
+            secret_fetched = self._update_secret(token_expiration_ns=token_expiration_ns, force_refetch=True)
 
             if secret_fetched:
                 try:
@@ -117,9 +126,10 @@ def _connect(self, props: Properties, connect_func: Callable) -> Connection:
                                                unhandled_error)) from unhandled_error
             raise AwsWrapperError(Messages.get_formatted("AwsSecretsManagerPlugin.FailedLogin", e)) from e
 
-    def _update_secret(self, force_refetch: bool = False) -> bool:
+    def _update_secret(self, token_expiration_ns: int, force_refetch: bool = False) -> bool:
         """
         Called to update credentials from the cache, or from the AWS Secrets Manager service.
+        :param token_expiration_ns: Expiration time in nanoseconds for secret stored in cache.
         :param force_refetch: Allows ignoring cached credentials and force fetches the latest credentials from the service.
         :return: `True`, if credentials were fetched from the service.
         """
@@ -135,7 +145,7 @@ def _update_secret(self, force_refetch: bool = False) -> bool:
                 try:
                     self._secret = self._fetch_latest_credentials()
                     if self._secret:
-                        AwsSecretsManagerPlugin._secrets_cache[self._secret_key] = self._secret
+                        AwsSecretsManagerPlugin._secrets_cache.put(self._secret_key, self._secret, token_expiration_ns)
                         fetched = True
                 except (ClientError, AttributeError) as e:
                     logger.debug("AwsSecretsManagerPlugin.FailedToFetchDbCredentials", e)

aws_advanced_python_wrapper/utils/cache_map.py

@@ -88,12 +88,12 @@ def _cleanup(self):
 
 
 class CacheItem(Generic[V]):
-    def __init__(self, item: V, expiration_time: int):
+    def __init__(self, item: V, expiration_time_ns: int):
         self.item = item
-        self._expiration_time = expiration_time
+        self._expiration_time_ns = expiration_time_ns
 
     def __str__(self):
-        return f"CacheItem [item={str(self.item)}, expiration_time={self._expiration_time}]"
+        return f"CacheItem [item={str(self.item)}, expiration_time={self._expiration_time_ns}]"
 
     def is_expired(self) -> bool:
-        return time.perf_counter_ns() > self._expiration_time
+        return time.perf_counter_ns() > self._expiration_time_ns

aws_advanced_python_wrapper/utils/iam_utils.py

@@ -31,6 +31,7 @@
     from aws_advanced_python_wrapper.hostinfo import HostInfo
     from aws_advanced_python_wrapper.plugin_service import PluginService
     from boto3 import Session
+    from types import SimpleNamespace
 
 from aws_advanced_python_wrapper.utils.properties import (Properties,
                                                           WrapperProperties)
@@ -132,3 +133,22 @@ def __init__(self, token: str, expiration: datetime):
 
     def is_expired(self) -> bool:
         return datetime.now() > self._expiration
+
+
+class SecretInfo:
+    @property
+    def secret(self):
+        return self._secret
+
+    @property
+    def expiration(self):
+        return self._expiration
+
+    def __init__(self, secret: SimpleNamespace, expiration: Optional[datetime] = None):
+        self._secret = secret
+        self._expiration = expiration
+
+    def is_expired(self) -> bool:
+        if self._expiration is None:
+            return False
+        return datetime.now() > self._expiration

aws_advanced_python_wrapper/utils/properties.py

@@ -146,6 +146,10 @@ class WrapperProperties:
     SECRETS_MANAGER_ENDPOINT = WrapperProperty(
         "secrets_manager_endpoint",
         "The endpoint of the secret to retrieve.")
+    SECRETS_MANAGER_EXPIRATION = WrapperProperty(
+        "secrets_manager_expiration",
+        "Secret cache expiration in seconds",
+        60 * 60 * 24 * 365)
 
     DIALECT = WrapperProperty("wrapper_dialect", "A unique identifier for the supported database dialect.")
     AUXILIARY_QUERY_TIMEOUT_SEC = WrapperProperty(
@@ -264,7 +268,8 @@ class WrapperProperties:
         True)
 
     # Host Selector
-    ROUND_ROBIN_DEFAULT_WEIGHT = WrapperProperty("round_robin_default_weight", "The default weight for any hosts that have not been " +
+    ROUND_ROBIN_DEFAULT_WEIGHT = WrapperProperty("round_robin_default_weight",
+                                                 "The default weight for any hosts that have not been " +
                                                  "configured with the `round_robin_host_weight_pairs` parameter.",
                                                  1)
 

docs/using-the-python-driver/using-plugins/UsingTheAwsSecretsManagerPlugin.md

@@ -24,6 +24,7 @@ The following properties are required for the AWS Secrets Manager Connection Plu
 | `secrets_manager_endpoint`        | String |                             No                              | Set this value to be the endpoint override to retrieve your secret from. This parameter value should be in the form of a URL, with a valid protocol (ex. `http://`) and domain (ex. `localhost`). A port number is not required. | `http://localhost:1234` | `None`        |
 | `secrets_manager_secret_username` | String |                             No                              | Set this value to be the key in the JSON secret that contains the username for database connection.                                                                                                                              | `username_key`          | `username`    |
 | `secrets_manager_secret_password` | String |                             No                              | SSet this value to be the key in the JSON secret that contains the password for database connection.                                                                                                                             | `password_key`          | `password`    |
+| `secrets_manager_expiration`      |  int   |                             No                              | Set this value to be the expiration time the secret is stored in the cache. If the value is below 0, sets the expiration time to one year.                                                                                       | 500                     | 31536000      |
 
 *NOTE* A Secret ARN has the following format: `arn:aws:secretsmanager:<Region>:<AccountId>:secret:Secre78tName-6RandomCharacters`
 

tests/unit/test_secrets_manager_plugin.py

@@ -28,6 +28,8 @@
 
 from typing import TYPE_CHECKING
 
+from aws_advanced_python_wrapper.utils.cache_map import CacheMap
+
 from aws_advanced_python_wrapper.aws_secrets_manager_plugin import \
     AwsSecretsManagerPlugin
 
@@ -38,7 +40,7 @@
     from aws_advanced_python_wrapper.plugin_service import PluginService
 
 from types import SimpleNamespace
-from typing import Callable, Dict, Tuple
+from typing import Callable, Tuple
 from unittest import TestCase
 from unittest.mock import MagicMock, patch
 
@@ -66,6 +68,7 @@ class TestAwsSecretsManagerPlugin(TestCase):
     _SECRET_CACHE_KEY = (_TEST_SECRET_ID, _TEST_REGION, _TEST_ENDPOINT)
     _TEST_HOST_INFO = HostInfo(_TEST_HOST, _TEST_PORT)
     _TEST_SECRET = SimpleNamespace(username="testUser", password="testPassword")
+    _ONE_YEAR_IN_NANOSECONDS = 60 * 60 * 24 * 365 * 1000
 
     _MYSQL_HOST_INFO = HostInfo("mysql.testdb.us-east-2.rds.amazonaws.com")
     _PG_HOST_INFO = HostInfo("pg.testdb.us-east-2.rds.amazonaws.com")
@@ -82,7 +85,7 @@ class TestAwsSecretsManagerPlugin(TestCase):
         }
     }, "some_operation")
 
-    _secrets_cache: Dict[Tuple, SimpleNamespace] = {}
+    _secrets_cache: CacheMap[Tuple, SimpleNamespace] = CacheMap()
 
     _mock_func: Callable
     _mock_plugin_service: PluginService
@@ -113,7 +116,7 @@ def setUp(self):
 
     @patch("aws_advanced_python_wrapper.aws_secrets_manager_plugin.AwsSecretsManagerPlugin._secrets_cache", _secrets_cache)
     def test_connect_with_cached_secrets(self):
-        self._secrets_cache[self._SECRET_CACHE_KEY] = self._TEST_SECRET
+        self._secrets_cache.put(self._SECRET_CACHE_KEY, self._TEST_SECRET, self._ONE_YEAR_IN_NANOSECONDS)
         target_plugin: AwsSecretsManagerPlugin = AwsSecretsManagerPlugin(self._mock_plugin_service,
                                                                          self._properties,
                                                                          self._mock_session)