[airflow] 02/03: Allow passing backend_kwargs to AWS SSM client (#8802)

[ka...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

kaxilnaik pushed a commit to branch v1-10-test
in repository https://gitbox.apache.org/repos/asf/airflow.git

commit ec3038e952cfbaff932ccb9c7c881093c013b077
Author: Kaxil Naik <ka...@gmail.com>
AuthorDate: Sun May 10 12:18:39 2020 +0100

    Allow passing backend_kwargs to AWS SSM client (#8802)
    
    (cherry picked from commit cbebed2b4d0bd1e0984c331c0270e83bf8df8540)
---
 airflow/contrib/secrets/aws_systems_manager.py    | 13 ++++++++++---
 tests/contrib/secrets/test_aws_systems_manager.py | 18 ++++++++++++++++++
 tests/secrets/test_secrets.py                     | 14 ++++++++++++++
 3 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/airflow/contrib/secrets/aws_systems_manager.py b/airflow/contrib/secrets/aws_systems_manager.py
index 862fa96..d85c736 100644
--- a/airflow/contrib/secrets/aws_systems_manager.py
+++ b/airflow/contrib/secrets/aws_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=None,  # type: Optional[str]
         **kwargs
     ):
+        super(SystemsManagerParameterStoreBackend, self).__init__(**kwargs)
         self.connections_prefix = connections_prefix.rstrip("/")
         self.variables_prefix = variables_prefix.rstrip('/')
         self.profile_name = profile_name
-        super(SystemsManagerParameterStoreBackend, self).__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):
         # type: (str) -> Optional[str]
@@ -74,7 +82,6 @@ class SystemsManagerParameterStoreBackend(BaseSecretsBackend, LoggingMixin):
         :type conn_id: str
         :rtype: str
         """
-
         return self._get_secret(self.connections_prefix, conn_id)
 
     def get_variable(self, key):
diff --git a/tests/contrib/secrets/test_aws_systems_manager.py b/tests/contrib/secrets/test_aws_systems_manager.py
index 9801f19..b7b48e1 100644
--- a/tests/contrib/secrets/test_aws_systems_manager.py
+++ b/tests/contrib/secrets/test_aws_systems_manager.py
@@ -22,6 +22,8 @@ from moto import mock_ssm
 
 from airflow.contrib.secrets.aws_systems_manager import SystemsManagerParameterStoreBackend
 from tests.compat import mock
+from airflow.secrets import initialize_secrets_backends
+from tests.test_utils.config import conf_vars
 
 
 class TestSystemsManagerParameterStoreBackend(unittest.TestCase):
@@ -108,3 +110,19 @@ class TestSystemsManagerParameterStoreBackend(unittest.TestCase):
         ssm_backend.client.put_parameter(**param)
 
         self.assertIsNone(ssm_backend.get_variable("test_mysql"))
+
+    @conf_vars({
+        ('secrets', 'backend'): 'airflow.contrib.secrets.aws_systems_manager.'
+                                'SystemsManagerParameterStoreBackend',
+        ('secrets', 'backend_kwargs'): '{"use_ssl": false}'
+    })
+    @mock.patch("airflow.contrib.secrets.aws_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 93de118..1f78c11 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.contrib.secrets.aws_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.contrib.secrets.aws_systems_manager.SystemsManagerParameterStoreBackend",
         ("secrets", "backend_kwargs"): '{"connections_prefix": "/airflow", "profile_name": null}',
     })
     @mock.patch.dict('os.environ', {