diff --git a/airflow/providers/apache/pinot/hooks/pinot.py b/airflow/providers/apache/pinot/hooks/pinot.py index 0246a5d3a0feed..0fe11a04e4b82b 100644 --- a/airflow/providers/apache/pinot/hooks/pinot.py +++ b/airflow/providers/apache/pinot/hooks/pinot.py @@ -56,6 +56,11 @@ class PinotAdminHook(BaseHook): "Exception" is in the output message. """ + conn_name_attr = "conn_id" + default_conn_name = "pinot_admin_default" + conn_type = "pinot_admin" + hook_name = "Pinot Admin" + def __init__( self, conn_id: str = "pinot_admin_default", @@ -258,6 +263,8 @@ class PinotDbApiHook(DbApiHook): conn_name_attr = "pinot_broker_conn_id" default_conn_name = "pinot_broker_default" + conn_type = "pinot" + hook_name = "Pinot Broker" supports_autocommit = False def get_conn(self) -> Any: diff --git a/airflow/providers/apache/pinot/provider.yaml b/airflow/providers/apache/pinot/provider.yaml index 51146a1e6aa69a..3c10b9c05f8185 100644 --- a/airflow/providers/apache/pinot/provider.yaml +++ b/airflow/providers/apache/pinot/provider.yaml @@ -59,5 +59,7 @@ hooks: - airflow.providers.apache.pinot.hooks.pinot connection-types: - - hook-class-name: airflow.providers.apache.pinot.hooks.pinot.PinotHook + - hook-class-name: airflow.providers.apache.pinot.hooks.pinot.PinotDbApiHook connection-type: pinot + - hook-class-name: airflow.providers.apache.pinot.hooks.pinot.PinotAdminHook + connection-type: pinot_admin