[airflow] branch main updated: Sanitize host in drill hook (#30215)

[potiuk]

This is an automated email from the ASF dual-hosted git repository.

potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 63d9b24aad Sanitize host in drill hook (#30215)
63d9b24aad is described below

commit 63d9b24aad0b4b9397682ddac1ea5824354789b3
Author: Jarek Potiuk <ja...@potiuk.com>
AuthorDate: Wed Mar 22 09:33:39 2023 +0100

    Sanitize host in drill hook (#30215)
    
    The host passed in drill connection might contain some invalid
    characters. We should sanitize and reject them.
---
 airflow/providers/apache/drill/hooks/drill.py    |  2 ++
 tests/providers/apache/drill/hooks/test_drill.py | 26 +++++++++++++++++++++++-
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/airflow/providers/apache/drill/hooks/drill.py 
b/airflow/providers/apache/drill/hooks/drill.py
index 1d847e27e6..5785cb0515 100644
--- a/airflow/providers/apache/drill/hooks/drill.py
+++ b/airflow/providers/apache/drill/hooks/drill.py
@@ -49,6 +49,8 @@ class DrillHook(DbApiHook):
         """Establish a connection to Drillbit."""
         conn_md = self.get_connection(getattr(self, self.conn_name_attr))
         creds = f"{conn_md.login}:{conn_md.password}@" if conn_md.login else ""
+        if "/" in conn_md.host or "&" in conn_md.host:
+            raise ValueError("Drill host should not contain '/&' characters")
         engine = create_engine(
             f'{conn_md.extra_dejson.get("dialect_driver", 
"drill+sadrill")}://{creds}'
             f"{conn_md.host}:{conn_md.port}/"
diff --git a/tests/providers/apache/drill/hooks/test_drill.py 
b/tests/providers/apache/drill/hooks/test_drill.py
index c873b1ceef..241f50fce5 100644
--- a/tests/providers/apache/drill/hooks/test_drill.py
+++ b/tests/providers/apache/drill/hooks/test_drill.py
@@ -17,11 +17,35 @@
 # under the License.
 from __future__ import annotations
 
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
+
+import pytest
 
 from airflow.providers.apache.drill.hooks.drill import DrillHook
 
 
+@pytest.mark.parametrize(
+    "host, expect_error", [("host_with/", True), ("host_with&", True), 
("good_host", False)]
+)
+def test_get_host(host, expect_error):
+    with patch(
+        "airflow.providers.apache.drill.hooks.drill.DrillHook.get_connection"
+    ) as mock_get_connection, 
patch("sqlalchemy.engine.base.Engine.raw_connection") as raw_connection:
+        raw_connection.return_value = MagicMock()
+        mock_get_connection.return_value = MagicMock(
+            host=host, port=80, login="drill_user", password="secret"
+        )
+        mock_get_connection.return_value.extra_dejson = {
+            "dialect_driver": "drill+sadrill",
+            "storage_plugin": "dfs",
+        }
+        if expect_error:
+            with pytest.raises(ValueError):
+                DrillHook().get_conn()
+        else:
+            assert DrillHook().get_conn()
+
+
 class TestDrillHook:
     def setup_method(self):
         self.cur = MagicMock(rowcount=0)