This is an automated email from the ASF dual-hosted git repository. potiuk pushed a commit to branch v1-10-test in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 2d0f194d742fcb26373bf91a8d6b2c921d70f38d Author: Kaxil Naik <[email protected]> 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', {
