From c9434db67be6f05e57c5d5ef21da49944741c341 Mon Sep 17 00:00:00 2001 From: Abdel Jaidi Date: Tue, 5 Oct 2021 18:40:49 +0100 Subject: [PATCH] Minor - Adding Secrets Manager Endpoint --- awswrangler/_config.py | 11 +++++++++++ awswrangler/_utils.py | 2 ++ tests/test_config.py | 4 ++++ 3 files changed, 17 insertions(+) diff --git a/awswrangler/_config.py b/awswrangler/_config.py index eca691770..54b77e653 100644 --- a/awswrangler/_config.py +++ b/awswrangler/_config.py @@ -44,6 +44,7 @@ class _ConfigArg(NamedTuple): "kms_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True), "emr_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True), "dynamodb_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True), + "secretsmanager_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True), # Botocore config "botocore_config": _ConfigArg(dtype=botocore.config.Config, nullable=True), } @@ -63,6 +64,7 @@ def __init__(self) -> None: self.kms_endpoint_url = None self.emr_endpoint_url = None self.dynamodb_endpoint_url = None + self.secretsmanager_endpoint_url = None self.botocore_config = None for name in _CONFIG_ARGS: self._load_config(name=name) @@ -363,6 +365,15 @@ def dynamodb_endpoint_url(self) -> Optional[str]: def dynamodb_endpoint_url(self, value: Optional[str]) -> None: self._set_config_value(key="dynamodb_endpoint_url", value=value) + @property + def secretsmanager_endpoint_url(self) -> Optional[str]: + """Property secretsmanager_endpoint_url.""" + return cast(Optional[str], self["secretsmanager_endpoint_url"]) + + @secretsmanager_endpoint_url.setter + def secretsmanager_endpoint_url(self, value: Optional[str]) -> None: + self._set_config_value(key="secretsmanager_endpoint_url", value=value) + @property def botocore_config(self) -> botocore.config.Config: """Property botocore_config.""" diff --git a/awswrangler/_utils.py b/awswrangler/_utils.py index 961f2a480..bf04f642c 100644 --- a/awswrangler/_utils.py +++ b/awswrangler/_utils.py @@ -95,6 +95,8 @@ def _get_endpoint_url(service_name: str) -> Optional[str]: endpoint_url = _config.config.emr_endpoint_url elif service_name == "dynamodb" and _config.config.dynamodb_endpoint_url is not None: endpoint_url = _config.config.dynamodb_endpoint_url + elif service_name == "secretsmanager" and _config.config.secretsmanager_endpoint_url is not None: + endpoint_url = _config.config.secretsmanager_endpoint_url return endpoint_url diff --git a/tests/test_config.py b/tests/test_config.py index c796a4b8b..86ea3eccc 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -29,6 +29,8 @@ def wrapper(self, **kwarg): assert url == wr.config.s3_endpoint_url elif name == "glue": assert url == wr.config.glue_endpoint_url + elif name == "secretsmanager": + assert url == wr.config.secretsmanager_endpoint_url return original(self, **kwarg) with patch("botocore.client.ClientCreator.create_client", new=wrapper): @@ -112,11 +114,13 @@ def test_basics(path, glue_database, glue_table, workgroup0, workgroup1): wr.config.s3_endpoint_url = f"https://s3.{region}.amazonaws.com" wr.config.athena_endpoint_url = f"https://athena.{region}.amazonaws.com" wr.config.glue_endpoint_url = f"https://glue.{region}.amazonaws.com" + wr.config.secretsmanager_endpoint_url = f"https://secretsmanager.{region}.amazonaws.com" _urls_test(glue_database) os.environ["WR_STS_ENDPOINT_URL"] = f"https://sts.{region}.amazonaws.com" os.environ["WR_S3_ENDPOINT_URL"] = f"https://s3.{region}.amazonaws.com" os.environ["WR_ATHENA_ENDPOINT_URL"] = f"https://athena.{region}.amazonaws.com" os.environ["WR_GLUE_ENDPOINT_URL"] = f"https://glue.{region}.amazonaws.com" + os.environ["WR_SECRETSMANAGER_ENDPOINT_URL"] = f"https://secretsmanager.{region}.amazonaws.com" wr.config.reset() _urls_test(glue_database)