[AIRFLOW-7105] Unify Secrets Backend method interfaces (#7830)

apache · Mar 23, 2020 · eef87b9 · eef87b9
1 parent 3cc37b1
commit eef87b9
Show file tree

Hide file tree

Showing 8 changed files with 81 additions and 105 deletions.
diff --git a/airflow/providers/amazon/aws/secrets/ssm.py b/airflow/providers/amazon/aws/secrets/ssm.py
@@ -18,12 +18,11 @@
 """
 Objects relating to sourcing connections from AWS SSM Parameter Store
 """
-from typing import List, Optional
+from typing import Optional
 
 import boto3
 from cached_property import cached_property
 
-from airflow.models import Connection
 from airflow.secrets import BaseSecretsBackend
 from airflow.utils.log.logging_mixin import LoggingMixin
 
@@ -62,16 +61,6 @@ def client(self):
  session = boto3.Session(profile_name=self.profile_name)
  return session.client("ssm")
 
- def build_ssm_path(self, conn_id: str):
- """
- Given conn_id, build SSM path.
-
- :param conn_id: connection id
- :type conn_id: str
- """
- param_path = self.connections_prefix + "/" + conn_id
- return param_path
-
  def get_conn_uri(self, conn_id: str) -> Optional[str]:
  """
  Get param value
@@ -80,7 +69,7 @@ def get_conn_uri(self, conn_id: str) -> Optional[str]:
  :type conn_id: str
  """
 
- ssm_path = self.build_ssm_path(conn_id=conn_id)
+ ssm_path = self.build_path(connections_prefix=self.connections_prefix, conn_id=conn_id)
  try:
  response = self.client.get_parameter(
  Name=ssm_path, WithDecryption=False
@@ -93,16 +82,3 @@ def get_conn_uri(self, conn_id: str) -> Optional[str]:
  "Parameter %s not found.", ssm_path
  )
  return None
-
- def get_connections(self, conn_id: str) -> List[Connection]:
- """
- Create connection object.
-
- :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/secrets/secrets_manager.py b/airflow/providers/google/cloud/secrets/secrets_manager.py
@@ -18,15 +18,14 @@
 """
 Objects relating to sourcing connections from GCP Secrets Manager
 """
-from typing import List, Optional
+from typing import 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,
 )
@@ -87,24 +86,14 @@ def client(self) -> SecretManagerServiceClient:
  )
  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)
+ secret_id = self.build_path(connections_prefix=self.connections_prefix, 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)
@@ -117,16 +106,3 @@ def get_conn_uri(self, conn_id: str) -> Optional[str]:
  "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/hashicorp/secrets/vault.py b/airflow/providers/hashicorp/secrets/vault.py
@@ -18,14 +18,13 @@
 """
 Objects relating to sourcing connections from Hashicorp Vault
 """
-from typing import List, Optional
+from typing import Optional
 
 import hvac
 from cached_property import cached_property
 from hvac.exceptions import InvalidPath, VaultError
 
 from airflow import AirflowException
-from airflow.models import Connection
 from airflow.secrets import BaseSecretsBackend
 from airflow.utils.log.logging_mixin import LoggingMixin
 
@@ -144,23 +143,14 @@ def client(self) -> hvac.Client:
  else:
  raise VaultError("Vault Authentication Error!")
 
- def build_path(self, conn_id: str):
- """
- Given conn_id, build path for Vault Secret
-
- :param conn_id: connection id
- :type conn_id: str
- """
- return self.connections_path + "/" + conn_id
-
  def get_conn_uri(self, conn_id: str) -> Optional[str]:
  """
  Get secret value from Vault. Store the secret in the form of URI
 
  :param conn_id: connection id
  :type conn_id: str
  """
- secret_path = self.build_path(conn_id=conn_id)
+ secret_path = self.build_path(connections_prefix=self.connections_path, conn_id=conn_id)
 
  try:
  if self.kv_engine_version == 1:
@@ -176,16 +166,3 @@ def get_conn_uri(self, conn_id: str) -> Optional[str]:
 
  return_data = response["data"] if self.kv_engine_version == 1 else response["data"]["data"]
  return return_data.get("conn_uri")
-
- def get_connections(self, conn_id: str) -> List[Connection]:
- """
- Get connections with a specific ID
-
- :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/secrets/__init__.py b/airflow/secrets/__init__.py
@@ -25,13 +25,13 @@
 __all__ = ['BaseSecretsBackend', 'get_connections']
 
 import json
-from abc import ABC, abstractmethod
 from json import JSONDecodeError
 from typing import List
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.models import Connection
+from airflow.secrets.base_secrets import BaseSecretsBackend
 from airflow.utils.module_loading import import_string
 
 CONFIG_SECTION = "secrets"
@@ -41,24 +41,6 @@
 ]
 
 
