Skip to content

Commit 05d8810

Browse files
massimiliano96Massimiliano Riva
massimiliano96
and
Massimiliano Riva
authored
feat: allow custom secret keys for database credentials retrieval (#843)
Co-authored-by: Massimiliano Riva <[email protected]>
1 parent cc4af78 commit 05d8810

File tree

5 files changed

+68
-10
lines changed

5 files changed

+68
-10
lines changed

‎README.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,8 @@ The following table lists the connection properties used with the AWS Advanced P
118118
| `secrets_manager_secret_id` | [Secrets Manager Plugin](./docs/using-the-python-driver/using-plugins/UsingTheAwsSecretsManagerPlugin.md) |
119119
| `secrets_manager_region` | [Secrets Manager Plugin](./docs/using-the-python-driver/using-plugins/UsingTheAwsSecretsManagerPlugin.md) |
120120
| `secrets_manager_endpoint` | [Secrets Manager Plugin](./docs/using-the-python-driver/using-plugins/UsingTheAwsSecretsManagerPlugin.md) |
121+
| `secrets_manager_secret_username` | [Secrets Manager Plugin](./docs/using-the-python-driver/using-plugins/UsingTheAwsSecretsManagerPlugin.md) |
122+
| `secrets_manager_secret_password` | [Secrets Manager Plugin](./docs/using-the-python-driver/using-plugins/UsingTheAwsSecretsManagerPlugin.md) |
121123
| `reader_host_selector_strategy` | [Connection Strategy](./docs/using-the-python-driver/using-plugins/UsingTheReadWriteSplittingPlugin.md#connection-strategies) |
122124
| `db_user` | [Federated Authentication Plugin](./docs/using-the-python-driver/using-plugins/UsingTheFederatedAuthenticationPlugin.md) |
123125
| `idp_username` | [Federated Authentication Plugin](./docs/using-the-python-driver/using-plugins/UsingTheFederatedAuthenticationPlugin.md) |

‎aws_advanced_python_wrapper/aws_secrets_manager_plugin.py

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -188,12 +188,15 @@ def _apply_secret_to_properties(self, properties: Properties):
188188
"""
189189
Updates credentials in provided properties. Other plugins in the plugin chain may change them if needed.
190190
Eventually, credentials will be used to open a new connection in :py:class:`DefaultConnectionPlugin`.
191-
192-
:param properties: Properties to store credentials.
193191
"""
194192
if self._secret:
195-
WrapperProperties.USER.set(properties, self._secret.username)
196-
WrapperProperties.PASSWORD.set(properties, self._secret.password)
193+
username_key = WrapperProperties.SECRETS_MANAGER_SECRET_USERNAME_KEY.get(properties)
194+
username_value = getattr(self._secret, str(username_key))
195+
WrapperProperties.USER.set(properties, username_value)
196+
197+
password_key = WrapperProperties.SECRETS_MANAGER_SECRET_PASSWORD_KEY.get(properties)
198+
password_value = getattr(self._secret, str(password_key))
199+
WrapperProperties.PASSWORD.set(properties, password_value)
197200

198201
def _get_rds_region(self, secret_id: str, props: Properties) -> str:
199202
session = self._session if self._session else boto3.Session()

‎aws_advanced_python_wrapper/utils/properties.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,15 @@ class WrapperProperties:
130130
SECRETS_MANAGER_SECRET_ID = WrapperProperty(
131131
"secrets_manager_secret_id",
132132
"The name or the ARN of the secret to retrieve.")
133+
SECRETS_MANAGER_SECRET_USERNAME_KEY = WrapperProperty(
134+
"secrets_manager_secret_username_key",
135+
"The key of the secret to retrieve, which contains the username.",
136+
"username")
137+
SECRETS_MANAGER_SECRET_PASSWORD_KEY = WrapperProperty(
138+
"secrets_manager_secret_password_key",
139+
"The key of the secret to retrieve, which contains the password.",
140+
"password"
141+
)
133142
SECRETS_MANAGER_REGION = WrapperProperty(
134143
"secrets_manager_region",
135144
"The region of the secret to retrieve.",

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

Lines changed: 28 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,16 +17,18 @@ The following properties are required for the AWS Secrets Manager Connection Plu
1717
> [!IMPORTANT]\
1818
>To use this plugin, you will need to set the following AWS Secrets Manager specific parameters.
1919
20-
| Parameter | Value | Required | Description | Example | Default Value |
21-
|-----------------------------|:------:|:-----------------------------------------------------------:|:---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|:------------------------|---------------|
22-
| `secrets_manager_secret_id` | String | Yes | Set this value to be the secret name or the secret ARN. | `secret_id` | `None` |
23-
| `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` |
24-
| `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` |
20+
| Parameter | Value | Required | Description | Example | Default Value |
21+
|-----------------------------------|:------:|:-----------------------------------------------------------:|:---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|:------------------------|---------------|
22+
| `secrets_manager_secret_id` | String | Yes | Set this value to be the secret name or the secret ARN. | `secret_id` | `None` |
23+
| `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` |
24+
| `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` |
25+
| `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` |
26+
| `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` |
2527

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

2830
## Secret Data
29-
The plugin assumes that the secret contains the following properties: `username` and `password`.
31+
The secret stored in the AWS Secrets Manager should be a JSON object containing the properties `username` and `password`. If the secret contains different key names, you can specify them with the `secrets_manager_secret_username` and `secrets_manager_secret_password` parameters.
3032

3133
### Example
3234

@@ -56,4 +58,24 @@ awsconn = AwsWrapperConnection.connect(
5658
)
5759
```
5860

61+
If you specify `secrets_manager_secret_username` and `secrets_manager_secret_password`, the AWS Advanced Python Driver will parse the secret searching for those specified keys.
62+
```python
63+
awsconn = AwsWrapperConnection.connect(
64+
psycopg.Connection.connect,
65+
host="database.cluster-xyz.us-east-1.rds.amazonaws.com",
66+
dbname="postgres",
67+
secrets_manager_secret_id="secret_name",
68+
secrets_manager_secret_username="custom_username_key",
69+
secrets_manager_secret_password="custom_password_key",
70+
plugins="aws_secrets_manager"
71+
)
72+
```
73+
In this case the secret should have the following format:
74+
```json
75+
{
76+
"custom_username_key": "the database username",
77+
"custom_password_key": "the database password"
78+
}
79+
```
80+
5981
You can find a full example for [PostgreSQL](../../examples/PGSecretsManager.py), and a full example for [MySQL](../../examples/MySQLSecretsManager.py).

‎tests/unit/test_secrets_manager_plugin.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,8 @@ class TestAwsSecretsManagerPlugin(TestCase):
5757
_TEST_ENDPOINT = None
5858
_TEST_USERNAME = "testUser"
5959
_TEST_PASSWORD = "testPassword"
60+
_TEST_USERNAME_KEY = "testUserKey"
61+
_TEST_PASSWORD_KEY = "testPasswordKey"
6062
_TEST_PORT = 5432
6163
_VALID_SECRET_STRING = {'SecretString': f'{{"username":"{_TEST_USERNAME}","password":"{_TEST_PASSWORD}"}}'}
6264
_INVALID_SECRET_STRING = {'SecretString': {"username": "invalid", "password": "invalid"}}
@@ -239,3 +241,23 @@ def test_connection_with_region_parameter_and_arn(self, arn: str, parsed_region:
239241
# The region specified in `secrets_manager_region` should override the region parsed from ARN.
240242
self._mock_session.client.assert_called_with('secretsmanager', region_name=expected_region, endpoint_url=None)
241243
self._mock_client.get_secret_value.assert_called_with(SecretId=arn)
244+
245+
@patch("aws_advanced_python_wrapper.aws_secrets_manager_plugin.AwsSecretsManagerPlugin._secrets_cache", _secrets_cache)
246+
def test_connect_with_different_secret_keys(self):
247+
self._properties["secrets_manager_secret_username_key"] = self._TEST_USERNAME_KEY
248+
self._properties["secrets_manager_secret_password_key"] = self._TEST_PASSWORD_KEY
249+
self._mock_client.get_secret_value.return_value = {
250+
'SecretString': f'{{"{self._TEST_USERNAME_KEY}":"{self._TEST_USERNAME}","{self._TEST_PASSWORD_KEY}":"{self._TEST_PASSWORD}"}}'
251+
}
252+
253+
target_plugin: AwsSecretsManagerPlugin = AwsSecretsManagerPlugin(self._mock_plugin_service,
254+
self._properties,
255+
self._mock_session)
256+
target_plugin.connect(
257+
MagicMock(), MagicMock(), self._TEST_HOST_INFO, self._properties, True, self._mock_func)
258+
259+
assert 1 == len(self._secrets_cache)
260+
self._mock_client.get_secret_value.assert_called_once()
261+
self._mock_func.assert_called_once()
262+
assert self._TEST_USERNAME == self._properties.get("user")
263+
assert self._TEST_PASSWORD == self._properties.get("password")

0 commit comments

Comments
 (0)