diff --git a/airflow/providers/apache/hive/hooks/hive.py b/airflow/providers/apache/hive/hooks/hive.py index cb66d362740ec..cd508b21866ce 100644 --- a/airflow/providers/apache/hive/hooks/hive.py +++ b/airflow/providers/apache/hive/hooks/hive.py @@ -116,6 +116,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 +131,7 @@ def get_connection_form_widgets(cls) -> dict[str, Any]: "principal": StringField( lazy_gettext("Principal"), widget=BS3TextFieldWidget(), default="hive/_HOST@EXAMPLE.COM" ), + "high_availability": BooleanField(lazy_gettext("High Availability"), default=False), } @classmethod @@ -159,7 +161,12 @@ def _prepare_cli_cmd(self) -> list[Any]: if self.use_beeline: 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) + else: + jdbc_url = f"jdbc:hive2://{conn.host}:{conn.port}/{conn.schema}" + self.log.info("High Availability not set, setting JDBC url as %s", jdbc_url) if conf.get("core", "security") == "kerberos": template = conn.extra_dejson.get("principal", "hive/_HOST@EXAMPLE.COM") if "_HOST" in template: @@ -170,6 +177,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 not jdbc_url.endswith(";"): + jdbc_url += ";" + jdbc_url += "serviceDiscoveryMode=zooKeeper;ssl=true;zooKeeperNamespace=hiveserver2" elif self.auth: jdbc_url += ";auth=" + self.auth @@ -186,7 +197,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)" ) diff --git a/docs/apache-airflow-providers-apache-hive/connections/hive_cli.rst b/docs/apache-airflow-providers-apache-hive/connections/hive_cli.rst index cc52f1db92be2..5e88df971d989 100644 --- a/docs/apache-airflow-providers-apache-hive/connections/hive_cli.rst +++ b/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. diff --git a/tests/providers/apache/hive/hooks/test_hive.py b/tests/providers/apache/hive/hooks/test_hive.py index b69f68b149632..e444697bf69bd 100644 --- a/tests/providers/apache/hive/hooks/test_hive.py +++ b/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]