[AIRFLOW-7104] Add Secret backend for GCP Secrets Manager (#7795)

apache · Mar 23, 2020 · cdf1809 · cdf1809
1 parent d372f23
commit cdf1809
Show file tree

Hide file tree

Showing 14 changed files with 583 additions and 125 deletions.
diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst
@@ -413,6 +413,7 @@ apache.livy http
 dingding http
 discord http
 google amazon,apache.cassandra,cncf.kubernetes,microsoft.azure,microsoft.mssql,mysql,postgres,presto,sftp
+hashicorp google
 microsoft.azure oracle
 microsoft.mssql odbc
 mysql amazon,presto,vertica

diff --git a/airflow/providers/dependencies.json b/airflow/providers/dependencies.json
@@ -38,6 +38,9 @@
  "presto",
  "sftp"
  ],
+ "hashicorp": [
+ "google"
+ ],
  "microsoft.azure": [
  "oracle"
  ],

diff --git a/airflow/providers/google/cloud/hooks/base.py b/airflow/providers/google/cloud/hooks/base.py
@@ -44,13 +44,14 @@
 from airflow import version
 from airflow.exceptions import AirflowException
 from airflow.hooks.base_hook import BaseHook
+from airflow.providers.google.cloud.utils.credentials_provider import (
+ _get_scopes, get_credentials_and_project_id,
+)
 from airflow.utils.process_utils import patch_environ
 
 log = logging.getLogger(__name__)
 
 
