This is an automated email from the ASF dual-hosted git repository.
kaxilnaik pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/master by this push:
new cbebed2 Allow passing backend_kwargs to AWS SSM client (#8802)
cbebed2 is described below
commit cbebed2b4d0bd1e0984c331c0270e83bf8df8540
Author: Kaxil Naik <[email protected]>
AuthorDate: Sun May 10 12:18:39 2020 +0100
Allow passing backend_kwargs to AWS SSM client (#8802)
---
.../providers/amazon/aws/secrets/systems_manager.py | 13 ++++++++++---
.../amazon/aws/secrets/test_systems_manager.py | 18 ++++++++++++++++++
tests/secrets/test_secrets.py | 14 ++++++++++++++
3 files changed, 42 insertions(+), 3 deletions(-)
diff --git a/airflow/providers/amazon/aws/secrets/systems_manager.py
b/airflow/providers/amazon/aws/secrets/systems_manager.py
index e2dcb78..4d2c384 100644
--- a/airflow/providers/amazon/aws/secrets/systems_manager.py
+++ b/airflow/providers/amazon/aws/secrets/systems_manager.py
@@ -43,6 +43,13 @@ class
SystemsManagerParameterStoreBackend(BaseSecretsBackend, LoggingMixin):
if you provide ``{"connections_prefix": "/airflow/connections"}`` and
request conn_id ``smtp_default``.
And if ssm path is ``/airflow/variables/hello``, this would be accessible
if you provide ``{"variables_prefix": "/airflow/variables"}`` and request
conn_id ``hello``.
+
+ :param connections_prefix: Specifies the prefix of the secret to read to
get Connections.
+ :type connections_prefix: str
+ :param variables_prefix: Specifies the prefix of the secret to read to get
Variables.
+ :type variables_prefix: str
+ :param profile_name: The name of a profile to use. If not given, then the
default profile is used.
+ :type profile_name: str
"""
def __init__(
@@ -52,10 +59,11 @@ class
SystemsManagerParameterStoreBackend(BaseSecretsBackend, LoggingMixin):
profile_name: Optional[str] = None,
**kwargs
):
+ super().__init__(**kwargs)
self.connections_prefix = connections_prefix.rstrip("/")
self.variables_prefix = variables_prefix.rstrip('/')
self.profile_name = profile_name
- super().__init__(**kwargs)
+ self.kwargs = kwargs
@cached_property
def client(self):
@@ -63,7 +71,7 @@ class SystemsManagerParameterStoreBackend(BaseSecretsBackend,
LoggingMixin):
Create a SSM client
"""
session = boto3.Session(profile_name=self.profile_name)
- return session.client("ssm")
+ return session.client("ssm", **self.kwargs)
def get_conn_uri(self, conn_id: str) -> Optional[str]:
"""
@@ -72,7 +80,6 @@ class SystemsManagerParameterStoreBackend(BaseSecretsBackend,
LoggingMixin):
:param conn_id: connection id
:type conn_id: str
"""
-
return self._get_secret(self.connections_prefix, conn_id)
def get_variable(self, key: str) -> Optional[str]:
diff --git a/tests/providers/amazon/aws/secrets/test_systems_manager.py
b/tests/providers/amazon/aws/secrets/test_systems_manager.py
index a953128..4adea15 100644
--- a/tests/providers/amazon/aws/secrets/test_systems_manager.py
+++ b/tests/providers/amazon/aws/secrets/test_systems_manager.py
@@ -20,6 +20,8 @@ from unittest import TestCase, mock
from moto import mock_ssm
from airflow.providers.amazon.aws.secrets.systems_manager import
SystemsManagerParameterStoreBackend
+from airflow.secrets import initialize_secrets_backends
+from tests.test_utils.config import conf_vars
class TestSsmSecrets(TestCase):
@@ -94,3 +96,19 @@ class TestSsmSecrets(TestCase):
ssm_backend.client.put_parameter(**param)
self.assertIsNone(ssm_backend.get_variable("test_mysql"))
+
+ @conf_vars({
+ ('secrets', 'backend'):
'airflow.providers.amazon.aws.secrets.systems_manager.'
+ 'SystemsManagerParameterStoreBackend',
+ ('secrets', 'backend_kwargs'): '{"use_ssl": false}'
+ })
+
@mock.patch("airflow.providers.amazon.aws.secrets.systems_manager.boto3.Session.client")
+ def test_passing_client_kwargs(self, mock_ssm_client):
+ backends = initialize_secrets_backends()
+ systems_manager = [
+ backend for backend in backends
+ if backend.__class__.__name__ ==
'SystemsManagerParameterStoreBackend'
+ ][0]
+
+ systems_manager.client
+ mock_ssm_client.assert_called_once_with('ssm', use_ssl=False)
diff --git a/tests/secrets/test_secrets.py b/tests/secrets/test_secrets.py
index 709aa8c..7bb1ed3 100644
--- a/tests/secrets/test_secrets.py
+++ b/tests/secrets/test_secrets.py
@@ -57,6 +57,20 @@ class TestConnectionsFromSecrets(unittest.TestCase):
@conf_vars({
("secrets", "backend"):
"airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend",
+ ("secrets", "backend_kwargs"): '{"use_ssl": false}',
+ })
+ def test_backends_kwargs(self):
+ backends = initialize_secrets_backends()
+ systems_manager = [
+ backend for backend in backends
+ if backend.__class__.__name__ ==
'SystemsManagerParameterStoreBackend'
+ ][0]
+
+ self.assertEqual(systems_manager.kwargs, {'use_ssl': False})
+
+ @conf_vars({
+ ("secrets", "backend"):
+
"airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend",
("secrets", "backend_kwargs"): '{"connections_prefix": "/airflow",
"profile_name": null}',
})
@mock.patch.dict('os.environ', {
