Add optional features in providers. (#21074)

Some features in providers can be optional, depending on the presence of some libraries. Since Providers Manager tries to import the right classes that are exposed via providers it should not - in this case - log warning message for those optional features. Previously, all ImportErrors were turned into debug log but now we only turn them in debug log when creator of the provider deliberately raised an AirflowOptionalProviderFeatureException. Instructions on how to raise such exception in the way to keep backwards compatibility were updated in proider's documentation. Fixes: #20709
apache · Jan 27, 2022 · cb73053 · cb73053
1 parent 6405d8f
commit cb73053
Show file tree

Hide file tree

Showing 32 changed files with 183 additions and 74 deletions.
diff --git a/airflow/exceptions.py b/airflow/exceptions.py
@@ -98,6 +98,10 @@ class AirflowFailException(AirflowException):
  """Raise when the task should be failed without retrying."""
 
 
+class AirflowOptionalProviderFeatureException(AirflowException):
+ """Raise by providers when imports are missing for optional provider features."""
+
+
 class UnmappableXComTypePushed(AirflowException):
  """Raise when an unmappable type is pushed as a mapped downstream's dependency."""
 

diff --git a/airflow/hooks/base.py b/airflow/hooks/base.py
@@ -97,6 +97,14 @@ def get_conn(self) -> Any:
  """Returns connection for the hook."""
  raise NotImplementedError()
 
+ @classmethod
+ def get_connection_form_widgets(cls) -> Dict[str, Any]:
+ ...
+
+ @classmethod
+ def get_ui_field_behaviour(cls) -> Dict[str, Any]:
+ ...
+
 
 class DiscoverableHook(Protocol):
  """
@@ -159,7 +167,7 @@ def get_connection_form_widgets() -> Dict[str, Any]:
  ...
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """
  Returns dictionary describing customizations to implement in javascript handling the
  connection form. Should be compliant with airflow/customized_form_field_behaviours.schema.json'

diff --git a/airflow/providers/apache/spark/hooks/spark_submit.py b/airflow/providers/apache/spark/hooks/spark_submit.py
@@ -86,7 +86,7 @@ class SparkSubmitHook(BaseHook, LoggingMixin):
  hook_name = 'Spark'
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """Returns custom field behaviour"""
  return {
  "hidden_fields": ['schema', 'login', 'password'],

diff --git a/airflow/providers/asana/hooks/asana.py b/airflow/providers/asana/hooks/asana.py
@@ -62,7 +62,7 @@ def get_connection_form_widgets() -> Dict[str, Any]:
  }
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """Returns custom field behaviour"""
  return {
  "hidden_fields": ["port", "host", "login", "schema"],

diff --git a/airflow/providers/cloudant/hooks/cloudant.py b/airflow/providers/cloudant/hooks/cloudant.py
@@ -16,7 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Hook for Cloudant"""
-from typing import Dict
+from typing import Any, Dict
 
 from cloudant import cloudant
 
@@ -39,7 +39,7 @@ class CloudantHook(BaseHook):
  hook_name = 'Cloudant'
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """Returns custom field behaviour"""
  return {
  "hidden_fields": ['port', 'extra'],

diff --git a/airflow/providers/cncf/kubernetes/hooks/kubernetes.py b/airflow/providers/cncf/kubernetes/hooks/kubernetes.py
@@ -92,7 +92,7 @@ def get_connection_form_widgets() -> Dict[str, Any]:
  }
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """Returns custom field behaviour"""
  return {
  "hidden_fields": ['host', 'schema', 'login', 'password', 'port', 'extra'],

diff --git a/airflow/providers/docker/hooks/docker.py b/airflow/providers/docker/hooks/docker.py
@@ -15,7 +15,7 @@
 # KIND, either express or implied. See the License for the
 # specific language governing permissions and limitations
 # under the License.
-from typing import Dict, Optional
+from typing import Any, Dict, Optional
 
 from docker import APIClient
 from docker.errors import APIError
@@ -39,7 +39,7 @@ class DockerHook(BaseHook, LoggingMixin):
  hook_name = 'Docker'
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """Returns custom field behaviour"""
  return {
  "hidden_fields": ['schema'],

diff --git a/airflow/providers/google/cloud/hooks/compute_ssh.py b/airflow/providers/google/cloud/hooks/compute_ssh.py
@@ -90,7 +90,7 @@ class ComputeEngineSSHHook(SSHHook):
  hook_name = 'Google Cloud SSH'
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  return {
  "hidden_fields": ['host', 'schema', 'login', 'password', 'port', 'extra'],
  "relabeling": {},

diff --git a/airflow/providers/google/common/hooks/base_google.py b/airflow/providers/google/common/hooks/base_google.py
@@ -197,7 +197,7 @@ def get_connection_form_widgets() -> Dict[str, Any]:
  }
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """Returns custom field behaviour"""
  return {
  "hidden_fields": ['host', 'schema', 'login', 'password', 'port', 'extra'],

diff --git a/airflow/providers/google/leveldb/hooks/leveldb.py b/airflow/providers/google/leveldb/hooks/leveldb.py
@@ -17,11 +17,24 @@
 """Hook for Level DB"""
 from typing import List, Optional
 
-import plyvel
-from plyvel import DB
-
-from airflow.exceptions import AirflowException
-from airflow.hooks.base import BaseHook
+try:
+ import plyvel
+ from plyvel import DB
+
+ from airflow.exceptions import AirflowException
+ from airflow.hooks.base import BaseHook
+
+except ImportError as e:
+ # Plyvel is an optional feature and if imports are missing, it should be silently ignored
+ # As of Airflow 2.3 and above the operator can throw OptionalProviderFeatureException
+ try:
+ from airflow.exceptions import AirflowOptionalProviderFeatureException
+ except ImportError:
+ # However, in order to keep backwards-compatibility with Airflow 2.1 and 2.2, if the
+ # 2.3 exception cannot be imported, the original ImportError should be raised.
+ # This try/except can be removed when the provider depends on Airflow >= 2.3.0
+ raise e from None
+ raise AirflowOptionalProviderFeatureException(e)
 
 DB_NOT_INITIALIZED_BEFORE = "The `get_conn` method should be called before!"
 

diff --git a/airflow/providers/jdbc/hooks/jdbc.py b/airflow/providers/jdbc/hooks/jdbc.py
@@ -54,7 +54,7 @@ def get_connection_form_widgets() -> Dict[str, Any]:
  }
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """Returns custom field behaviour"""
  return {
  "hidden_fields": ['port', 'schema', 'extra'],

diff --git a/airflow/providers/microsoft/azure/hooks/adx.py b/airflow/providers/microsoft/azure/hooks/adx.py
@@ -97,7 +97,7 @@ def get_connection_form_widgets() -> Dict[str, Any]:
  }
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """Returns custom field behaviour"""
  return {
  "hidden_fields": ['schema', 'port', 'extra'],

diff --git a/airflow/providers/microsoft/azure/hooks/base_azure.py b/airflow/providers/microsoft/azure/hooks/base_azure.py
@@ -57,7 +57,7 @@ def get_connection_form_widgets() -> Dict[str, Any]:
  }
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """Returns custom field behaviour"""
  import json
 

diff --git a/airflow/providers/microsoft/azure/hooks/batch.py b/airflow/providers/microsoft/azure/hooks/batch.py
@@ -56,7 +56,7 @@ def get_connection_form_widgets() -> Dict[str, Any]:
  }
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """Returns custom field behaviour"""
  return {
  "hidden_fields": ['schema', 'port', 'host', 'extra'],

diff --git a/airflow/providers/microsoft/azure/hooks/container_registry.py b/airflow/providers/microsoft/azure/hooks/container_registry.py
@@ -17,7 +17,7 @@
 # under the License.
 """Hook for Azure Container Registry"""
 
-from typing import Dict
+from typing import Any, Dict
 
 from azure.mgmt.containerinstance.models import ImageRegistryCredential
 
@@ -39,7 +39,7 @@ class AzureContainerRegistryHook(BaseHook):
  hook_name = 'Azure Container Registry'
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """Returns custom field behaviour"""
  return {
  "hidden_fields": ['schema', 'port', 'extra'],

diff --git a/airflow/providers/microsoft/azure/hooks/container_volume.py b/airflow/providers/microsoft/azure/hooks/container_volume.py
@@ -54,7 +54,7 @@ def get_connection_form_widgets() -> Dict[str, Any]:
  }
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """Returns custom field behaviour"""
  return {
  "hidden_fields": ['schema', 'port', 'host', "extra"],

diff --git a/airflow/providers/microsoft/azure/hooks/cosmos.py b/airflow/providers/microsoft/azure/hooks/cosmos.py
@@ -67,7 +67,7 @@ def get_connection_form_widgets() -> Dict[str, Any]:
  }
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """Returns custom field behaviour"""
  return {
  "hidden_fields": ['schema', 'port', 'host', 'extra'],

diff --git a/airflow/providers/microsoft/azure/hooks/data_factory.py b/airflow/providers/microsoft/azure/hooks/data_factory.py
@@ -144,7 +144,7 @@ def get_connection_form_widgets() -> Dict[str, Any]:
  }
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """Returns custom field behaviour"""
  return {
  "hidden_fields": ['schema', 'port', 'host', 'extra'],

diff --git a/airflow/providers/microsoft/azure/hooks/data_lake.py b/airflow/providers/microsoft/azure/hooks/data_lake.py
@@ -65,7 +65,7 @@ def get_connection_form_widgets() -> Dict[str, Any]:
  }
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """Returns custom field behaviour"""
  return {
  "hidden_fields": ['schema', 'port', 'host', 'extra'],

diff --git a/airflow/providers/microsoft/azure/hooks/fileshare.py b/airflow/providers/microsoft/azure/hooks/fileshare.py
@@ -64,7 +64,7 @@ def get_connection_form_widgets() -> Dict[str, Any]:
  }
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """Returns custom field behaviour"""
  return {
  "hidden_fields": ['schema', 'port', 'host', 'extra'],

diff --git a/airflow/providers/microsoft/azure/hooks/wasb.py b/airflow/providers/microsoft/azure/hooks/wasb.py
@@ -81,7 +81,7 @@ def get_connection_form_widgets() -> Dict[str, Any]:
  }
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """Returns custom field behaviour"""
  return {
  "hidden_fields": ['schema', 'port'],

diff --git a/airflow/providers/pagerduty/hooks/pagerduty.py b/airflow/providers/pagerduty/hooks/pagerduty.py
@@ -49,7 +49,7 @@ class PagerdutyHook(BaseHook):
  hook_name = "Pagerduty"
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """Returns custom field behaviour"""
  return {
  "hidden_fields": ['port', 'login', 'schema', 'host'],

diff --git a/airflow/providers/pagerduty/hooks/pagerduty_events.py b/airflow/providers/pagerduty/hooks/pagerduty_events.py
@@ -41,7 +41,7 @@ class PagerdutyEventsHook(BaseHook):
  hook_name = "Pagerduty Events"
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """Returns custom field behaviour"""
  return {
  "hidden_fields": ['port', 'login', 'schema', 'host', 'extra'],

diff --git a/airflow/providers/qubole/hooks/qubole.py b/airflow/providers/qubole/hooks/qubole.py
@@ -22,7 +22,7 @@
 import os
 import pathlib
 import time
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
 
 from qds_sdk.commands import (
  Command,
@@ -119,7 +119,7 @@ class QuboleHook(BaseHook):
  hook_name = 'Qubole'
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """Returns custom field behaviour"""
  return {
  "hidden_fields": ['login', 'schema', 'port', 'extra'],

diff --git a/airflow/providers/salesforce/hooks/salesforce.py b/airflow/providers/salesforce/hooks/salesforce.py
@@ -119,7 +119,7 @@ def get_connection_form_widgets() -> Dict[str, Any]:
  }
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """Returns custom field behaviour"""
  return {
  "hidden_fields": ["schema", "port", "extra", "host"],

diff --git a/airflow/providers/sftp/hooks/sftp.py b/airflow/providers/sftp/hooks/sftp.py
@@ -19,7 +19,7 @@
 import datetime
 import stat
 import warnings
-from typing import Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 import pysftp
 import tenacity
@@ -61,7 +61,7 @@ class SFTPHook(SSHHook):
  hook_name = 'SFTP'
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  return {
  "hidden_fields": ['schema'],
  "relabeling": {

diff --git a/airflow/providers/snowflake/hooks/snowflake.py b/airflow/providers/snowflake/hooks/snowflake.py
@@ -97,7 +97,7 @@ def get_connection_form_widgets() -> Dict[str, Any]:
  }
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """Returns custom field behaviour"""
  import json
 

diff --git a/airflow/providers/ssh/hooks/ssh.py b/airflow/providers/ssh/hooks/ssh.py
@@ -88,7 +88,7 @@ class SSHHook(BaseHook):
  hook_name = 'SSH'
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """Returns custom field behaviour"""
  return {
  "hidden_fields": ['schema'],

diff --git a/airflow/providers/yandex/hooks/yandex.py b/airflow/providers/yandex/hooks/yandex.py
@@ -95,7 +95,7 @@ def provider_user_agent(cls) -> Optional[str]:
  return None
 
  @staticmethod
- def get_ui_field_behaviour() -> Dict:
+ def get_ui_field_behaviour() -> Dict[str, Any]:
  """Returns custom field behaviour"""
  return {
  "hidden_fields": ['host', 'schema', 'login', 'password', 'port', 'extra'],