Adding support to hive hook for high availability Hive installations

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

@@ -100,6 +100,9 @@ def __init__(
     ) -> None:
         super().__init__()
         conn = self.get_connection(hive_cli_conn_id)
+
+        print("conn is from __init__", conn)
+
         self.hive_cli_params: str = hive_cli_params
         self.use_beeline: bool = conn.extra_dejson.get("use_beeline", False)
         self.auth = auth
@@ -116,6 +119,7 @@ def __init__(
         self.mapred_queue_priority = mapred_queue_priority
         self.mapred_job_name = mapred_job_name
         self.proxy_user = proxy_user
+        self.high_availability = self.conn.extra_dejson.get("high_availability", False)
 
     @classmethod
     def get_connection_form_widgets(cls) -> dict[str, Any]:
@@ -130,6 +134,7 @@ def get_connection_form_widgets(cls) -> dict[str, Any]:
             "principal": StringField(
                 lazy_gettext("Principal"), widget=BS3TextFieldWidget(), default="hive/[email protected]"
             ),
+            "high_availability": BooleanField(lazy_gettext("High Availability"), default=False),
         }
 
     @classmethod
@@ -160,6 +165,9 @@ def _prepare_cli_cmd(self) -> list[Any]:
             hive_bin = "beeline"
             self._validate_beeline_parameters(conn)
             jdbc_url = f"jdbc:hive2://{conn.host}:{conn.port}/{conn.schema}"
+            if self.high_availability:
+                jdbc_url = f"jdbc:hive2://{conn.host}/{conn.schema}"
+                self.log.info("High Availability set, setting JDBC url as %s", jdbc_url)
             if conf.get("core", "security") == "kerberos":
                 template = conn.extra_dejson.get("principal", "hive/[email protected]")
                 if "_HOST" in template:
@@ -170,6 +178,10 @@ def _prepare_cli_cmd(self) -> list[Any]:
                 if ";" in proxy_user:
                     raise RuntimeError("The proxy_user should not contain the ';' character")
                 jdbc_url += f";principal={template};{proxy_user}"
+                if self.high_availability:
+                    if proxy_user:
+                        jdbc_url += ";"
+                    jdbc_url += "serviceDiscoveryMode=zooKeeper;ssl=true;zooKeeperNamespace=hiveserver2"
             elif self.auth:
                 jdbc_url += ";auth=" + self.auth
 
@@ -186,7 +198,13 @@ def _prepare_cli_cmd(self) -> list[Any]:
         return [hive_bin, *cmd_extra, *hive_params_list]
 
     def _validate_beeline_parameters(self, conn):
-        if ":" in conn.host or "/" in conn.host or ";" in conn.host:
+        if self.high_availability:
+            if ";" in conn.schema:
+                raise Exception(
+                    f"The schema used in beeline command ({conn.schema}) should not contain ';' character)"
+                )
+            return
+        elif ":" in conn.host or "/" in conn.host or ";" in conn.host:
             raise Exception(
                 f"The host used in beeline command ({conn.host}) should not contain ':/;' characters)"
             )

docs/apache-airflow-providers-apache-hive/connections/hive_cli.rst

@@ -73,6 +73,10 @@ Proxy User (optional)
 Principal (optional)
     Specify the JDBC Hive principal to be used with Hive Beeline.
 
+High Availability (optional)
+    Specify as ``True`` if you want to connect to a Hive installation running in high
+    availability mode. Specify host accordingly.
+
 
 When specifying the connection in environment variable you should specify
 it using URI syntax.

tests/providers/apache/hive/hooks/test_hive.py

@@ -901,3 +901,43 @@ def test_get_wrong_principal(self):
         # Run
         with pytest.raises(RuntimeError, match="The principal should not contain the ';' character"):
             hook._prepare_cli_cmd()
+
+    @pytest.mark.parametrize(
+        "extra_dejson, expected_keys",
+        [
+            (
+                {"high_availability": "true"},
+                "serviceDiscoveryMode=zooKeeper;ssl=true;zooKeeperNamespace=hiveserver2",
+            ),
+            (
+                {"high_availability": "false"},
+                "serviceDiscoveryMode=zooKeeper;ssl=true;zooKeeperNamespace=hiveserver2",
+            ),
+            ({}, "serviceDiscoveryMode=zooKeeper;ssl=true;zooKeeperNamespace=hiveserver2"),
+            # with proxy user
+            (
+                {"proxy_user": "a_user_proxy", "high_availability": "true"},
+                "hive.server2.proxy.user=a_user_proxy;"
+                "serviceDiscoveryMode=zooKeeper;ssl=true;zooKeeperNamespace=hiveserver2",
+            ),
+        ],
+    )
+    def test_high_availability(self, extra_dejson, expected_keys):
+        hook = MockHiveCliHook()
+        returner = mock.MagicMock()
+        returner.extra_dejson = extra_dejson
+        returner.login = "admin"
+        hook.use_beeline = True
+        hook.conn = returner
+        hook.high_availability = (
+            True
+            if ("high_availability" in extra_dejson and extra_dejson["high_availability"] == "true")
+            else False
+        )
+
+        result = hook._prepare_cli_cmd()
+
+        if hook.high_availability:
+            assert expected_keys in result[2]
+        else:
+            assert expected_keys not in result[2]