Adds support for Connection/Hook discovery from providers (#12466)

* Adds support for Hook discovery from providers

This PR extends providers discovery with the mechanism
of retrieving mapping of connections from type to hook.

Fixes #12456

* fixup! Adds support for Hook discovery from providers

* fixup! fixup! Adds support for Hook discovery from providers

Author: potiuk

.pre-commit-config.yaml

@@ -457,7 +457,7 @@ repos:
         entry: ./scripts/ci/pre_commit/pre_commit_check_provider_yaml_files.py
         language: python
         require_serial: true
-        files: provider.yaml$
+        files: provider.yaml$|scripts/ci/pre_commit/pre_commit_check_provider_yaml_files.py$
         additional_dependencies: ['PyYAML==5.3.1', 'jsonschema==3.2.0', 'tabulate==0.8.7']
       - id: mermaid
         name: Generate mermaid images

airflow/cli/cli_parser.py

@@ -1153,6 +1153,12 @@ class GroupCommand(NamedTuple):
     ),
 )
 PROVIDERS_COMMANDS = (
+    ActionCommand(
+        name='hooks',
+        help='List registered provider hooks',
+        func=lazy_load_command('airflow.cli.commands.provider_command.hooks_list'),
+        args=(ARG_OUTPUT,),
+    ),
     ActionCommand(
         name='list',
         help='List installed providers',

airflow/cli/commands/provider_command.py

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Providers sub-commands"""
-from typing import Dict, List
+from typing import Dict, List, Tuple
 
 import pygments
 import yaml
@@ -44,8 +44,7 @@ def provider_get(args):
     """Get a provider info."""
     providers = ProvidersManager().providers
     if args.provider_name in providers:
-        provider_version = providers[args.provider_name][0]
-        provider_info = providers[args.provider_name][1]
+        provider_version, provider_info = providers[args.provider_name]
         print("#")
         print(f"# Provider: {args.provider_name}")
         print(f"# Version: {provider_version}")
@@ -64,3 +63,23 @@ def provider_get(args):
 def providers_list(args):
     """Lists all providers at the command line"""
     print(_tabulate_providers(ProvidersManager().providers.values(), args.output))
+
+
+def _tabulate_hooks(hook_items: Tuple[str, Tuple[str, str]], tablefmt: str):
+    tabulate_data = [
+        {
+            'Connection type': hook_item[0],
+            'Hook class': hook_item[1][0],
+            'Hook connection attribute name': hook_item[1][1],
+        }
+        for hook_item in hook_items
+    ]
+
+    msg = tabulate(tabulate_data, tablefmt=tablefmt, headers='keys')
+    return msg
+
+
+def hooks_list(args):
+    """Lists all hooks at the command line"""
+    msg = _tabulate_hooks(ProvidersManager().hooks.items(), args.output)
+    print(msg)

airflow/models/connection.py

@@ -30,70 +30,10 @@
 from airflow.exceptions import AirflowException, AirflowNotFoundException
 from airflow.models.base import ID_LEN, Base
 from airflow.models.crypto import get_fernet
+from airflow.providers_manager import ProvidersManager
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.module_loading import import_string
 
-# A map that assigns a connection type to a tuple that contains
-# the path of the class and the name of the conn_id key parameter.
-# PLEASE KEEP BELOW LIST IN ALPHABETICAL ORDER.
-CONN_TYPE_TO_HOOK = {
-    "azure_batch": (
-        "airflow.providers.microsoft.azure.hooks.azure_batch.AzureBatchHook",
-        "azure_batch_conn_id",
-    ),
-    "azure_cosmos": (
-        "airflow.providers.microsoft.azure.hooks.azure_cosmos.AzureCosmosDBHook",
-        "azure_cosmos_conn_id",
-    ),
-    "azure_data_lake": (
-        "airflow.providers.microsoft.azure.hooks.azure_data_lake.AzureDataLakeHook",
-        "azure_data_lake_conn_id",
-    ),
-    "cassandra": ("airflow.providers.apache.cassandra.hooks.cassandra.CassandraHook", "cassandra_conn_id"),
-    "cloudant": ("airflow.providers.cloudant.hooks.cloudant.CloudantHook", "cloudant_conn_id"),
-    "dataprep": ("airflow.providers.google.cloud.hooks.dataprep.GoogleDataprepHook", "dataprep_default"),
-    "docker": ("airflow.providers.docker.hooks.docker.DockerHook", "docker_conn_id"),
-    "elasticsearch": (
-        "airflow.providers.elasticsearch.hooks.elasticsearch.ElasticsearchHook",
-        "elasticsearch_conn_id",
-    ),
-    "exasol": ("airflow.providers.exasol.hooks.exasol.ExasolHook", "exasol_conn_id"),
-    "gcpcloudsql": (
-        "airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLDatabaseHook",
-        "gcp_cloudsql_conn_id",
-    ),
-    "gcpssh": (
-        "airflow.providers.google.cloud.hooks.compute_ssh.ComputeEngineSSHHook",
-        "gcp_conn_id",
-    ),
-    "google_cloud_platform": (
-        "airflow.providers.google.cloud.hooks.bigquery.BigQueryHook",
-        "bigquery_conn_id",
-    ),
-    "grpc": ("airflow.providers.grpc.hooks.grpc.GrpcHook", "grpc_conn_id"),
-    "hive_cli": ("airflow.providers.apache.hive.hooks.hive.HiveCliHook", "hive_cli_conn_id"),
-    "hiveserver2": ("airflow.providers.apache.hive.hooks.hive.HiveServer2Hook", "hiveserver2_conn_id"),
-    "imap": ("airflow.providers.imap.hooks.imap.ImapHook", "imap_conn_id"),
-    "jdbc": ("airflow.providers.jdbc.hooks.jdbc.JdbcHook", "jdbc_conn_id"),
-    "jira": ("airflow.providers.jira.hooks.jira.JiraHook", "jira_conn_id"),
-    "kubernetes": ("airflow.providers.cncf.kubernetes.hooks.kubernetes.KubernetesHook", "kubernetes_conn_id"),
-    "mongo": ("airflow.providers.mongo.hooks.mongo.MongoHook", "conn_id"),
-    "mssql": ("airflow.providers.odbc.hooks.odbc.OdbcHook", "odbc_conn_id"),
-    "mysql": ("airflow.providers.mysql.hooks.mysql.MySqlHook", "mysql_conn_id"),
-    "odbc": ("airflow.providers.odbc.hooks.odbc.OdbcHook", "odbc_conn_id"),
-    "oracle": ("airflow.providers.oracle.hooks.oracle.OracleHook", "oracle_conn_id"),
-    "pig_cli": ("airflow.providers.apache.pig.hooks.pig.PigCliHook", "pig_cli_conn_id"),
-    "postgres": ("airflow.providers.postgres.hooks.postgres.PostgresHook", "postgres_conn_id"),
-    "presto": ("airflow.providers.presto.hooks.presto.PrestoHook", "presto_conn_id"),
-    "redis": ("airflow.providers.redis.hooks.redis.RedisHook", "redis_conn_id"),
-    "snowflake": ("airflow.providers.snowflake.hooks.snowflake.SnowflakeHook", "snowflake_conn_id"),
-    "sqlite": ("airflow.providers.sqlite.hooks.sqlite.SqliteHook", "sqlite_conn_id"),
-    "tableau": ("airflow.providers.salesforce.hooks.tableau.TableauHook", "tableau_conn_id"),
-    "vertica": ("airflow.providers.vertica.hooks.vertica.VerticaHook", "vertica_conn_id"),
-    "wasb": ("airflow.providers.microsoft.azure.hooks.wasb.WasbHook", "wasb_conn_id"),
-}
-# PLEASE KEEP ABOVE LIST IN ALPHABETICAL ORDER.
-
 
 def parse_netloc_to_hostname(*args, **kwargs):
     """This method is deprecated."""
@@ -326,7 +266,8 @@ def rotate_fernet_key(self):
 
     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))
+        hook_class_name, conn_id_param = ProvidersManager().hooks.get(self.conn_type, (None, None))
+
         if not hook_class_name:
             raise AirflowException(f'Unknown hook type "{self.conn_type}"')
         hook_class = import_string(hook_class_name)

airflow/plugins_manager.py

@@ -29,6 +29,7 @@
 import importlib_metadata
 
 from airflow import settings
+from airflow.utils.entry_points import entry_points_with_dist
 from airflow.utils.file import find_path_from_directory
 
 if TYPE_CHECKING:
@@ -169,23 +170,6 @@ def is_valid_plugin(plugin_obj):
     return False
 
 
-def entry_points_with_dist(group: str):
-    """
-    Return EntryPoint objects of the given group, along with the distribution information.
-
-    This is like the ``entry_points()`` function from importlib.metadata,
-    except it also returns the distribution the entry_point was loaded from.
-
-    :param group: FIlter results to only this entrypoint group
-    :return: Generator of (EntryPoint, Distribution) objects for the specified groups
-    """
-    for dist in importlib_metadata.distributions():
-        for e in dist.entry_points:
-            if e.group != group:
-                continue
-            yield (e, dist)
-
-
 def load_entrypoint_plugins():
     """
     Load and register plugins AirflowPlugin subclasses from the entrypoints.

airflow/provider.yaml.schema.json

@@ -173,6 +173,13 @@
           "python-module"
         ]
       }
