You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ka...@apache.org on 2020/11/07 07:09:57 UTC

[airflow] branch master updated: In AWS Secrets backend, a lookup is optional (#12143)

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

kamilbregula 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 fb6bddb  In AWS Secrets backend, a lookup is optional (#12143)
fb6bddb is described below

commit fb6bddba0c9e3e7ef2610b4fb3f73622e48d7ea0
Author: Faisal <fh...@users.noreply.github.com>
AuthorDate: Sat Nov 7 02:08:48 2020 -0500

    In AWS Secrets backend, a lookup is optional (#12143)
---
 .../amazon/aws/secrets/secrets_manager.py          | 27 ++++++++++--
 .../amazon/aws/secrets/systems_manager.py          | 27 ++++++++++--
 .../amazon/aws/secrets/test_secrets_manager.py     | 42 ++++++++++++++++++
 .../amazon/aws/secrets/test_systems_manager.py     | 51 ++++++++++++++++++++++
 4 files changed, 141 insertions(+), 6 deletions(-)

diff --git a/airflow/providers/amazon/aws/secrets/secrets_manager.py b/airflow/providers/amazon/aws/secrets/secrets_manager.py
index e4786c5..37005e2 100644
--- a/airflow/providers/amazon/aws/secrets/secrets_manager.py
+++ b/airflow/providers/amazon/aws/secrets/secrets_manager.py
@@ -50,10 +50,13 @@ class SecretsManagerBackend(BaseSecretsBackend, LoggingMixin):
     or ``region_name`` to this class and they would be passed on to Boto3 client.
 
     :param connections_prefix: Specifies the prefix of the secret to read to get Connections.
+        If set to None (null), requests for connections will not be sent to AWS Secrets Manager
     :type connections_prefix: str
     :param variables_prefix: Specifies the prefix of the secret to read to get Variables.
+        If set to None (null), requests for variables will not be sent to AWS Secrets Manager
     :type variables_prefix: str
     :param config_prefix: Specifies the prefix of the secret to read to get Variables.
+        If set to None (null), requests for configurations will not be sent to AWS Secrets Manager
     :type config_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
@@ -71,9 +74,18 @@ class SecretsManagerBackend(BaseSecretsBackend, LoggingMixin):
         **kwargs,
     ):
         super().__init__()
-        self.connections_prefix = connections_prefix.rstrip("/")
-        self.variables_prefix = variables_prefix.rstrip('/')
-        self.config_prefix = config_prefix.rstrip('/')
+        if connections_prefix is not None:
+            self.connections_prefix = connections_prefix.rstrip("/")
+        else:
+            self.connections_prefix = connections_prefix
+        if variables_prefix is not None:
+            self.variables_prefix = variables_prefix.rstrip('/')
+        else:
+            self.variables_prefix = variables_prefix
+        if config_prefix is not None:
+            self.config_prefix = config_prefix.rstrip('/')
+        else:
+            self.config_prefix = config_prefix
         self.profile_name = profile_name
         self.sep = sep
         self.kwargs = kwargs
@@ -93,6 +105,9 @@ class SecretsManagerBackend(BaseSecretsBackend, LoggingMixin):
         :param conn_id: connection id
         :type conn_id: str
         """
+        if self.connections_prefix is None:
+            return None
+
         return self._get_secret(self.connections_prefix, conn_id)
 
     def get_variable(self, key: str) -> Optional[str]:
@@ -102,6 +117,9 @@ class SecretsManagerBackend(BaseSecretsBackend, LoggingMixin):
         :param key: Variable Key
         :return: Variable Value
         """
+        if self.variables_prefix is None:
+            return None
+
         return self._get_secret(self.variables_prefix, key)
 
     def get_config(self, key: str) -> Optional[str]:
@@ -111,6 +129,9 @@ class SecretsManagerBackend(BaseSecretsBackend, LoggingMixin):
         :param key: Configuration Option Key
         :return: Configuration Option Value
         """
+        if self.config_prefix is None:
+            return None
+
         return self._get_secret(self.config_prefix, key)
 
     def _get_secret(self, path_prefix: str, secret_id: str) -> Optional[str]:
diff --git a/airflow/providers/amazon/aws/secrets/systems_manager.py b/airflow/providers/amazon/aws/secrets/systems_manager.py
index c7b59ff..cb15f5a 100644
--- a/airflow/providers/amazon/aws/secrets/systems_manager.py
+++ b/airflow/providers/amazon/aws/secrets/systems_manager.py
@@ -43,10 +43,13 @@ class SystemsManagerParameterStoreBackend(BaseSecretsBackend, LoggingMixin):
     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.
+        If set to None (null), requests for connections will not be sent to AWS SSM Parameter Store.
     :type connections_prefix: str
     :param variables_prefix: Specifies the prefix of the secret to read to get Variables.
+        If set to None (null), requests for variables will not be sent to AWS SSM Parameter Store.
     :type variables_prefix: str
     :param config_prefix: Specifies the prefix of the secret to read to get Variables.
+        If set to None (null), requests for configurations will not be sent to AWS SSM Parameter Store.
     :type config_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
@@ -61,9 +64,18 @@ class SystemsManagerParameterStoreBackend(BaseSecretsBackend, LoggingMixin):
         **kwargs,
     ):
         super().__init__()
-        self.connections_prefix = connections_prefix.rstrip("/")
-        self.variables_prefix = variables_prefix.rstrip('/')
-        self.config_prefix = config_prefix.rstrip('/')
+        if connections_prefix is not None:
+            self.connections_prefix = connections_prefix.rstrip("/")
+        else:
+            self.connections_prefix = connections_prefix
+        if variables_prefix is not None:
+            self.variables_prefix = variables_prefix.rstrip('/')
+        else:
+            self.variables_prefix = variables_prefix
+        if config_prefix is not None:
+            self.config_prefix = config_prefix.rstrip('/')
+        else:
+            self.config_prefix = config_prefix
         self.profile_name = profile_name
         self.kwargs = kwargs
 
@@ -80,6 +92,9 @@ class SystemsManagerParameterStoreBackend(BaseSecretsBackend, LoggingMixin):
         :param conn_id: connection id
         :type conn_id: str
         """
+        if self.connections_prefix is None:
+            return None
+
         return self._get_secret(self.connections_prefix, conn_id)
 
     def get_variable(self, key: str) -> Optional[str]:
@@ -89,6 +104,9 @@ class SystemsManagerParameterStoreBackend(BaseSecretsBackend, LoggingMixin):
         :param key: Variable Key
         :return: Variable Value
         """
+        if self.variables_prefix is None:
+            return None
+
         return self._get_secret(self.variables_prefix, key)
 
     def get_config(self, key: str) -> Optional[str]:
@@ -98,6 +116,9 @@ class SystemsManagerParameterStoreBackend(BaseSecretsBackend, LoggingMixin):
         :param key: Configuration Option Key
         :return: Configuration Option Value
         """
+        if self.config_prefix is None:
+            return None
+
         return self._get_secret(self.config_prefix, key)
 
     def _get_secret(self, path_prefix: str, secret_id: str) -> Optional[str]:
diff --git a/tests/providers/amazon/aws/secrets/test_secrets_manager.py b/tests/providers/amazon/aws/secrets/test_secrets_manager.py
index c4dd8e7..bb6e272 100644
--- a/tests/providers/amazon/aws/secrets/test_secrets_manager.py
+++ b/tests/providers/amazon/aws/secrets/test_secrets_manager.py
@@ -83,3 +83,45 @@ class TestSecretsManagerBackend(TestCase):
         secrets_manager_backend.client.put_secret_value(**param)
 
         self.assertIsNone(secrets_manager_backend.get_variable("test_mysql"))
+
+    @mock.patch("airflow.providers.amazon.aws.secrets.secrets_manager.SecretsManagerBackend._get_secret")
+    def test_connection_prefix_none_value(self, mock_get_secret):
+        """
+        Test that if Variable key is not present in AWS Secrets Manager,
+        SecretsManagerBackend.get_conn_uri should return None,
+        SecretsManagerBackend._get_secret should not be called
+        """
+        kwargs = {'connections_prefix': None}
+
+        secrets_manager_backend = SecretsManagerBackend(**kwargs)
+
+        self.assertIsNone(secrets_manager_backend.get_conn_uri("test_mysql"))
+        mock_get_secret.assert_not_called()
+
+    @mock.patch("airflow.providers.amazon.aws.secrets.secrets_manager.SecretsManagerBackend._get_secret")
+    def test_variable_prefix_none_value(self, mock_get_secret):
+        """
+        Test that if Variable key is not present in AWS Secrets Manager,
+        SecretsManagerBackend.get_variables should return None,
+        SecretsManagerBackend._get_secret should not be called
+        """
+        kwargs = {'variables_prefix': None}
+
+        secrets_manager_backend = SecretsManagerBackend(**kwargs)
+
+        self.assertIsNone(secrets_manager_backend.get_variable("hello"))
+        mock_get_secret.assert_not_called()
+
+    @mock.patch("airflow.providers.amazon.aws.secrets.secrets_manager.SecretsManagerBackend._get_secret")
+    def test_config_prefix_none_value(self, mock_get_secret):
+        """
+        Test that if Variable key is not present in AWS Secrets Manager,
+        SecretsManagerBackend.get_config should return None,
+        SecretsManagerBackend._get_secret should not be called
+        """
+        kwargs = {'config_prefix': None}
+
+        secrets_manager_backend = SecretsManagerBackend(**kwargs)
+
+        self.assertIsNone(secrets_manager_backend.get_config("config"))
+        mock_get_secret.assert_not_called()
diff --git a/tests/providers/amazon/aws/secrets/test_systems_manager.py b/tests/providers/amazon/aws/secrets/test_systems_manager.py
index cd48533..01a4669 100644
--- a/tests/providers/amazon/aws/secrets/test_systems_manager.py
+++ b/tests/providers/amazon/aws/secrets/test_systems_manager.py
@@ -131,3 +131,54 @@ class TestSsmSecrets(TestCase):
 
         systems_manager.client
         mock_ssm_client.assert_called_once_with('ssm', use_ssl=False)
+
+    @mock.patch(
+        "airflow.providers.amazon.aws.secrets.systems_manager."
+        "SystemsManagerParameterStoreBackend._get_secret"
+    )
+    def test_connection_prefix_none_value(self, mock_get_secret):
+        """
+        Test that if Variable key is not present in SSM,
+        SystemsManagerParameterStoreBackend.get_conn_uri should return None,
+        SystemsManagerParameterStoreBackend._get_secret should not be called
+        """
+        kwargs = {'connections_prefix': None}
+
+        ssm_backend = SystemsManagerParameterStoreBackend(**kwargs)
+
+        self.assertIsNone(ssm_backend.get_conn_uri("test_mysql"))
+        mock_get_secret.assert_not_called()
+
+    @mock.patch(
+        "airflow.providers.amazon.aws.secrets.systems_manager."
+        "SystemsManagerParameterStoreBackend._get_secret"
+    )
+    def test_variable_prefix_none_value(self, mock_get_secret):
+        """
+        Test that if Variable key is not present in SSM,
+        SystemsManagerParameterStoreBackend.get_variables should return None,
+        SystemsManagerParameterStoreBackend._get_secret should not be called
+        """
+        kwargs = {'variables_prefix': None}
+
+        ssm_backend = SystemsManagerParameterStoreBackend(**kwargs)
+
+        self.assertIsNone(ssm_backend.get_variable("hello"))
+        mock_get_secret.assert_not_called()
+
+    @mock.patch(
+        "airflow.providers.amazon.aws.secrets.systems_manager."
+        "SystemsManagerParameterStoreBackend._get_secret"
+    )
+    def test_config_prefix_none_value(self, mock_get_secret):
+        """
+        Test that if Variable key is not present in SSM,
+        SystemsManagerParameterStoreBackend.get_config should return None,
+        SystemsManagerParameterStoreBackend._get_secret should not be called
+        """
+        kwargs = {'config_prefix': None}
+
+        ssm_backend = SystemsManagerParameterStoreBackend(**kwargs)
+
+        self.assertIsNone(ssm_backend.get_config("config"))
+        mock_get_secret.assert_not_called()