From 6724eeb6210d5965937eaf9dae3e476eb30f8268 Mon Sep 17 00:00:00 2001 From: Jarek Potiuk Date: Sun, 18 Jun 2023 08:50:12 +0200 Subject: [PATCH] Sanitize beeline principal parameter (#31983) Similar to other parameters of Beeline, principal should not contain the semicolon. --- airflow/providers/apache/hive/hooks/hive.py | 4 ++-- tests/providers/apache/hive/hooks/test_hive.py | 11 +++++++++++ 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/airflow/providers/apache/hive/hooks/hive.py b/airflow/providers/apache/hive/hooks/hive.py index fea826a8c39b4..945e6ba308bbc 100644 --- a/airflow/providers/apache/hive/hooks/hive.py +++ b/airflow/providers/apache/hive/hooks/hive.py @@ -150,9 +150,9 @@ def _prepare_cli_cmd(self) -> list[Any]: template = conn.extra_dejson.get("principal", "hive/_HOST@EXAMPLE.COM") if "_HOST" in template: template = utils.replace_hostname_pattern(utils.get_components(template)) - proxy_user = self._get_proxy_user() - + if ";" in template: + raise RuntimeError("The principal should not contain the ';' character") jdbc_url += f";principal={template};{proxy_user}" elif self.auth: jdbc_url += ";auth=" + self.auth diff --git a/tests/providers/apache/hive/hooks/test_hive.py b/tests/providers/apache/hive/hooks/test_hive.py index 5c0940d4097dc..9f55e870b7713 100644 --- a/tests/providers/apache/hive/hooks/test_hive.py +++ b/tests/providers/apache/hive/hooks/test_hive.py @@ -891,3 +891,14 @@ def test_get_proxy_user_value(self): # Verify assert "hive.server2.proxy.user=a_user_proxy" in result[2] + + def test_get_wrong_principal(self): + hook = MockHiveCliHook() + returner = mock.MagicMock() + returner.extra_dejson = {"principal": "principal with ; semicolon"} + hook.use_beeline = True + hook.conn = returner + + # Run + with pytest.raises(RuntimeError, match="The principal should not contain the ';' character"): + hook._prepare_cli_cmd()