feat: allow custom secret keys for database credentials retrieval

README.md

@@ -120,6 +120,8 @@ The following table lists the connection properties used with the AWS Advanced P
 | `secrets_manager_secret_id`                  |           [Secrets Manager Plugin](./docs/using-the-python-driver/using-plugins/UsingTheAwsSecretsManagerPlugin.md)           |
 | `secrets_manager_region`                     |           [Secrets Manager Plugin](./docs/using-the-python-driver/using-plugins/UsingTheAwsSecretsManagerPlugin.md)           |
 | `secrets_manager_endpoint`                   |           [Secrets Manager Plugin](./docs/using-the-python-driver/using-plugins/UsingTheAwsSecretsManagerPlugin.md)           |
+| `secrets_manager_secret_username`            |           [Secrets Manager Plugin](./docs/using-the-python-driver/using-plugins/UsingTheAwsSecretsManagerPlugin.md)           |
+| `secrets_manager_secret_password`            |           [Secrets Manager Plugin](./docs/using-the-python-driver/using-plugins/UsingTheAwsSecretsManagerPlugin.md)           |
 | `reader_host_selector_strategy`              | [Connection Strategy](./docs/using-the-python-driver/using-plugins/UsingTheReadWriteSplittingPlugin.md#connection-strategies) |
 | `db_user`                                    |   [Federated Authentication Plugin](./docs/using-the-python-driver/using-plugins/UsingTheFederatedAuthenticationPlugin.md)    |
 | `idp_username`                               |   [Federated Authentication Plugin](./docs/using-the-python-driver/using-plugins/UsingTheFederatedAuthenticationPlugin.md)    |

aws_advanced_python_wrapper/aws_secrets_manager_plugin.py

@@ -188,12 +188,15 @@ def _apply_secret_to_properties(self, properties: Properties):
         """
         Updates credentials in provided properties. Other plugins in the plugin chain may change them if needed.
         Eventually, credentials will be used to open a new connection in :py:class:`DefaultConnectionPlugin`.
-
-        :param properties: Properties to store credentials.
         """
         if self._secret:
-            WrapperProperties.USER.set(properties, self._secret.username)
-            WrapperProperties.PASSWORD.set(properties, self._secret.password)
+            username_key = WrapperProperties.SECRETS_MANAGER_SECRET_USERNAME_KEY.get(properties)
+            username_value = getattr(self._secret, str(username_key))
+            WrapperProperties.USER.set(properties, username_value)
+
+            password_key = WrapperProperties.SECRETS_MANAGER_SECRET_PASSWORD_KEY.get(properties)
+            password_value = getattr(self._secret, str(password_key))
+            WrapperProperties.PASSWORD.set(properties, password_value)
 
     def _get_rds_region(self, secret_id: str, props: Properties) -> str:
         session = self._session if self._session else boto3.Session()

aws_advanced_python_wrapper/utils/properties.py

@@ -130,6 +130,15 @@ class WrapperProperties:
     SECRETS_MANAGER_SECRET_ID = WrapperProperty(
         "secrets_manager_secret_id",
         "The name or the ARN of the secret to retrieve.")
+    SECRETS_MANAGER_SECRET_USERNAME_KEY = WrapperProperty(
+        "secrets_manager_secret_username_key",
+        "The key of the secret to retrieve, which contains the username.",
+        "username")
+    SECRETS_MANAGER_SECRET_PASSWORD_KEY = WrapperProperty(
+        "secrets_manager_secret_password_key",
+        "The key of the secret to retrieve, which contains the password.",
+        "password"
+    )
     SECRETS_MANAGER_REGION = WrapperProperty(
         "secrets_manager_region",
         "The region of the secret to retrieve.",

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

@@ -17,16 +17,18 @@ The following properties are required for the AWS Secrets Manager Connection Plu
 > [!IMPORTANT]\
 >To use this plugin, you will need to set the following AWS Secrets Manager specific parameters.
 
-| Parameter                   | Value  |                          Required                           | Description                                                                                                                                                                                                                      | Example                 | Default Value |
-|-----------------------------|:------:|:-----------------------------------------------------------:|:---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|:------------------------|---------------|
-| `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`        |
+| Parameter                         | Value  |                          Required                           | Description                                                                                                                                                                                                                      | Example                 | Default Value |
+|-----------------------------------|:------:|:-----------------------------------------------------------:|:---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|:------------------------|---------------|
+| `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_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`    |
 
 *NOTE* A Secret ARN has the following format: `arn:aws:secretsmanager:<Region>:<AccountId>:secret:Secre78tName-6RandomCharacters`
 
 ## Secret Data
-The plugin assumes that the secret contains the following properties: `username` and `password`.
+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.
 
 ### Example
 
@@ -56,4 +58,24 @@ awsconn = AwsWrapperConnection.connect(
 )
 ```
 
+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.
+```python
+awsconn = AwsWrapperConnection.connect(
+        psycopg.Connection.connect,
+        host="database.cluster-xyz.us-east-1.rds.amazonaws.com",
+        dbname="postgres",
+        secrets_manager_secret_id="secret_name",
+        secrets_manager_secret_username="custom_username_key",
+        secrets_manager_secret_password="custom_password_key",
+        plugins="aws_secrets_manager"
+)
+```
+In this case the secret should have the following format:
+```json
+{
+  "custom_username_key": "the database username",
+  "custom_password_key": "the database password"
+}
+```
+
 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

@@ -57,6 +57,8 @@ class TestAwsSecretsManagerPlugin(TestCase):
     _TEST_ENDPOINT = None
     _TEST_USERNAME = "testUser"
     _TEST_PASSWORD = "testPassword"
+    _TEST_USERNAME_KEY = "testUserKey"
+    _TEST_PASSWORD_KEY = "testPasswordKey"
     _TEST_PORT = 5432
     _VALID_SECRET_STRING = {'SecretString': f'{{"username":"{_TEST_USERNAME}","password":"{_TEST_PASSWORD}"}}'}
     _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:
         # The region specified in `secrets_manager_region` should override the region parsed from ARN.
         self._mock_session.client.assert_called_with('secretsmanager', region_name=expected_region, endpoint_url=None)
         self._mock_client.get_secret_value.assert_called_with(SecretId=arn)
+
+    @patch("aws_advanced_python_wrapper.aws_secrets_manager_plugin.AwsSecretsManagerPlugin._secrets_cache", _secrets_cache)
+    def test_connect_with_different_secret_keys(self):
+        self._properties["secrets_manager_secret_username_key"] = self._TEST_USERNAME_KEY
+        self._properties["secrets_manager_secret_password_key"] = self._TEST_PASSWORD_KEY
+        self._mock_client.get_secret_value.return_value = {
+            'SecretString': f'{{"{self._TEST_USERNAME_KEY}":"{self._TEST_USERNAME}","{self._TEST_PASSWORD_KEY}":"{self._TEST_PASSWORD}"}}'
+        }
+
+        target_plugin: AwsSecretsManagerPlugin = AwsSecretsManagerPlugin(self._mock_plugin_service,
+                                                                         self._properties,
+                                                                         self._mock_session)
+        target_plugin.connect(
+            MagicMock(), MagicMock(), self._TEST_HOST_INFO, self._properties, True, self._mock_func)
+
+        assert 1 == len(self._secrets_cache)
+        self._mock_client.get_secret_value.assert_called_once()
+        self._mock_func.assert_called_once()
+        assert self._TEST_USERNAME == self._properties.get("user")
+        assert self._TEST_PASSWORD == self._properties.get("password")