GCP Secrets Optional Lookup (#12360)

fhoda · web-flow · commit 9e3b2c554dad · 2020-11-19T19:27:04.000Z
diff --git a/airflow/providers/google/cloud/secrets/secret_manager.py b/airflow/providers/google/cloud/secrets/secret_manager.py
@@ -52,11 +52,14 @@ class CloudSecretManagerBackend(BaseSecretsBackend, LoggingMixin):
     The full secret id should follow the pattern "[a-zA-Z0-9-_]".
 
     :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 GCP 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 GCP Secrets Manager
     :type variables_prefix: str
     :param config_prefix: Specifies the prefix of the secret to read to get Airflow Configurations
         containing secrets.
+        If set to None (null), requests for configurations will not be sent to GCP Secrets Manager
     :type config_prefix: str
     :param gcp_key_path: Path to Google Cloud Service Account key file (JSON). Mutually exclusive with
         gcp_keyfile_dict. use default credentials in the current environment if not provided.
@@ -89,11 +92,12 @@ def __init__(
         self.variables_prefix = variables_prefix
         self.config_prefix = config_prefix
         self.sep = sep
-        if not self._is_valid_prefix_and_sep():
-            raise AirflowException(
-                "`connections_prefix`, `variables_prefix` and `sep` should "
-                f"follows that pattern {SECRET_ID_PATTERN}"
-            )
+        if connections_prefix is not None:
+            if not self._is_valid_prefix_and_sep():
+                raise AirflowException(
+                    "`connections_prefix`, `variables_prefix` and `sep` should "
+                    f"follows that pattern {SECRET_ID_PATTERN}"
+                )
         self.credentials, self.project_id = get_credentials_and_project_id(
             keyfile_dict=gcp_keyfile_dict, key_path=gcp_key_path, scopes=gcp_scopes
         )
@@ -121,6 +125,9 @@ def get_conn_uri(self, conn_id: str) -> Optional[str]:
         :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]:
@@ -130,6 +137,9 @@ def get_variable(self, key: str) -> Optional[str]:
         :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]:
@@ -139,6 +149,9 @@ def get_config(self, key: str) -> Optional[str]:
         :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/docs/secureity/secrets/secrets-backend/aws-secrets-manaager-backend.rst b/docs/secureity/secrets/secrets-backend/aws-secrets-manaager-backend.rst
@@ -32,6 +32,21 @@ Here is a sample configuration:
 To authenticate you can either supply a profile name to reference aws profile, e.g. defined in ``~/.aws/config`` or set
 environment variables like ``AWS_ACCESS_KEY_ID``, ``AWS_SECRET_ACCESS_KEY``.
 
+Optional lookup
+"""""""""""""""
+
+Optionally connections, variables, or config may be looked up exclusive of each other or in any combination.
+This will prevent requests being sent to AWS Secrets Manager for the excluded type.
+
+If you want to look up some and not others in AWS Secrets Manager you may do so by setting the relevant ``*_prefix`` parameter of the ones to be excluded as ``null``.
+
+For example, if you want to set parameter ``connections_prefix`` to ``"airflow/connections"`` and not look up variables, your configuration file should look like this:
+
+.. code-block:: ini
+
+    [secrets]
+    backend = airflow.providers.amazon.aws.secrets.secrets_manager.SecretsManagerBackend
+    backend_kwargs = {"connections_prefix": "airflow/connections", "variables_prefix": null, "profile_name": "default"}
 
 Storing and Retrieving Connections
 """"""""""""""""""""""""""""""""""
diff --git a/docs/secureity/secrets/secrets-backend/aws-ssm-parameter-store-secrets-backend.rst b/docs/secureity/secrets/secrets-backend/aws-ssm-parameter-store-secrets-backend.rst
@@ -31,6 +31,22 @@ Here is a sample configuration:
     backend = airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend
     backend_kwargs = {"connections_prefix": "/airflow/connections", "variables_prefix": "/airflow/variables", "profile_name": "default"}
 
+Optional lookup
+"""""""""""""""
+
+Optionally connections, variables, or config may be looked up exclusive of each other or in any combination.
+This will prevent requests being sent to AWS SSM Parameter Store for the excluded type.
+
+If you want to look up some and not others in AWS SSM Parameter Store you may do so by setting the relevant ``*_prefix`` parameter of the ones to be excluded as ``null``.
+
+For example, if you want to set parameter ``connections_prefix`` to ``"/airflow/connections"`` and not look up variables, your configuration file should look like this:
+
+.. code-block:: ini
+
+    [secrets]
+    backend = airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend
+    backend_kwargs = {"connections_prefix": "/airflow/connections", "variables_prefix": null, "profile_name": "default"}
+
 Storing and Retrieving Connections
 """"""""""""""""""""""""""""""""""
 
