Add optional features in providers.

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."""
 

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'

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'],

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"],

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'],

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'],

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'],

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": {},

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'],

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!"
 

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'],

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'],

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
 

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'],

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'],

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"],

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'],

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'],

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'],

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'],

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'],

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'],

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'],

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'],

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"],

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": {

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
 

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'],

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'],