POC/WIP on provider package discovery mechanism (1.10)

airflow/bin/airflow

@@ -24,13 +24,21 @@ import argcomplete
 
 from airflow.configuration import conf
 from airflow.bin.cli import CLIFactory
+from airflow.get_provider_info import get_provider_info
+from airflow.models.connection import Connection
+
+from airflow.models.connection import CONN_TYPE_TO_HOOK
 
 if __name__ == '__main__':
 
     if conf.get("core", "security") == 'kerberos':
         os.environ['KRB5CCNAME'] = conf.get('kerberos', 'ccache')
         os.environ['KRB5_KTNAME'] = conf.get('kerberos', 'keytab')
 
+    get_provider_info(CONN_TYPE_TO_HOOK, Connection._types)
+    print(CONN_TYPE_TO_HOOK)
+    print(Connection._types)
+
     parser = CLIFactory.get_parser()
     argcomplete.autocomplete(parser)
     args = parser.parse_args()

airflow/get_provider_info.py

@@ -0,0 +1,80 @@
+# 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.
+
+import importlib
+import pkgutil
+
+
+def get_provider_info(conn_type_to_hook, connection_types):
+    """
+    Retrieves provider information and adds information to hooks/connection
+    :param conn_type_to_hook: dictionary of mapping for connections -> hooks
+    :param connection_types: array of connection types
+    :return: dictionary metadata about all providers installed
+    """
+    def ignore(_):
+        pass
+
+    def update_or_replace_connection_type(core_connection_types, provider_connection_types):
+        """
+        Updates or replaces all core connection types with those that come from a provider.
+        If same connection type is found in core types it is replaced with the provider one,
+        otherwise the provider connection type is appended to the list.
+
+        :param core_connection_types: all core connection tepes
+        :type core_connection_types List[Tuple]
+        :param provider_connection_types: provider connection types
+        :type provider_connection_types List[Tuple]
+        :return: None
+        """
+        for provider_connection_type in provider_connection_types:
+            for index, core_connection_type in enumerate(connection_types):
+                if core_connection_type[0] == provider_connection_type[0]:
+                    connection_types[index] = provider_connection_type
+                    break
+            core_connection_types.append(provider_connection_type)
+    try:
+        from airflow import providers
+    except ImportError:
+        print("No providers are available!")
+        return
+
+    providers_path = providers.__path__
+    providers_name = providers.__name__
+    provider_dict = {}
+
+    for (_, name, ispkg) in pkgutil.walk_packages(path=providers_path,
+                                                  prefix=providers_name + ".",
+                                                  onerror=ignore):
+        try:
+            if ispkg:
+                provider_info_module = importlib.import_module(".provider_info", package=name)
+                conn_type_to_hook.update(provider_info_module.CONN_TYPE_TO_HOOK)
+                update_or_replace_connection_type(connection_types, provider_info_module.connection_types)
+                provider_metadata = {
+                    'name': provider_info_module.PROVIDER_NAME,
+                    'version': provider_info_module.PROVIDER_VERSION,
+                    'url': provider_info_module.PROVIDER_URL,
+                    'docs': provider_info_module.PROVIDER_DOCS,
+                }
+                provider_dict[provider_info_module.PROVIDER_NAME] = provider_metadata
+                print(provider_metadata)
+        except ModuleNotFoundError:
+            pass
+        except Exception as e:  # noqa pylint: disable=broad-except
+            print("Provider {} could not be loaded because of {}".format(name, e))
+    return providers

airflow/models/connection.py

@@ -18,6 +18,10 @@
 # under the License.
 
 import json
+
+from typing import Dict, Tuple
+
+from airflow.utils.module_loading import import_string
 from builtins import bytes
 from urllib.parse import parse_qsl, quote, unquote, urlencode, urlparse
 
@@ -45,6 +49,9 @@ def parse_netloc_to_hostname(uri_parts):
     return hostname
 
 
+CONN_TYPE_TO_HOOK = {}  # type: Dict[str, Tuple[str, str]]
+
+
 class Connection(Base, LoggingMixin):
     """
     Placeholder to store information about different database instances
@@ -236,7 +243,7 @@ def rotate_fernet_key(self):
         if self._extra and self.is_extra_encrypted:
             self._extra = fernet.rotate(self._extra.encode('utf-8')).decode()
 
-    def get_hook(self):
+    def get_ariflow_1_10_hook(self):
         if self.conn_type == 'mysql':
             from airflow.hooks.mysql_hook import MySqlHook
             return MySqlHook(mysql_conn_id=self.conn_id)
@@ -308,6 +315,14 @@ def get_hook(self):
             return GrpcHook(grpc_conn_id=self.conn_id)
         raise AirflowException("Unknown hook type {}".format(self.conn_type))
 
+    def get_hook(self):
+        """Return hook based on conn_type."""
+        hook_class_name, conn_id_param = CONN_TYPE_TO_HOOK.get(self.conn_type, (None, None))
+        if not hook_class_name:
+            return self.get_ariflow_1_10_hook()
+        hook_class = import_string(hook_class_name)
+        return hook_class(**{conn_id_param: self.conn_id})
+
     def __repr__(self):
         return self.conn_id