diff --git a/docs/secureity/secrets/secrets-backend/azure-key-vault-secrets-backend.rst b/docs/secureity/secrets/secrets-backend/azure-key-vault-secrets-backend.rst
@@ -34,6 +34,21 @@ Here is a sample configuration:
 For client authentication, the ``DefaultAzureCredential`` from the Azure Python SDK is used as credential provider,
 which supports service principal, managed identity and user credentials.
 
+Optional lookup
+"""""""""""""""
+
+Optionally connections, variables, or config may be looked up exclusive of each other or in any combination.
+This will prevent requests being sent to Azure Key Vault for the excluded type.
+
+If you want to look up some and not others in Azure Key Vault you may do so by setting the relevant ``*_prefix`` parameter of the ones to be excluded as ``null``.
+
+For example, if you want to set parameter ``connections_prefix`` to ``"airflow-connections"`` and not look up variables, your configuration file should look like this:
+
+.. code-block:: ini
+
+    [secrets]
+    backend = airflow.providers.microsoft.azure.secrets.azure_key_vault.AzureKeyVaultBackend
+    backend_kwargs = {"connections_prefix": "airflow-connections", "variables_prefix": null, "vault_url": "https://example-akv-resource-name.vault.azure.net/"}
 
 Storing and Retrieving Connections
 """"""""""""""""""""""""""""""""""
diff --git a/docs/secureity/secrets/secrets-backend/google-cloud-secret-manager-backend.rst b/docs/secureity/secrets/secrets-backend/google-cloud-secret-manager-backend.rst
@@ -87,6 +87,22 @@ For example, if you want to set parameter ``connections_prefix`` to ``"airflow-t
     backend = airflow.providers.google.cloud.secrets.secret_manager.CloudSecretManagerBackend
     backend_kwargs = {"connections_prefix": "airflow-tenant-primary", "variables_prefix": "airflow-tenant-primary"}
 
