[AIRFLOW-986] HiveCliHook ignores 'proxy_user' value in a connection'…

…s extra parameter (#5305) fix HiveCliHook ignores 'proxy_user' value in a connection's extra parameter
apache · Jun 3, 2019 · e8c5c7a · e8c5c7a
1 parent 67263d2
commit e8c5c7a
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 5 deletions.
diff --git a/airflow/hooks/hive_hooks.py b/airflow/hooks/hive_hooks.py
@@ -103,6 +103,21 @@ def __init__(
         self.mapred_queue_priority = mapred_queue_priority
         self.mapred_job_name = mapred_job_name
 
+    def _get_proxy_user(self):
+        """
+        This function set the proper proxy_user value in case the user overwtire the default.
+        """
+        conn = self.conn
+
+        proxy_user_value = conn.extra_dejson.get('proxy_user', "")
+        if proxy_user_value == "login" and conn.login:
+            return "hive.server2.proxy.user={0}".format(conn.login)
+        if proxy_user_value == "owner" and self.run_as:
+            return "hive.server2.proxy.user={0}".format(self.run_as)
+        if proxy_user_value != "":  # There is a custom proxy user
+            return "hive.server2.proxy.user={0}".format(proxy_user_value)
+        return proxy_user_value  # The default proxy user (undefined)
+
     def _prepare_cli_cmd(self):
         """
         This function creates the command list from available information
@@ -122,11 +137,7 @@ def _prepare_cli_cmd(self):
                     template = utils.replace_hostname_pattern(
                         utils.get_components(template))
 
-                proxy_user = ""  # noqa
-                if conn.extra_dejson.get('proxy_user') == "login" and conn.login:
-                    proxy_user = "hive.server2.proxy.user={0}".format(conn.login)
-                elif conn.extra_dejson.get('proxy_user') == "owner" and self.run_as:
-                    proxy_user = "hive.server2.proxy.user={0}".format(self.run_as)
+                proxy_user = self._get_proxy_user()
 
                 jdbc_url += ";principal={template};{proxy_user}".format(
                     template=template, proxy_user=proxy_user)

diff --git a/tests/operators/test_hive_operator.py b/tests/operators/test_hive_operator.py
@@ -60,6 +60,32 @@ def setUp(self):
         """
 
 
+class HiveCliTest(unittest.TestCase):
+
+    def setUp(self):
+        configuration.load_test_config()
+        self.nondefault_schema = "nondefault"
+        os.environ["AIRFLOW__CORE__SECURITY"] = "kerberos"
+
+    def tearDown(self):
+        del os.environ["AIRFLOW__CORE__SECURITY"]
+
+    def test_get_proxy_user_value(self):
+        from airflow.hooks.hive_hooks import HiveCliHook
+
+        hook = HiveCliHook()
+        returner = mock.MagicMock()
+        returner.extra_dejson = {'proxy_user': 'a_user_proxy'}
+        hook.use_beeline = True
+        hook.conn = returner
+
+        # Run
+        result = hook._prepare_cli_cmd()
+
+        # Verify
+        self.assertIn('hive.server2.proxy.user=a_user_proxy', result[2])
+
+
 class HiveOperatorConfigTest(HiveEnvironmentTest):
 
     def test_hive_airflow_default_config_queue(self):