-class BaseSecretsBackend(ABC):
- """
- Abstract base class to retrieve secrets given a conn_id and construct a Connection object
- """
-
- def __init__(self, **kwargs):
- pass
-
- @abstractmethod
- def get_connections(self, conn_id) -> List[Connection]:
- """
- Return list of connection objects matching a given ``conn_id``.
-
- :param conn_id: connection id to search for
- :return:
- """
-
-
 def get_connections(conn_id: str) -> List[Connection]:
  """
  Get all connections as an iterable.

diff --git a/airflow/secrets/base_secrets.py b/airflow/secrets/base_secrets.py
@@ -0,0 +1,64 @@
+# 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.
+
+from abc import ABC
+from typing import List, Optional
+
+from airflow.models import Connection
+
+
+class BaseSecretsBackend(ABC):
+ """
+ Abstract base class to retrieve secrets given a conn_id and construct a Connection object
+ """
+
+ def __init__(self, **kwargs):
+ pass
+
+ @staticmethod
+ def build_path(connections_prefix: str, conn_id: str) -> str:
+ """
+ Given conn_id, build path for Secrets Backend
+
+ :param connections_prefix: prefix of the secret to read to get Connections
+ :type connections_prefix: str
+ :param conn_id: connection id
+ :type conn_id: str
+ """
+ return f"{connections_prefix}/{conn_id}"
+
+ def get_conn_uri(self, conn_id: str) -> Optional[str]:
+ """
+ Get conn_uri from Secrets Backend
+
+ :param conn_id: connection id
+ :type conn_id: str
+ """
+ raise NotImplementedError()
+
+ def get_connections(self, conn_id: str) -> List[Connection]:
+ """
+ Return connection object with a given ``conn_id``.
+
+ :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/secrets/environment_variables.py b/airflow/secrets/environment_variables.py
@@ -20,9 +20,8 @@
 """
 
 import os
-from typing import List
+from typing import Optional
 
-from airflow.models import Connection
 from airflow.secrets import BaseSecretsBackend
 
 CONN_ENV_PREFIX = "AIRFLOW_CONN_"
@@ -34,10 +33,6 @@ class EnvironmentVariablesSecretsBackend(BaseSecretsBackend):
  """
 
  # pylint: disable=missing-docstring
- def get_connections(self, conn_id) -> List[Connection]:
+ def get_conn_uri(self, conn_id: str) -> Optional[str]:
  environment_uri = os.environ.get(CONN_ENV_PREFIX + conn_id.upper())
- if environment_uri:
- conn = Connection(conn_id=conn_id, uri=environment_uri)
- return [conn]
- else:
- return []
+ return environment_uri
diff --git a/docs/howto/use-alternative-secrets-backend.rst b/docs/howto/use-alternative-secrets-backend.rst
@@ -187,6 +187,12 @@ Roll your own secrets backend
 A secrets backend is a subclass of :py:class:`airflow.secrets.BaseSecretsBackend`, and just has to implement the
 :py:meth:`~airflow.secrets.BaseSecretsBackend.get_connections` method.
 
+There are two options:
+
+* Option 1: a base implmentation of the :py:meth:`~airflow.secrets.BaseSecretsBackend.get_connections` is provided, you just need to implement the
+ :py:meth:`~airflow.secrets.BaseSecretsBackend.get_conn_uri` method to make it functional.
+* Option 2: simply override the :py:meth:`~airflow.secrets.BaseSecretsBackend.get_connections` method.
+
 Just create your class, and put the fully qualified class name in ``backend`` key in the ``[secrets]``
 section of ``airflow.cfg``. You can you can also pass kwargs to ``__init__`` by supplying json to the
 ``backend_kwargs`` config param. See :ref:`Configuration <secrets_backend_configuration>` for more details,

diff --git a/tests/providers/google/cloud/secrets/test_secrets_manager.py b/tests/providers/google/cloud/secrets/test_secrets_manager.py
@@ -33,7 +33,7 @@
 MODULE_NAME = "airflow.providers.google.cloud.secrets.secrets_manager"
 
 
-class TestGcpSecretsManagerBackend(TestCase):
+class TestCloudSecretsManagerBackend(TestCase):
  @parameterized.expand([
  "airflow/connections",
  "connections",