diff --git a/airflow/bin/airflow b/airflow/bin/airflow index 2f712ec005ecb..6874246e9070b 100755 --- a/airflow/bin/airflow +++ b/airflow/bin/airflow @@ -24,6 +24,10 @@ 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__': @@ -31,6 +35,10 @@ if __name__ == '__main__': 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() diff --git a/airflow/get_provider_info.py b/airflow/get_provider_info.py new file mode 100644 index 0000000000000..0c457e52eb760 --- /dev/null +++ b/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 diff --git a/airflow/models/connection.py b/airflow/models/connection.py index 0ab305f3ddbae..47477fc2349dd 100644 --- a/airflow/models/connection.py +++ b/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