From 945b0b7bf23d19735b4374d3de2d15dbbd90f710 Mon Sep 17 00:00:00 2001 From: PoAn Yang Date: Wed, 26 Aug 2026 22:47:35 +0900 Subject: [PATCH] Add missing Druid Ingest connection type Signed-off-by: PoAn Yang --- airflow-core/src/airflow/utils/db.py | 2 +- airflow-core/tests/unit/utils/test_db.py | 5 ++ .../apache/druid/docs/connections/druid.rst | 69 ++++++++++++++++ .../druid/docs/connections/druid_ingest.rst | 78 +++++++++++++++++++ .../apache/druid/docs/connections/index.rst | 28 +++++++ providers/apache/druid/docs/index.rst | 1 + providers/apache/druid/provider.yaml | 3 + .../apache/druid/get_provider_info.py | 7 +- .../providers/apache/druid/hooks/druid.py | 18 +++-- .../unit/apache/druid/hooks/test_druid.py | 30 +++++++ 10 files changed, 234 insertions(+), 7 deletions(-) create mode 100644 providers/apache/druid/docs/connections/druid.rst create mode 100644 providers/apache/druid/docs/connections/druid_ingest.rst create mode 100644 providers/apache/druid/docs/connections/index.rst diff --git a/airflow-core/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py index 71ec2d2004b6d..c11199bfca0f1 100644 --- a/airflow-core/src/airflow/utils/db.py +++ b/airflow-core/src/airflow/utils/db.py @@ -288,7 +288,7 @@ def get_default_connections(): ), 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/airflow-core/tests/unit/utils/test_db.py b/airflow-core/tests/unit/utils/test_db.py index a33f7b667f43b..6de60339fb048 100644 --- a/airflow-core/tests/unit/utils/test_db.py +++ b/airflow-core/tests/unit/utils/test_db.py @@ -47,6 +47,7 @@ compare_type, create_default_connections, downgrade, + get_default_connections, initdb, resetdb, upgradedb, @@ -234,6 +235,10 @@ def test_default_connections_sort(self): src = pattern.findall(source) assert sorted(src) == src + def test_default_druid_ingest_connection_uses_the_ingest_conn_type(self): + conn = next(c for c in get_default_connections() if c.conn_id == "druid_ingest_default") + assert conn.conn_type == "druid_ingest" + @pytest.mark.usefixtures("initialized_db") def test_check_migrations(self): # Should run without error. Can't easily test the behaviour, but we can check it works diff --git a/providers/apache/druid/docs/connections/druid.rst b/providers/apache/druid/docs/connections/druid.rst new file mode 100644 index 0000000000000..a8b62b33f6987 --- /dev/null +++ b/providers/apache/druid/docs/connections/druid.rst @@ -0,0 +1,69 @@ + .. 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: + +Apache Druid Connection +======================= + +The Apache Druid connection type enables connection to a Druid broker, in order to query data. +It is used by :class:`~airflow.providers.apache.druid.hooks.druid.DruidDbApiHook`. + +To submit ingestion tasks to a Druid overlord instead, use the +:ref:`Apache Druid Ingest connection `. + +Default Connection IDs +---------------------- + +The Druid broker hook uses the parameter ``druid_broker_conn_id`` for Connection IDs and the value of the +parameter as ``druid_broker_default`` by default. + +Configuring the Connection +-------------------------- +Host (required) + The host of the Druid broker, without a scheme, for example ``druid-broker``. + +Port (required) + The port the Druid broker listens on, ``8082`` by default. + +Login (optional) + The username used for basic authentication, when the cluster is secured with the + ``druid-basic-security`` extension. + +Password (optional) + The password used for basic authentication. + +Extra (optional, connection parameters) + Specify the extra parameters (as json dictionary) that can be used in the Druid connection. + The following parameters are supported: + + * ``endpoint`` - Path of the SQL endpoint of the broker. Defaults to ``/druid/v2/sql``. + * ``schema`` - The URL scheme used to reach the broker, either ``http`` or ``https``. + Defaults to ``http``. This is an extra parameter, not the connection's own Schema + field, which this hook does not read. + * ``ssl_verify_cert`` - Whether the broker TLS certificate is verified. Defaults to ``true``. + +When specifying the connection in environment variable you should specify +it using URI syntax. + +Note that all components of the URI should be URL-encoded. + +For example: + +.. code-block:: bash + + export AIRFLOW_CONN_DRUID_BROKER_DEFAULT='druid://druid-broker:8082/?endpoint=druid%2Fv2%2Fsql' diff --git a/providers/apache/druid/docs/connections/druid_ingest.rst b/providers/apache/druid/docs/connections/druid_ingest.rst new file mode 100644 index 0000000000000..50a2535468a77 --- /dev/null +++ b/providers/apache/druid/docs/connections/druid_ingest.rst @@ -0,0 +1,78 @@ + .. 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 Ingest connection type enables connection to a Druid overlord, which accepts +indexing (ingestion) tasks. It is used by :class:`~airflow.providers.apache.druid.hooks.druid.DruidHook` +and by the operators built on top of it, such as +:class:`~airflow.providers.apache.druid.operators.druid.DruidOperator`. + +To query a Druid broker instead, use the :ref:`Apache Druid connection `. + +Default Connection IDs +---------------------- + +The Druid Ingest hook, ``DruidOperator`` and ``HiveToDruidOperator`` use the parameter +``druid_ingest_conn_id`` for Connection IDs and the value of the parameter as +``druid_ingest_default`` by default. + +Configuring the Connection +-------------------------- +Host (required) + The host of the Druid overlord, without a scheme, for example ``druid-overlord``. + +Port (required) + The port the Druid overlord listens on, ``8081`` by default. + +Schema (optional) + The URL scheme used to reach the overlord, either ``http`` or ``https``. Defaults to ``http``. + +Login (optional) + The username used for basic authentication, when the cluster is secured with the + ``druid-basic-security`` extension. + +Password (optional) + The password used for basic authentication. Set Login and Password together: if either one + is empty the hook sends its requests without authentication. + +Extra (optional, connection parameters) + Specify the extra parameters (as json dictionary) that can be used in the Druid Ingest connection. + The following parameters are supported: + + * ``endpoint`` - Path of the endpoint that accepts native batch ingestion tasks, for example + ``druid/indexer/v1/task``. + * ``msq_endpoint`` - Path of the endpoint that accepts SQL-based (MSQ) ingestion tasks, for example + ``druid/v2/sql/task``. + * ``status_endpoint`` - Path of the endpoint used to poll the status of an SQL-based ingestion task. + Defaults to ``druid/indexer/v1/task``. + * ``ca_bundle_path`` - Path to a CA bundle used to verify the connection. Only used when the hook is + created with ``verify_ssl=False``. + +When specifying the connection in environment variable you should specify +it using URI syntax. + +Note that all components of the URI should be URL-encoded. + +For example: + +.. code-block:: bash + + export AIRFLOW_CONN_DRUID_INGEST_DEFAULT='druid-ingest://druid-overlord:8081/?endpoint=druid%2Findexer%2Fv1%2Ftask' diff --git a/providers/apache/druid/docs/connections/index.rst b/providers/apache/druid/docs/connections/index.rst new file mode 100644 index 0000000000000..878036ccf8d56 --- /dev/null +++ b/providers/apache/druid/docs/connections/index.rst @@ -0,0 +1,28 @@ + .. 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. + + + +Apache Druid Connections +======================== + + +.. toctree:: + :maxdepth: 1 + :glob: + + * diff --git a/providers/apache/druid/docs/index.rst b/providers/apache/druid/docs/index.rst index b3991a754686f..12c0f4e3ffbba 100644 --- a/providers/apache/druid/docs/index.rst +++ b/providers/apache/druid/docs/index.rst @@ -33,6 +33,7 @@ :maxdepth: 1 :caption: Guides + Connection types Operators .. toctree:: diff --git a/providers/apache/druid/provider.yaml b/providers/apache/druid/provider.yaml index 3af222d4b1ba5..f682d5b31ad54 100644 --- a/providers/apache/druid/provider.yaml +++ b/providers/apache/druid/provider.yaml @@ -99,6 +99,9 @@ connection-types: - hook-class-name: airflow.providers.apache.druid.hooks.druid.DruidDbApiHook hook-name: "Druid" connection-type: druid + - hook-class-name: airflow.providers.apache.druid.hooks.druid.DruidHook + hook-name: "Druid Ingest" + connection-type: druid_ingest transfers: - source-integration-name: Apache Hive diff --git a/providers/apache/druid/src/airflow/providers/apache/druid/get_provider_info.py b/providers/apache/druid/src/airflow/providers/apache/druid/get_provider_info.py index 0093743726464..26ec4f65dc21d 100644 --- a/providers/apache/druid/src/airflow/providers/apache/druid/get_provider_info.py +++ b/providers/apache/druid/src/airflow/providers/apache/druid/get_provider_info.py @@ -52,7 +52,12 @@ def get_provider_info(): "hook-class-name": "airflow.providers.apache.druid.hooks.druid.DruidDbApiHook", "hook-name": "Druid", "connection-type": "druid", - } + }, + { + "hook-class-name": "airflow.providers.apache.druid.hooks.druid.DruidHook", + "hook-name": "Druid Ingest", + "connection-type": "druid_ingest", + }, ], "transfers": [ { diff --git a/providers/apache/druid/src/airflow/providers/apache/druid/hooks/druid.py b/providers/apache/druid/src/airflow/providers/apache/druid/hooks/druid.py index 298733a201fec..7986657875fd9 100644 --- a/providers/apache/druid/src/airflow/providers/apache/druid/hooks/druid.py +++ b/providers/apache/druid/src/airflow/providers/apache/druid/hooks/druid.py @@ -61,9 +61,14 @@ class DruidHook(BaseHook): connection information for path to a CA bundle to use. Defaults to True """ + 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, verify_ssl: bool = True, @@ -87,10 +92,13 @@ def conn(self) -> Connection: @property def get_connection_type(self) -> str: if self.conn.schema: - conn_type = self.conn.schema - else: - conn_type = self.conn.conn_type or "http" - return conn_type + return self.conn.schema + # This value becomes the URL scheme, so a connection whose type is "druid_ingest" + # would build druid_ingest:// and fail. Use http for it. Any other conn_type passes + # through, since that is how a connection selects https. + if self.conn.conn_type and self.conn.conn_type != self.conn_type: + return self.conn.conn_type + return "http" def get_conn_url(self, ingestion_type: IngestionType = IngestionType.BATCH) -> str: """Get Druid connection url.""" diff --git a/providers/apache/druid/tests/unit/apache/druid/hooks/test_druid.py b/providers/apache/druid/tests/unit/apache/druid/hooks/test_druid.py index c2e82b84e5570..fc02ee341bd08 100644 --- a/providers/apache/druid/tests/unit/apache/druid/hooks/test_druid.py +++ b/providers/apache/druid/tests/unit/apache/druid/hooks/test_druid.py @@ -325,6 +325,36 @@ def test_get_status_url(self, mock_get_connection): hook = DruidHook(timeout=1, max_ingestion_time=5) assert hook.get_status_url(IngestionType.MSQ) == "https://test_host:1/druid/indexer/v1/task" + def test_connection_class_attributes(self): + assert DruidHook.conn_type == "druid_ingest" + assert DruidHook.conn_name_attr == "druid_ingest_conn_id" + assert DruidHook.default_conn_name == "druid_ingest_default" + assert DruidHook.hook_name == "Druid Ingest" + assert DruidHook().druid_ingest_conn_id == DruidHook.default_conn_name + + @pytest.mark.parametrize( + ("conn_type", "schema", "expected"), + [ + pytest.param("druid_ingest", None, "http", id="own-conn-type-is-not-a-scheme"), + pytest.param("druid_ingest", "https", "https", id="schema-still-wins"), + pytest.param("https", None, "https", id="other-conn-type-kept-as-scheme"), + pytest.param("druid", None, "druid", id="broker-conn-type-left-untouched"), + pytest.param(None, None, "http", id="no-conn-type-falls-back-to-http"), + ], + ) + @patch("airflow.providers.apache.druid.hooks.druid.DruidHook.get_connection") + def test_get_connection_type(self, mock_get_connection, conn_type, schema, expected): + get_conn_value = MagicMock() + get_conn_value.host = "test_host" + get_conn_value.port = "1" + get_conn_value.conn_type = conn_type + get_conn_value.schema = schema + get_conn_value.extra_dejson = {"endpoint": "ingest"} + mock_get_connection.return_value = get_conn_value + hook = DruidHook() + assert hook.get_connection_type == expected + assert hook.get_conn_url() == f"{expected}://test_host:1/ingest" + @patch("airflow.providers.apache.druid.hooks.druid.DruidHook.get_connection") def test_get_auth(self, mock_get_connection): get_conn_value = MagicMock()