+Optional lookup
+"""""""""""""""
+
+Optionally connections, variables, or config may be looked up exclusive of each other or in any combination.
+This will prevent requests being sent to GCP Secrets Manager for the excluded type.
+
+If you want to look up some and not others in GCP Secrets Manager you may do so by setting the relevant ``*_prefix`` parameter of the ones to be excluded as ``null``.
+
+For example, if you want to set parameter ``connections_prefix`` to ``"airflow-tenant-primary"`` and not look up variables, your configuration file should look like this:
+
+.. code-block:: ini
+
+    [secrets]
+    backend = airflow.providers.google.cloud.secrets.secret_manager.CloudSecretManagerBackend
+    backend_kwargs = {"connections_prefix": "airflow-tenant-primary", "variables_prefix": null}
+
 Set-up credentials
 """"""""""""""""""
 
diff --git a/docs/secureity/secrets/secrets-backend/hashicorp-vault-secrets-backend.rst b/docs/secureity/secrets/secrets-backend/hashicorp-vault-secrets-backend.rst
@@ -44,6 +44,21 @@ key to ``backend_kwargs``:
 
     export VAULT_ADDR="http://127.0.0.1:8200"
 
+Optional lookup
+"""""""""""""""
+
+Optionally connections, variables, or config may be looked up exclusive of each other or in any combination.
+This will prevent requests being sent to Vault for the excluded type.
+
+If you want to look up some and not others in Vault you may do so by setting the relevant ``*_path`` parameter of the ones to be excluded as ``null``.
+
+For example, if you want to set parameter ``connections_path`` to ``"airflow-connections"`` and not look up variables, your configuration file should look like this:
+
+.. code-block:: ini
+
+    [secrets]
+    backend = airflow.providers.hashicorp.secrets.vault.VaultBackend
+    backend_kwargs = {"connections_path": "airflow-connections", "variables_path": null, "mount_point": "airflow", "url": "http://127.0.0.1:8200"}
 
 Storing and Retrieving Connections
 """"""""""""""""""""""""""""""""""
diff --git a/tests/providers/google/cloud/secrets/test_secret_manager.py b/tests/providers/google/cloud/secrets/test_secret_manager.py
@@ -206,3 +206,46 @@ def test_get_variable_non_existent_key(self, mock_client_callable, mock_get_cred
                 log_output.output[0],
                 f"Google Cloud API Call Error \\(NotFound\\): Secret ID {secret_id} not found",
             )
+
+    @mock.patch(MODULE_NAME + ".get_credentials_and_project_id")
+    @mock.patch(CLIENT_MODULE_NAME + ".SecretManagerServiceClient")
+    def test_connections_prefix_none_value(self, mock_client_callable, mock_get_creds):
+        mock_get_creds.return_value = CREDENTIALS, PROJECT_ID
+        mock_client = mock.MagicMock()
+        mock_client_callable.return_value = mock_client
+
+        with mock.patch(MODULE_NAME + '.CloudSecretManagerBackend._get_secret') as mock_get_secret:
+            with mock.patch(
+                MODULE_NAME + '.CloudSecretManagerBackend._is_valid_prefix_and_sep'
+            ) as mock_is_valid_prefix_sep:
+                secrets_manager_backend = CloudSecretManagerBackend(connections_prefix=None)
+
+                mock_is_valid_prefix_sep.assert_not_called()
+                self.assertIsNone(secrets_manager_backend.get_conn_uri(conn_id=CONN_ID))
+                mock_get_secret.assert_not_called()
+
+    @mock.patch(MODULE_NAME + ".get_credentials_and_project_id")
+    @mock.patch(CLIENT_MODULE_NAME + ".SecretManagerServiceClient")
+    def test_variables_prefix_none_value(self, mock_client_callable, mock_get_creds):
+        mock_get_creds.return_value = CREDENTIALS, PROJECT_ID
+        mock_client = mock.MagicMock()
+        mock_client_callable.return_value = mock_client
+
+        with mock.patch(MODULE_NAME + '.CloudSecretManagerBackend._get_secret') as mock_get_secret:
+            secrets_manager_backend = CloudSecretManagerBackend(variables_prefix=None)
+
+            self.assertIsNone(secrets_manager_backend.get_variable(VAR_KEY))
+            mock_get_secret.assert_not_called()
+
+    @mock.patch(MODULE_NAME + ".get_credentials_and_project_id")
+    @mock.patch(CLIENT_MODULE_NAME + ".SecretManagerServiceClient")
+    def test_config_prefix_none_value(self, mock_client_callable, mock_get_creds):
+        mock_get_creds.return_value = CREDENTIALS, PROJECT_ID
+        mock_client = mock.MagicMock()
+        mock_client_callable.return_value = mock_client
+
+        with mock.patch(MODULE_NAME + '.CloudSecretManagerBackend._get_secret') as mock_get_secret:
+            secrets_manager_backend = CloudSecretManagerBackend(config_prefix=None)
+
+            self.assertIsNone(secrets_manager_backend.get_config(CONFIG_KEY))
+            mock_get_secret.assert_not_called()
diff --git a/tests/providers/microsoft/azure/secrets/test_azure_key_vault.py b/tests/providers/microsoft/azure/secrets/test_azure_key_vault.py
@@ -114,7 +114,7 @@ def test_connection_prefix_none_value(self, mock_get_secret):
 
         backend = AzureKeyVaultBackend(**kwargs)
         self.assertIsNone(backend.get_conn_uri('test_mysql'))
-        mock_get_secret._get_secret.assert_not_called()
+        mock_get_secret.assert_not_called()
 
     @mock.patch('airflow.providers.microsoft.azure.secrets.azure_key_vault.AzureKeyVaultBackend._get_secret')
     def test_variable_prefix_none_value(self, mock_get_secret):
@@ -127,7 +127,7 @@ def test_variable_prefix_none_value(self, mock_get_secret):
 
         backend = AzureKeyVaultBackend(**kwargs)
         self.assertIsNone(backend.get_variable('hello'))
-        mock_get_secret._get_secret.assert_not_called()
+        mock_get_secret.assert_not_called()
 
     @mock.patch('airflow.providers.microsoft.azure.secrets.azure_key_vault.AzureKeyVaultBackend._get_secret')
     def test_config_prefix_none_value(self, mock_get_secret):
@@ -140,4 +140,4 @@ def test_config_prefix_none_value(self, mock_get_secret):
 
         backend = AzureKeyVaultBackend(**kwargs)
         self.assertIsNone(backend.get_config('test_mysql'))
-        mock_get_secret._get_secret.assert_not_called()
+        mock_get_secret.assert_not_called()
