Skip to content

Commit 59dfd8b

Browse files
committed
feat: adding expiration time for secret cache in secret manager plugin
# Conflicts: # docs/using-the-python-driver/using-plugins/UsingTheAwsSecretsManagerPlugin.md
1 parent 0f65cba commit 59dfd8b

File tree

6 files changed

+53
-14
lines changed

6 files changed

+53
-14
lines changed

aws_advanced_python_wrapper/aws_secrets_manager_plugin.py

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,13 @@
1717
from json import JSONDecodeError, loads
1818
from re import search
1919
from types import SimpleNamespace
20-
from typing import TYPE_CHECKING, Callable, Dict, Optional, Set, Tuple
20+
from typing import TYPE_CHECKING, Callable, Optional, Set, Tuple
2121

2222
import boto3
2323
from botocore.exceptions import ClientError, EndpointConnectionError
2424

25+
from aws_advanced_python_wrapper.utils.cache_map import CacheMap
26+
2527
if TYPE_CHECKING:
2628
from boto3 import Session
2729
from aws_advanced_python_wrapper.driver_dialect import DriverDialect
@@ -46,8 +48,10 @@ class AwsSecretsManagerPlugin(Plugin):
4648
_SUBSCRIBED_METHODS: Set[str] = {"connect", "force_connect"}
4749

4850
_SECRETS_ARN_PATTERN = r"^arn:aws:secretsmanager:(?P<region>[^:\n]*):[^:\n]*:([^:/\n]*[:/])?(.*)$"
51+
_ONE_YEAR_IN_SECONDS = 60 * 60 * 24 * 365
4952

50-
_secrets_cache: Dict[Tuple, SimpleNamespace] = {}
53+
_secret: Optional[SimpleNamespace] = None
54+
_secrets_cache: CacheMap[Tuple, SimpleNamespace] = CacheMap()
5155
_secret_key: Tuple = ()
5256

5357
@property
@@ -94,7 +98,13 @@ def force_connect(
9498
return self._connect(props, force_connect_func)
9599

96100
def _connect(self, props: Properties, connect_func: Callable) -> Connection:
97-
secret_fetched: bool = self._update_secret()
101+
token_expiration_sec: int = WrapperProperties.SECRETS_MANAGER_EXPIRATION.get_int(props)
102+
# if value is less than 0, default to one year
103+
if token_expiration_sec < 0:
104+
token_expiration_sec = AwsSecretsManagerPlugin._ONE_YEAR_IN_SECONDS
105+
token_expiration_ns = token_expiration_sec * 1000
106+
107+
secret_fetched: bool = self._update_secret(token_expiration_ns=token_expiration_ns)
98108

99109
try:
100110
self._apply_secret_to_properties(props)
@@ -105,7 +115,7 @@ def _connect(self, props: Properties, connect_func: Callable) -> Connection:
105115
raise AwsWrapperError(
106116
Messages.get_formatted("AwsSecretsManagerPlugin.ConnectException", e)) from e
107117

108-
secret_fetched = self._update_secret(True)
118+
secret_fetched = self._update_secret(token_expiration_ns=token_expiration_ns, force_refetch=True)
109119

110120
if secret_fetched:
111121
try:
@@ -117,9 +127,10 @@ def _connect(self, props: Properties, connect_func: Callable) -> Connection:
117127
unhandled_error)) from unhandled_error
118128
raise AwsWrapperError(Messages.get_formatted("AwsSecretsManagerPlugin.FailedLogin", e)) from e
119129

120-
def _update_secret(self, force_refetch: bool = False) -> bool:
130+
def _update_secret(self, token_expiration_ns: int, force_refetch: bool = False) -> bool:
121131
"""
122132
Called to update credentials from the cache, or from the AWS Secrets Manager service.
133+
:param token_expiration_ns: Expiration time in nanoseconds for secret stored in cache.
123134
:param force_refetch: Allows ignoring cached credentials and force fetches the latest credentials from the service.
124135
:return: `True`, if credentials were fetched from the service.
125136
"""
@@ -135,7 +146,7 @@ def _update_secret(self, force_refetch: bool = False) -> bool:
135146
try:
136147
self._secret = self._fetch_latest_credentials()
137148
if self._secret:
138-
AwsSecretsManagerPlugin._secrets_cache[self._secret_key] = self._secret
149+
AwsSecretsManagerPlugin._secrets_cache.put(self._secret_key, self._secret, token_expiration_ns)
139150
fetched = True
140151
except (ClientError, AttributeError) as e:
141152
logger.debug("AwsSecretsManagerPlugin.FailedToFetchDbCredentials", e)

aws_advanced_python_wrapper/utils/cache_map.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -88,12 +88,12 @@ def _cleanup(self):
8888

8989

9090
class CacheItem(Generic[V]):
91-
def __init__(self, item: V, expiration_time: int):
91+
def __init__(self, item: V, expiration_time_ns: int):
9292
self.item = item
93-
self._expiration_time = expiration_time
93+
self._expiration_time_ns = expiration_time_ns
9494

9595
def __str__(self):
96-
return f"CacheItem [item={str(self.item)}, expiration_time={self._expiration_time}]"
96+
return f"CacheItem [item={str(self.item)}, expiration_time={self._expiration_time_ns}]"
9797

9898
def is_expired(self) -> bool:
99-
return time.perf_counter_ns() > self._expiration_time
99+
return time.perf_counter_ns() > self._expiration_time_ns

