Refactoring secrets

Author: leandrodamascena

aws_lambda_powertools/utilities/parameters/secrets.py

@@ -120,13 +120,23 @@ def _get_multiple(self, path: str, **sdk_options) -> Dict[str, str]:
         """
         raise NotImplementedError()
 
-    def _set(
+    def _create_secret(self, name: str, **sdk_options):
+        try:
+            sdk_options["Name"] = name
+            return self.client.create_secret(**sdk_options)
+        except Exception as exc:
+            raise SetSecretError(f"Error setting secret - {str(exc)}") from exc
+
+    def _update_secret(self, name: str, **sdk_options):
+        sdk_options["SecretId"] = name
+        return self.client.put_secret_value(**sdk_options)
+
+    def set(
         self,
         name: str,
         value: Union[str, dict, bytes],
         *,  # force keyword arguments
         client_request_token: Optional[str] = None,
-        version_stages: Optional[list[str]] = None,
         **sdk_options,
     ) -> SetSecretResponse:
         """
@@ -143,8 +153,6 @@ def _set(
             a UUID-type value to ensure uniqueness within the specified secret.
             This value becomes the VersionId of the new version. This field is
             autopopulated if not provided.
-        version_stages: list[str], optional
-            Specifies a list of staging labels that are attached to this version of the secret.
         sdk_options: dict, optional
             Dictionary of options that will be passed to the Secrets Manager update_secret API call
 
@@ -157,8 +165,6 @@ def _set(
             https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/secretsmanager/client/put_secret_value.html
         """
 
-        sdk_options["SecretId"] = name
-
         if isinstance(value, dict):
             value = json.dumps(value)
 
@@ -167,13 +173,13 @@ def _set(
         else:
             sdk_options["SecretString"] = value
 
-        if version_stages:
-            sdk_options["VersionStages"] = version_stages
         if client_request_token:
             sdk_options["ClientRequestToken"] = client_request_token
 
         try:
-            return self.client.put_secret_value(**sdk_options)
+            return self._update_secret(name=name, **sdk_options)
+        except self.client.exceptions.ResourceNotFoundException:
+            return self._create_secret(name=name, **sdk_options)
         except Exception as exc:
             raise SetSecretError(f"Error setting secret - {str(exc)}") from exc
 
@@ -350,10 +356,9 @@ def set_secret(
     if "secrets" not in DEFAULT_PROVIDERS:
         DEFAULT_PROVIDERS["secrets"] = SecretsProvider()
 
-    return DEFAULT_PROVIDERS["secrets"]._set(
+    return DEFAULT_PROVIDERS["secrets"].set(
         name=name,
         value=value,
         client_request_token=client_request_token,
-        version_stages=version_stages,
         **sdk_options,
     )

aws_lambda_powertools/utilities/parameters/types.py

@@ -1,4 +1,6 @@
-from aws_lambda_powertools.shared.types import List, Literal, TypedDict
+from typing import Any, Optional
+
+from aws_lambda_powertools.shared.types import Dict, List, Literal, TypedDict
 
 TransformOptions = Literal["json", "binary", "auto", None]
 
@@ -13,5 +15,6 @@ class SetSecretResponse(TypedDict):
     ARN: str
     Name: str
     VersionId: str
-    VersionStages: List[str]
+    VersionStages: Optional[List[str]]
+    ReplicationStatus: Optional[List[Dict[str, Any]]]
     ResponseMetadata: dict

docs/utilities/parameters.md

@@ -32,10 +32,10 @@ This utility requires additional permissions to work as expected.
 | SSM       | **`get_parameter`**, **`SSMProvider.get`**                             | **`ssm:GetParameter`**                                                               |
 | SSM       | **`get_parameters`**, **`SSMProvider.get_multiple`**                   | **`ssm:GetParametersByPath`**                                                        |
 | SSM       | **`get_parameters_by_name`**, **`SSMProvider.get_parameters_by_name`** | **`ssm:GetParameter`** and **`ssm:GetParameters`**                                   |
-| SSM       | **`set_parameter`**                                                    | **`ssm:PutParameter`**                                                               |
+| SSM       | **`set_parameter`**, **`SSMProvider.set_parameter`**                   | **`ssm:PutParameter`**                                                               |
 | SSM       | If using **`decrypt=True`**                                            | You must add an additional permission **`kms:Decrypt`**                              |
 | Secrets   | **`get_secret`**, **`SecretsProvider.get`**                            | **`secretsmanager:GetSecretValue`**                                                  |
-| Secrets   | **`set_secret`**, **`SecretsProvider.get`**                            | **`secretsmanager:PutSecretValue`** and or **`secretsmanager:CreateSecret`**         |
+| Secrets   | **`set_secret`**, **`SecretsProvider.set`**                            | **`secretsmanager:PutSecretValue`** and or **`secretsmanager:CreateSecret`**         |
 | DynamoDB  | **`DynamoDBProvider.get`**                                             | **`dynamodb:GetItem`**                                                               |
 | DynamoDB  | **`DynamoDBProvider.get_multiple`**                                    | **`dynamodb:Query`**                                                                 |
 | AppConfig | **`get_app_config`**, **`AppConfigProvider.get_app_config`**           | **`appconfig:GetLatestConfiguration`** and **`appconfig:StartConfigurationSession`** |

examples/parameters/src/getting_started_setting_secret.py

@@ -9,7 +9,7 @@
 
 def access_token(client_id: str, client_secret: str, audience: str) -> str:
     # example function that returns a JWT Access Token
-    ...
+    # add your own logic here
     return f"{client_id}.{client_secret}.{audience}"
 
 
@@ -25,6 +25,6 @@ def lambda_handler(event: dict, context: LambdaContext):
         update_secret_version_id = parameters.set_secret(name="/aws-powertools/jwt_token", value=jwt_token)
 
         return {"access_token": "updated", "statusCode": 200, "update_secret_version_id": update_secret_version_id}
-    except parameters.exceptions.SetParameterError as error:
+    except parameters.exceptions.SetSecretError as error:
         logger.exception(error)
         return {"access_token": "updated", "statusCode": 400}

tests/functional/test_utilities_parameters.py

@@ -513,7 +513,7 @@ def test_ssm_provider_get(mock_name, mock_value, mock_version, config):
 
 def test_set_parameter(monkeypatch, mock_name, mock_value):
     """
-    Test get_parameter()
+    Test set_parameter()
     """
 
     class TestProvider(BaseProvider):
@@ -534,7 +534,7 @@ def _get_multiple(self, path: str, **kwargs) -> Dict[str, str]:
     assert value == mock_value
 
 
-def test_ssm_provider_set(mock_name, mock_value, mock_version, config):
+def test_ssm_provider_set_parameter(mock_name, mock_value, mock_version, config):
     """
     Test SSMProvider.set_parameter() with a non-cached value
     """
@@ -564,7 +564,7 @@ def test_ssm_provider_set(mock_name, mock_value, mock_version, config):
         stubber.deactivate()
 
 
-def test_ssm_provider_set_default_config(monkeypatch, mock_name, mock_value, mock_version):
+def test_ssm_provider_set_parameter_default_config(monkeypatch, mock_name, mock_value, mock_version):
     """
     Test SSMProvider._set() without specifying the config
     """
@@ -596,9 +596,9 @@ def test_ssm_provider_set_default_config(monkeypatch, mock_name, mock_value, moc
         stubber.deactivate()
 
 
-def test_ssm_provider_set_with_custom_options(monkeypatch, mock_name, mock_value, mock_version):
+def test_ssm_provider_set_parameter_with_custom_options(monkeypatch, mock_name, mock_value, mock_version):
     """
-    Test SSMProvider._set() without specifying the config
+    Test SSMProvider._set() with custom options
     """
 
     monkeypatch.setenv("AWS_DEFAULT_REGION", "us-east-2")
@@ -638,7 +638,7 @@ def test_ssm_provider_set_with_custom_options(monkeypatch, mock_name, mock_value
         stubber.deactivate()
 
 
-def test_ssm_provider_set_raise_on_failure(mock_name, mock_value, mock_version, config):
+def test_ssm_provider_set_parameter_raise_on_failure(mock_name, mock_value, mock_version, config):
     """
     Test SSMProvider.set_parameter() with failure
     """
@@ -669,6 +669,29 @@ def test_ssm_provider_set_raise_on_failure(mock_name, mock_value, mock_version,
             stubber.deactivate()
 
 
+def test_set_secret(monkeypatch, mock_name, mock_value):
+    """
+    Test set_secret()
+    """
+
+    class TestProvider(BaseProvider):
+        def set(self, name: str, value: Any, *, overwrite: bool = False, **kwargs) -> str:
+            assert name == mock_name
+            return mock_value
+
+        def _get(self, name: str, **kwargs) -> str:
+            raise NotImplementedError()
+
+        def _get_multiple(self, path: str, **kwargs) -> Dict[str, str]:
+            raise NotImplementedError()
+
+    monkeypatch.setitem(parameters.base.DEFAULT_PROVIDERS, "secrets", TestProvider())
+
+    value = parameters.set_secret(name=mock_name, value=mock_value)
+
+    assert value == mock_value
+
+
 def test_ssm_provider_get_with_custom_client(mock_name, mock_value, mock_version, config):
     """
     Test SSMProvider.get() with a non-cached value