-_DEFAULT_SCOPES = ('https://www.googleapis.com/auth/cloud-platform',) # type: Sequence[str]
-
 # Constants used by the mechanism of repeating requests in reaction to exceeding the temporary quota.
 INVALID_KEYS = [
  'DefaultRequestsPerMinutePerProject',
@@ -167,55 +168,21 @@ def _get_credentials_and_project_id(self) -> Tuple[google.auth.credentials.Crede
  if self._cached_credentials is not None:
  return self._cached_credentials, self._cached_project_id
 
- key_path = self._get_field('key_path', None) # type: Optional[str]
- keyfile_dict = self._get_field('keyfile_dict', None) # type: Optional[str]
- if key_path and keyfile_dict:
- raise AirflowException(
- "The `keyfile_dict` and `key_path` fields are mutually exclusive. "
- "Please provide only one value."
- )
- if not key_path and not keyfile_dict:
- self.log.info(
- 'Getting connection using `google.auth.default()` since no key file is defined for hook.'
- )
- credentials, project_id = google.auth.default(scopes=self.scopes)
- elif key_path:
- # Get credentials from a JSON file.
- if key_path.endswith('.json'):
- self.log.debug('Getting connection using JSON key file %s', key_path)
- credentials = (
- google.oauth2.service_account.Credentials.from_service_account_file(
- key_path, scopes=self.scopes)
- )
- project_id = credentials.project_id
- elif key_path.endswith('.p12'):
- raise AirflowException(
- 'Legacy P12 key file are not supported, use a JSON key file.'
- )
- else:
- raise AirflowException('Unrecognised extension for key file.')
- else:
- # Get credentials from JSON data provided in the UI.
- try:
- if not keyfile_dict:
- raise ValueError("The keyfile_dict should be set")
- keyfile_dict_json: Dict[str, str] = json.loads(keyfile_dict)
-
- # Depending on how the JSON was formatted, it may contain
- # escaped newlines. Convert those to actual newlines.
- keyfile_dict_json['private_key'] = keyfile_dict_json['private_key'].replace(
- '\\n', '\n')
-
- credentials = (
- google.oauth2.service_account.Credentials.from_service_account_info(
- keyfile_dict_json, scopes=self.scopes)
- )
- project_id = credentials.project_id
- except json.decoder.JSONDecodeError:
- raise AirflowException('Invalid key JSON.')
-
- if self.delegate_to:
- credentials = credentials.with_subject(self.delegate_to)
+ key_path: Optional[str] = self._get_field('key_path', None)
+ try:
+ keyfile_dict: Optional[str] = self._get_field('keyfile_dict', None)
+ keyfile_dict_json: Optional[Dict[str, str]] = None
+ if keyfile_dict:
+ keyfile_dict_json = json.loads(keyfile_dict)
+ except json.decoder.JSONDecodeError:
+ raise AirflowException('Invalid key JSON.')
+
+ credentials, project_id = get_credentials_and_project_id(
+ key_path=key_path,
+ keyfile_dict=keyfile_dict_json,
+ scopes=self.scopes,
+ delegate_to=self.delegate_to
+ )
 
  overridden_project_id = self._get_field('project')
  if overridden_project_id:
@@ -308,8 +275,7 @@ def scopes(self) -> Sequence[str]:
  """
  scope_value = self._get_field('scope', None) # type: Optional[str]
 
- return [s.strip() for s in scope_value.split(',')] \
- if scope_value else _DEFAULT_SCOPES
+ return _get_scopes(scope_value)
 
  @staticmethod
  def quota_retry(*args, **kwargs) -> Callable:

diff --git a/airflow/providers/google/cloud/secrets/__init__.py b/airflow/providers/google/cloud/secrets/__init__.py
@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
diff --git a/airflow/providers/google/cloud/secrets/secrets_manager.py b/airflow/providers/google/cloud/secrets/secrets_manager.py
@@ -0,0 +1,132 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Objects relating to sourcing connections from GCP Secrets Manager
+"""
+from typing import List, Optional
+
+from cached_property import cached_property
+from google.api_core.exceptions import NotFound
+from google.api_core.gapic_v1.client_info import ClientInfo
+from google.cloud.secretmanager_v1 import SecretManagerServiceClient
+
+from airflow import version
+from airflow.models import Connection
+from airflow.providers.google.cloud.utils.credentials_provider import (
+ _get_scopes, get_credentials_and_project_id,
+)
+from airflow.secrets import BaseSecretsBackend
+from airflow.utils.log.logging_mixin import LoggingMixin
+
+
+class CloudSecretsManagerSecretsBackend(BaseSecretsBackend, LoggingMixin):
+ """
+ Retrieves Connection object from GCP Secrets Manager
+
+ Configurable via ``airflow.cfg`` as follows:
+
+ .. code-block:: ini
+
+ [secrets]
+ backend = airflow.providers.google.cloud.secrets.secrets_manager.CloudSecretsManagerSecretsBackend
+ backend_kwargs = {"connections_prefix": "airflow/connections"}
+
+ For example, if secret id is ``airflow/connections/smtp_default``, this would be accessible
+ if you provide ``{"connections_prefix": "airflow/connections"}`` and request conn_id ``smtp_default``.
+
+ :param connections_prefix: Specifies the prefix of the secret to read to get Connections.
+ :type connections_prefix: str
+ :param gcp_key_path: Path to GCP Credential JSON file;
+ use default credentials in the current environment if not provided.
+ :type gcp_key_path: str
+ :param gcp_scopes: Comma-separated string containing GCP scopes
+ :type gcp_scopes: str
+ """
+ def __init__(
+ self,
+ connections_prefix: str = "airflow/connections",
+ gcp_key_path: Optional[str] = None,
+ gcp_scopes: Optional[str] = None,
+ **kwargs
+ ):
+ self.connections_prefix = connections_prefix.rstrip("/")
+ self.gcp_key_path = gcp_key_path
+ self.gcp_scopes = gcp_scopes
+ self.credentials: Optional[str] = None
+ self.project_id: Optional[str] = None
+ super().__init__(**kwargs)
+
+ @cached_property
+ def client(self) -> SecretManagerServiceClient:
+ """
+ Create an authenticated KMS client
+ """
+ scopes = _get_scopes(self.gcp_scopes)
+ self.credentials, self.project_id = get_credentials_and_project_id(
+ key_path=self.gcp_key_path,
+ scopes=scopes
+ )
+ _client = SecretManagerServiceClient(
+ credentials=self.credentials,
+ client_info=ClientInfo(client_library_version='airflow_v' + version.version)
+ )
+ return _client
+
+ def build_secret_id(self, conn_id: str) -> str:
+ """
+ Given conn_id, build path for Secrets Manager
+
+ :param conn_id: connection id
+ :type conn_id: str
+ """
+ secret_id = f"{self.connections_prefix}/{conn_id}"
+ return secret_id
+
+ def get_conn_uri(self, conn_id: str) -> Optional[str]:
+ """
+ Get secret value from Secrets Manager.
+
+ :param conn_id: connection id
+ :type conn_id: str
+ """
+ secret_id = self.build_secret_id(conn_id=conn_id)
+ # always return the latest version of the secret
+ secret_version = "latest"
+ name = self.client.secret_version_path(self.project_id, secret_id, secret_version)
+ try:
+ response = self.client.access_secret_version(name)
+ value = response.payload.data.decode('UTF-8')
+ return value
+ except NotFound:
+ self.log.error(
+ "GCP API Call Error (NotFound): Secret ID %s not found.", secret_id
+ )
+ return None
+
+ def get_connections(self, conn_id: str) -> List[Connection]:
+ """
+ Create connection object from GCP Secrets Manager
+
+ :param conn_id: connection id
+ :type conn_id: str
+ """
+ conn_uri = self.get_conn_uri(conn_id=conn_id)
+ if not conn_uri:
+ return []
+ conn = Connection(conn_id=conn_id, uri=conn_uri)
+ return [conn]
diff --git a/airflow/providers/google/cloud/utils/credentials_provider.py b/airflow/providers/google/cloud/utils/credentials_provider.py
@@ -20,17 +20,23 @@
 Google Cloud Platform authentication.
 """
 import json
+import logging
 import tempfile
 from contextlib import contextmanager
-from typing import Dict, Optional, Sequence
+from typing import Dict, Optional, Sequence, Tuple
 from urllib.parse import urlencode
 
+import google.auth
+import google.oauth2.service_account
 from google.auth.environment_vars import CREDENTIALS
 
 from airflow.exceptions import AirflowException
 from airflow.utils.process_utils import patch_environ
 
+log = logging.getLogger(__name__)
+
 AIRFLOW_CONN_GOOGLE_CLOUD_DEFAULT = "AIRFLOW_CONN_GOOGLE_CLOUD_DEFAULT"
+_DEFAULT_SCOPES: Sequence[str] = ('https://www.googleapis.com/auth/cloud-platform',)
 
 
 def build_gcp_conn(
@@ -158,3 +164,87 @@ def provide_gcp_conn_and_credentials(
  key_file_path, scopes, project_id
  ):
  yield
+
+
+def get_credentials_and_project_id(
+ key_path: Optional[str] = None,
+ keyfile_dict: Optional[Dict[str, str]] = None,
+ scopes: Optional[Sequence[str]] = None,
+ delegate_to: Optional[str] = None
+) -> Tuple[google.auth.credentials.Credentials, str]:
+ """
+ Returns the Credentials object for Google API and the associated project_id
+
+ Only either `key_path` or `keyfile_dict` should be provided, or an exception will
+ occur. If neither of them are provided, return default credentials for the current environment
+
+ :param key_path: Path to GCP Credential JSON file
+ :type key_path: str
+ :param key_dict: A dict representing GCP Credential as in the Credential JSON file
+ :type key_dict: Dict[str, str]
+ :param scopes: OAuth scopes for the connection
+ :type scopes: Sequence[str]
+ :param delegate_to: The account to impersonate, if any.
+ For this to work, the service account making the request must have
+ domain-wide delegation enabled.
+ :type delegate_to: str
+ :return: Google Auth Credentials
+ :type: google.auth.credentials.Credentials
+ """
+ if key_path and keyfile_dict:
+ raise AirflowException(
+ "The `keyfile_dict` and `key_path` fields are mutually exclusive. "
+ "Please provide only one value."
+ )
+ if not key_path and not keyfile_dict:
+ log.info(
+ 'Getting connection using `google.auth.default()` since no key file is defined for hook.'
+ )
+ credentials, project_id = google.auth.default(scopes=scopes)
+ elif key_path:
+ # Get credentials from a JSON file.
+ if key_path.endswith('.json'):
+ log.debug('Getting connection using JSON key file %s', key_path)
+ credentials = (
+ google.oauth2.service_account.Credentials.from_service_account_file(
+ key_path, scopes=scopes)
+ )
+ project_id = credentials.project_id
+ elif key_path.endswith('.p12'):
+ raise AirflowException(
+ 'Legacy P12 key file are not supported, use a JSON key file.'
+ )
+ else:
+ raise AirflowException('Unrecognised extension for key file.')
+ else:
+ if not keyfile_dict:
+ raise ValueError("The keyfile_dict should be set")
+ # Depending on how the JSON was formatted, it may contain
+ # escaped newlines. Convert those to actual newlines.
+ keyfile_dict['private_key'] = keyfile_dict['private_key'].replace(
+ '\\n', '\n')
+
+ credentials = (
+ google.oauth2.service_account.Credentials.from_service_account_info(
+ keyfile_dict, scopes=scopes)
+ )
+ project_id = credentials.project_id
+
+ if delegate_to:
+ credentials = credentials.with_subject(delegate_to)
+
+ return credentials, project_id
+
+
+def _get_scopes(scopes: Optional[str] = None) -> Sequence[str]:
+ """
+ Parse a comma-separated string containing GCP scopes if `scopes` is provided.
+ Otherwise, default scope will be returned.
+
+ :param scopes: A comma-separated string containing GCP scopes
+ :type scopes: Optional[str]
+ :return: Returns the scope defined in the connection configuration, or the default scope
+ :rtype: Sequence[str]
+ """
+ return [s.strip() for s in scopes.split(',')] \
+ if scopes else _DEFAULT_SCOPES