aws_advanced_python_wrapper/utils/iam_utils.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
from aws_advanced_python_wrapper.hostinfo import HostInfo
3232
from aws_advanced_python_wrapper.plugin_service import PluginService
3333
from boto3 import Session
34+
from types import SimpleNamespace
3435

3536
from aws_advanced_python_wrapper.utils.properties import (Properties,
3637
WrapperProperties)
@@ -132,3 +133,22 @@ def __init__(self, token: str, expiration: datetime):
132133

133134
def is_expired(self) -> bool:
134135
return datetime.now() > self._expiration
136+
137+
138+
class SecretInfo:
139+
@property
140+
def secret(self):
141+
return self._secret
142+
143+
@property
144+
def expiration(self):
145+
return self._expiration
146+
147+
def __init__(self, secret: SimpleNamespace, expiration: Optional[datetime] = None):
148+
self._secret = secret
149+
self._expiration = expiration
150+
151+
def is_expired(self) -> bool:
152+
if self._expiration is None:
153+
return False
154+
return datetime.now() > self._expiration

aws_advanced_python_wrapper/utils/properties.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,10 @@ class WrapperProperties:
146146
SECRETS_MANAGER_ENDPOINT = WrapperProperty(
147147
"secrets_manager_endpoint",
148148
"The endpoint of the secret to retrieve.")
149+
SECRETS_MANAGER_EXPIRATION = WrapperProperty(
150+
"secrets_manager_expiration",
151+
"Secret cache expiration in seconds",
152+
60 * 60 * 24 * 365)
149153

150154
DIALECT = WrapperProperty("wrapper_dialect", "A unique identifier for the supported database dialect.")
151155
AUXILIARY_QUERY_TIMEOUT_SEC = WrapperProperty(
@@ -264,7 +268,8 @@ class WrapperProperties:
264268
True)
265269

266270
# Host Selector
267-
ROUND_ROBIN_DEFAULT_WEIGHT = WrapperProperty("round_robin_default_weight", "The default weight for any hosts that have not been " +
271+
ROUND_ROBIN_DEFAULT_WEIGHT = WrapperProperty("round_robin_default_weight",
272+
"The default weight for any hosts that have not been " +
268273
"configured with the `round_robin_host_weight_pairs` parameter.",
269274
1)
270275

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ The following properties are required for the AWS Secrets Manager Connection Plu
2424
| `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` |
2525
| `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` |
2626
| `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` |
27+
| `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 |
2728

2829
*NOTE* A Secret ARN has the following format: `arn:aws:secretsmanager:<Region>:<AccountId>:secret:Secre78tName-6RandomCharacters`
2930

tests/unit/test_secrets_manager_plugin.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030

3131
from aws_advanced_python_wrapper.aws_secrets_manager_plugin import \
3232
AwsSecretsManagerPlugin
33+
from aws_advanced_python_wrapper.utils.cache_map import CacheMap
3334

3435
if TYPE_CHECKING:
3536
from boto3 import Session, client
@@ -38,7 +39,7 @@
3839
from aws_advanced_python_wrapper.plugin_service import PluginService
3940

4041
from types import SimpleNamespace
41-
from typing import Callable, Dict, Tuple
42+
from typing import Callable, Tuple
4243
from unittest import TestCase
4344
from unittest.mock import MagicMock, patch
4445

@@ -66,6 +67,7 @@ class TestAwsSecretsManagerPlugin(TestCase):
6667
_SECRET_CACHE_KEY = (_TEST_SECRET_ID, _TEST_REGION, _TEST_ENDPOINT)
6768
_TEST_HOST_INFO = HostInfo(_TEST_HOST, _TEST_PORT)
6869
_TEST_SECRET = SimpleNamespace(username="testUser", password="testPassword")
70+
_ONE_YEAR_IN_NANOSECONDS = 60 * 60 * 24 * 365 * 1000
6971

7072
_MYSQL_HOST_INFO = HostInfo("mysql.testdb.us-east-2.rds.amazonaws.com")
7173
_PG_HOST_INFO = HostInfo("pg.testdb.us-east-2.rds.amazonaws.com")
@@ -82,7 +84,7 @@ class TestAwsSecretsManagerPlugin(TestCase):
8284
}
8385
}, "some_operation")
8486

85-
_secrets_cache: Dict[Tuple, SimpleNamespace] = {}
87+
_secrets_cache: CacheMap[Tuple, SimpleNamespace] = CacheMap()
8688

8789
_mock_func: Callable
8890
_mock_plugin_service: PluginService
@@ -113,7 +115,7 @@ def setUp(self):
113115

114116
@patch("aws_advanced_python_wrapper.aws_secrets_manager_plugin.AwsSecretsManagerPlugin._secrets_cache", _secrets_cache)
115117
def test_connect_with_cached_secrets(self):
116-
self._secrets_cache[self._SECRET_CACHE_KEY] = self._TEST_SECRET
118+
self._secrets_cache.put(self._SECRET_CACHE_KEY, self._TEST_SECRET, self._ONE_YEAR_IN_NANOSECONDS)
117119
target_plugin: AwsSecretsManagerPlugin = AwsSecretsManagerPlugin(self._mock_plugin_service,
118120
self._properties,
119121
self._mock_session)

0 commit comments

Comments
 (0)