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

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 -1, default to one year
+        if token_expiration_sec == -1:
+            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

@@ -137,6 +137,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",
+        -1)
 
     DIALECT = WrapperProperty("wrapper_dialect", "A unique identifier for the supported database dialect.")
     AUXILIARY_QUERY_TIMEOUT_SEC = WrapperProperty(
@@ -255,7 +259,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/examples/MySQLSecretsManager.py

@@ -21,12 +21,24 @@
 if __name__ == "__main__":
     with AwsWrapperConnection.connect(
             mysql.connector.Connect,
-            host="database.cluster-xyz.us-east-1.rds.amazonaws.com",
+            host="atlas-mysql.cluster-cx422ywmsto6.us-east-2.rds.amazonaws.com",
             database="mysql",
-            secrets_manager_secret_id="arn:aws:secretsmanager:<Region>:<AccountId>:secret:Secre78tName-6RandomCharacters",
+            secrets_manager_secret_id="arn:aws:secretsmanager:us-east-2:851725167871:secret:mysql_test_1-ZjindE",
             secrets_manager_region="us-east-2",
             plugins="aws_secrets_manager"
     ) as awsconn, awsconn.cursor() as cursor:
         cursor.execute("SELECT @@aurora_server_id")
         for record in cursor.fetchone():
             print(record)
+    with AwsWrapperConnection.connect(
+            mysql.connector.Connect,
+            host="atlas-mysql.cluster-cx422ywmsto6.us-east-2.rds.amazonaws.com",
+            database="mysql",
+            secrets_manager_secret_id="arn:aws:secretsmanager:us-east-2:851725167871:secret:mysql_test_1-ZjindE",
+            secrets_manager_region="us-east-2",
+            plugins="aws_secrets_manager",
+            secrets_manager_expiration=50
+    ) as awsconn, awsconn.cursor() as cursor:
+        cursor.execute("SELECT @@aurora_server_id")
+        for record in cursor.fetchone():
+            print(record)

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

@@ -22,6 +22,7 @@ The following properties are required for the AWS Secrets Manager Connection Plu
 | `secrets_manager_secret_id` | String |                             Yes                             | Set this value to be the secret name or the secret ARN.                                                                                                                                                                          | `secret_id`             | `None`        |
 | `secrets_manager_region`    | String | Yes unless the `secrets_manager_secret_id` is a Secret ARN. | Set this value to be the region your secret is in.                                                                                                                                                                               | `us-east-2`             | `us-east-1`   |
 | `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_expiration`|  int   |                             No                              | Set this value to be the expiration time the secret is stored in the cache. If the value is -1, sets the expiration time to one year.                                                                                            | 500                     | `-1`          |
 
 *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 CacheItem, 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
 
@@ -64,6 +66,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")
@@ -80,7 +83,7 @@ class TestAwsSecretsManagerPlugin(TestCase):
         }
     }, "some_operation")
 
-    _secrets_cache: Dict[Tuple, SimpleNamespace] = {}
+    _secrets_cache: CacheMap[Tuple, SimpleNamespace] = CacheMap()
 
     _mock_func: Callable
     _mock_plugin_service: PluginService
@@ -111,7 +114,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)