From 7ccda26af0e5f1a8e987059704fe3a5701f527bd Mon Sep 17 00:00:00 2001
From: Hossein <blcksrx@pm.me>
Date: Sat, 28 Oct 2023 20:59:09 +0200
Subject: [PATCH] enhance druid connection parameters

---
 airflow/providers/apache/druid/hooks/druid.py | 24 +++++---
 airflow/providers/apache/druid/provider.yaml  |  4 +-
 airflow/utils/db.py                           |  4 +-
 .../connections/druid_broker.rst              | 49 +++++++++++++++++
 .../connections/druid_ingest.rst              | 55 +++++++++++++++++++
 .../connections/index.rst                     | 31 +++++++++++
 .../index.rst                                 |  1 +
 .../apache/druid/hooks/test_druid.py          | 26 ++++++++-
 8 files changed, 183 insertions(+), 11 deletions(-)
 create mode 100644 docs/apache-airflow-providers-apache-druid/connections/druid_broker.rst
 create mode 100644 docs/apache-airflow-providers-apache-druid/connections/druid_ingest.rst
 create mode 100644 docs/apache-airflow-providers-apache-druid/connections/index.rst

diff --git a/airflow/providers/apache/druid/hooks/druid.py b/airflow/providers/apache/druid/hooks/druid.py
index 455c780f76ae9..6a7009f63c8b9 100644
--- a/airflow/providers/apache/druid/hooks/druid.py
+++ b/airflow/providers/apache/druid/hooks/druid.py
@@ -55,9 +55,14 @@ class DruidHook(BaseHook):
     :param max_ingestion_time: The maximum ingestion time before assuming the job failed
     """
 
+    conn_name_attr = "druid_ingest_conn_id"
+    default_conn_name = "druid_ingest_default"
+    conn_type = "druid_ingest"
+    hook_name = "Druid Ingest"
+
     def __init__(
         self,
-        druid_ingest_conn_id: str = "druid_ingest_default",
+        druid_ingest_conn_id: str = default_conn_name,
         timeout: int = 1,
         max_ingestion_time: int | None = None,
     ) -> None:
@@ -75,7 +80,7 @@ def get_conn_url(self, ingestion_type: IngestionType = IngestionType.BATCH) -> s
         conn = self.get_connection(self.druid_ingest_conn_id)
         host = conn.host
         port = conn.port
-        conn_type = conn.conn_type or "http"
+        conn_type = (conn.schema or "http") if conn.conn_type == self.conn_type else conn.conn_type
         if ingestion_type == IngestionType.BATCH:
             endpoint = conn.extra_dejson.get("endpoint", "")
         else:
@@ -104,7 +109,6 @@ def submit_indexing_job(
 
         self.log.info("Druid ingestion spec: %s", json_index_spec)
         req_index = requests.post(url, data=json_index_spec, headers=self.header, auth=self.get_auth())
-
         code = req_index.status_code
         if code != 200:
             self.log.error("Error submitting the Druid job to %s (%s) %s", url, code, req_index.content)
@@ -163,8 +167,8 @@ class DruidDbApiHook(DbApiHook):
 
     conn_name_attr = "druid_broker_conn_id"
     default_conn_name = "druid_broker_default"
-    conn_type = "druid"
-    hook_name = "Druid"
+    conn_type = "druid_broker"
+    hook_name = "Druid Broker"
     supports_autocommit = False
 
     def __init__(self, context: dict | None = None, *args, **kwargs) -> None:
@@ -174,6 +178,7 @@ def __init__(self, context: dict | None = None, *args, **kwargs) -> None:
     def get_conn(self) -> connect:
         """Establish a connection to druid broker."""
         conn = self.get_connection(getattr(self, self.conn_name_attr))
+
         druid_broker_conn = connect(
             host=conn.host,
             port=conn.port,
@@ -182,7 +187,12 @@ def get_conn(self) -> connect:
             user=conn.login,
             password=conn.password,
             context=self.context,
+            header=conn.extra_dejson.get("header", False),
+            ssl_verify_cert=conn.extra_dejson.get("ssl_verify_cert", True),
+            ssl_client_cert=conn.extra_dejson.get("ssl_client_cert"),
+            proxies=conn.extra_dejson.get("proxies"),
         )
+
         self.log.info("Get the connection to druid broker on %s using user %s", conn.host, conn.login)
         return druid_broker_conn
 
@@ -200,7 +210,7 @@ def get_uri(self) -> str:
         endpoint = conn.extra_dejson.get("endpoint", "druid/v2/sql")
         return f"{conn_type}://{host}/{endpoint}"
 
-    def set_autocommit(self, conn: connect, autocommit: bool) -> NotImplementedError:
+    def set_autocommit(self, conn: connect, autocommit: bool) -> None:
         raise NotImplementedError()
 
     def insert_rows(
@@ -211,5 +221,5 @@ def insert_rows(
         commit_every: int = 1000,
         replace: bool = False,
         **kwargs: Any,
-    ) -> NotImplementedError:
+    ) -> None:
         raise NotImplementedError()
diff --git a/airflow/providers/apache/druid/provider.yaml b/airflow/providers/apache/druid/provider.yaml
index 0e982d67276b1..7cdcf60a73914 100644
--- a/airflow/providers/apache/druid/provider.yaml
+++ b/airflow/providers/apache/druid/provider.yaml
@@ -74,7 +74,9 @@ hooks:
 
 connection-types:
   - hook-class-name: airflow.providers.apache.druid.hooks.druid.DruidDbApiHook
-    connection-type: druid
+    connection-type: druid_broker
+  - hook-class-name: airflow.providers.apache.druid.hooks.druid.DruidHook
+    connection-type: druid_ingest
 
 transfers:
   - source-integration-name: Apache Hive
diff --git a/airflow/utils/db.py b/airflow/utils/db.py
index 39f344299775d..7097a9b291b90 100644
--- a/airflow/utils/db.py
+++ b/airflow/utils/db.py
@@ -225,7 +225,7 @@ def create_default_connections(session: Session = NEW_SESSION):
     merge_conn(
         Connection(
             conn_id="druid_broker_default",
-            conn_type="druid",
+            conn_type="druid_broker",
             host="druid-broker",
             port=8082,
             extra='{"endpoint": "druid/v2/sql"}',
@@ -235,7 +235,7 @@ def create_default_connections(session: Session = NEW_SESSION):
     merge_conn(
         Connection(
             conn_id="druid_ingest_default",
-            conn_type="druid",
+            conn_type="druid_ingest",
             host="druid-overlord",
             port=8081,
             extra='{"endpoint": "druid/indexer/v1/task"}',
diff --git a/docs/apache-airflow-providers-apache-druid/connections/druid_broker.rst b/docs/apache-airflow-providers-apache-druid/connections/druid_broker.rst
new file mode 100644
index 0000000000000..2de0cffe85de6
--- /dev/null
+++ b/docs/apache-airflow-providers-apache-druid/connections/druid_broker.rst
@@ -0,0 +1,49 @@
+ .. 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.
+
+
+
+.. _howto/connection:druid_broker:
+
+Apache Druid Broker Connection
+==============================
+
+The Apache Druid connection type configures a connection to Apache Druid Broker via the ``pydruid`` DBAPI Python package.
+
+
+Default Connection IDs
+----------------------
+
+Druid hooks and operators use ``druid_broker_default`` by default.
+
+Configuring the Connection
+--------------------------
+Host (required)
+    The host of the Druid broker server.
+
+Port (optional)
+    The port of the Druid broker to connect to.
+
+Extra (optional)
+     A JSON dictionary specifying the extra parameters that can be used in pydruid connection.
+
+    * ``endpoint`` - endpoint to connect druid broker. Defaults to ``/druid/v2/sql``
+    * ``schema`` - connection schema (http or https). Defaults to ``http``
+    * ``header`` - to get header or not. Defaults to ``False``
+    * ``ssl_verify_cert`` - verify ssl certificate or not. Defaults to ``True``
+    * ``ssl_client_cert`` - ssl client certificate
+    * ``proxies`` - map of proxies schema to proxy server. e.g ``{'http': 'localhost:9050'}``
diff --git a/docs/apache-airflow-providers-apache-druid/connections/druid_ingest.rst b/docs/apache-airflow-providers-apache-druid/connections/druid_ingest.rst
new file mode 100644
index 0000000000000..b2cdfabae0ae4
--- /dev/null
+++ b/docs/apache-airflow-providers-apache-druid/connections/druid_ingest.rst
@@ -0,0 +1,55 @@
+ .. 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.
+
+
+
+.. _howto/connection:druid_ingest:
+
+Apache Druid Ingest Connection
+==============================
+
+The Apache Druid connection type configures a connection to Apache Druid Overload
+via the RestAPI to submit Indexing Jobs.
+
+
+Default Connection IDs
+----------------------
+
+Druid hooks and operators use ``druid_ingest_default`` by default.
+
+Configuring the Connection
+--------------------------
+Host (required)
+    The host of the Druid Overload server.
+
+Port (optional)
+    The port of the Druid Overload to connect to.
+
+Login (optional)
+    The username to connect the Druid Overload server.
+
+Password (optional)
+    The password to connect the Druid Overload server.
+
+Schema (optional)
+    The connection schema. Defaults to ``http``
+
+Extra (optional)
+     A JSON dictionary specifying the extra parameters that can be used in pydruid connection.
+
+    * ``endpoint`` - endpoint to connect druid Overload (Batch Ingestion )
+    * ``msq_endpoint`` - endpoint to connect druid Overload
diff --git a/docs/apache-airflow-providers-apache-druid/connections/index.rst b/docs/apache-airflow-providers-apache-druid/connections/index.rst
new file mode 100644
index 0000000000000..9bd881701d7f0
--- /dev/null
+++ b/docs/apache-airflow-providers-apache-druid/connections/index.rst
@@ -0,0 +1,31 @@
+ .. 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.
+
+Connection Types
+================
+
+Apache Airflow implemented connection to Apache Druid in different ways that listed here:
+    * Connection that uses the druid ingestion type that makes it possible to submit index/ingestion jobs.
+      To use the ingestion job, choose ``Druid Ingest Connection``
+    * Connection that uses to Query data using SQL.
+      To Query data use ``Druid Broker Connection``
+
+.. toctree::
+    :maxdepth: 1
+    :glob:
+
+    *
diff --git a/docs/apache-airflow-providers-apache-druid/index.rst b/docs/apache-airflow-providers-apache-druid/index.rst
index 8d6f206ef64cb..0d5c8d41d4f1d 100644
--- a/docs/apache-airflow-providers-apache-druid/index.rst
+++ b/docs/apache-airflow-providers-apache-druid/index.rst
@@ -33,6 +33,7 @@
     :maxdepth: 1
     :caption: Guides
 
+    Connection types <connections/index>
     Operators <operators>
 
 .. toctree::
diff --git a/tests/providers/apache/druid/hooks/test_druid.py b/tests/providers/apache/druid/hooks/test_druid.py
index 221371eac6cec..5d45750e66b45 100644
--- a/tests/providers/apache/druid/hooks/test_druid.py
+++ b/tests/providers/apache/druid/hooks/test_druid.py
@@ -180,9 +180,22 @@ def get_conn_url(self, ingestion_type: IngestionType = IngestionType.BATCH):
 
     @patch("airflow.providers.apache.druid.hooks.druid.DruidHook.get_connection")
     def test_get_conn_url(self, mock_get_connection):
+        get_conn_value = MagicMock()
+        get_conn_value.host = "test_host"
+        get_conn_value.schema = "https"
+        get_conn_value.conn_type = "druid_ingest"
+        get_conn_value.port = "1"
+        get_conn_value.extra_dejson = {"endpoint": "ingest"}
+        mock_get_connection.return_value = get_conn_value
+        hook = DruidHook(timeout=1, max_ingestion_time=5)
+        assert hook.get_conn_url() == "https://test_host:1/ingest"
+
+    @patch("airflow.providers.apache.druid.hooks.druid.DruidHook.get_connection")
+    def test_get_legacy_conn_url(self, mock_get_connection):
         get_conn_value = MagicMock()
         get_conn_value.host = "test_host"
         get_conn_value.conn_type = "https"
+        get_conn_value.schema = None
         get_conn_value.port = "1"
         get_conn_value.extra_dejson = {"endpoint": "ingest"}
         mock_get_connection.return_value = get_conn_value
@@ -271,7 +284,14 @@ def test_get_conn_with_context(
         get_conn_value.login = "test_login"
         get_conn_value.password = "test_password"
         get_conn_value.port = 10000
-        get_conn_value.extra_dejson = {"endpoint": "/test/endpoint", "schema": "https"}
+        get_conn_value.extra_dejson = {
+            "endpoint": "/test/endpoint",
+            "schema": "https",
+            "header": True,
+            "ssl_verify_cert": False,
+            "ssl_client_cert": "cert",
+            "proxies": "localhost:9090",
+        }
         mock_get_connection.return_value = get_conn_value
         hook = DruidDbApiHook(context=specified_context)
         hook.get_conn()
@@ -283,6 +303,10 @@ def test_get_conn_with_context(
             user="test_login",
             password="test_password",
             context=passed_context,
+            header=True,
+            ssl_verify_cert=False,
+            ssl_client_cert="cert",
+            proxies="localhost:9090",
         )
 
     def test_get_uri(self):