+    },
+    "hook-class-names": {
+      "type": "array",
+      "description": "Hook class names that provide connection types to core",
+      "items": {
+        "type": "string"
+      }
     }
   },
   "additionalProperties": false,

airflow/providers/apache/cassandra/hooks/cassandra.py

@@ -83,7 +83,11 @@ class CassandraHook(BaseHook, LoggingMixin):
     For details of the Cluster config, see cassandra.cluster.
     """
 
-    def __init__(self, cassandra_conn_id: str = 'cassandra_default'):
+    conn_name_attr = 'cassandra_conn_id'
+    default_conn_name = 'cassandra_default'
+    conn_type = 'cassandra'
+
+    def __init__(self, cassandra_conn_id: str = default_conn_name):
         super().__init__()
         conn = self.get_connection(cassandra_conn_id)
 

airflow/providers/apache/cassandra/provider.yaml

@@ -41,3 +41,6 @@ hooks:
   - integration-name: Apache Cassandra
     python-modules:
       - airflow.providers.apache.cassandra.hooks.cassandra
+
+hook-class-names:
+  - airflow.providers.apache.cassandra.hooks.cassandra.CassandraHook

airflow/providers/apache/hive/hooks/hive.py

@@ -78,9 +78,13 @@ class HiveCliHook(BaseHook):
     :type  mapred_job_name: str
     """
 
+    conn_name_attr = 'hive_cli_conn_id'
+    default_conn_name = 'hive_cli_default'
+    conn_type = 'hive_cli'
+
     def __init__(
         self,
-        hive_cli_conn_id: str = "hive_cli_default",
+        hive_cli_conn_id: str = default_conn_name,
         run_as: Optional[str] = None,
         mapred_queue: Optional[str] = None,
         mapred_queue_priority: Optional[str] = None,
@@ -809,6 +813,7 @@ class HiveServer2Hook(DbApiHook):
 
     conn_name_attr = 'hiveserver2_conn_id'
     default_conn_name = 'hiveserver2_default'
+    conn_type = 'hiveserver2'
     supports_autocommit = False
 
     def get_conn(self, schema: Optional[str] = None) -> Any:

airflow/providers/apache/hive/provider.yaml

@@ -66,3 +66,7 @@ transfers:
   - source-integration-name: Microsoft SQL Server (MSSQL)
     target-integration-name: Apache Hive
     python-module: airflow.providers.apache.hive.transfers.mssql_to_hive
+
+hook-class-names:
+  - airflow.providers.apache.hive.hooks.hive.HiveCliHook
+  - airflow.providers.apache.hive.hooks.hive.HiveServer2Hook