From f49076563760352d37529e0ea56e857f9c90a4cd Mon Sep 17 00:00:00 2001
From: eladkal <45845474+eladkal@users.noreply.github.com>
Date: Tue, 28 Jun 2022 22:21:41 +0300
Subject: [PATCH] fix connection extra parameter `auth_mechanism` in
 `HiveMetastoreHook` and `HiveServer2Hook` closes:
 https://github.com/apache/airflow/issues/24692

---
 airflow/providers/apache/hive/hooks/hive.py   | 31 +++++++++++++++----
 .../providers/apache/hive/hooks/test_hive.py  |  2 +-
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/airflow/providers/apache/hive/hooks/hive.py b/airflow/providers/apache/hive/hooks/hive.py
index bccd279f77cf0..559c9727a8ad8 100644
--- a/airflow/providers/apache/hive/hooks/hive.py
+++ b/airflow/providers/apache/hive/hooks/hive.py
@@ -21,6 +21,7 @@
 import socket
 import subprocess
 import time
+import warnings
 from collections import OrderedDict
 from tempfile import NamedTemporaryFile, TemporaryDirectory
 from typing import Any, Dict, List, Optional, Union
@@ -492,10 +493,19 @@ def get_metastore_client(self) -> Any:
         if not host:
             raise AirflowException("Failed to locate the valid server.")
 
-        auth_mechanism = conn.extra_dejson.get('authMechanism', 'NOSASL')
+        if 'authMechanism' in conn.extra_dejson:
+            warnings.warn(
+                "The 'authMechanism' option is deprecated. Please use 'auth_mechanism'.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            conn.extra_dejson['auth_mechanism'] = conn.extra_dejson['authMechanism']
+            del conn.extra_dejson['authMechanism']
+
+        auth_mechanism = conn.extra_dejson.get('auth_mechanism', 'NOSASL')
 
         if conf.get('core', 'security') == 'kerberos':
-            auth_mechanism = conn.extra_dejson.get('authMechanism', 'GSSAPI')
+            auth_mechanism = conn.extra_dejson.get('auth_mechanism', 'GSSAPI')
             kerberos_service_name = conn.extra_dejson.get('kerberos_service_name', 'hive')
 
         conn_socket = TSocket.TSocket(host, conn.port)
@@ -779,7 +789,7 @@ class HiveServer2Hook(DbApiHook):
     Wrapper around the pyhive library
 
     Notes:
-    * the default authMechanism is PLAIN, to override it you
+    * the default auth_mechanism is PLAIN, to override it you
     can specify it in the ``extra`` of your connection in the UI
     * the default for run_set_variable_statements is true, if you
     are using impala you may need to set it to false in the
@@ -803,19 +813,28 @@ def get_conn(self, schema: Optional[str] = None) -> Any:
 
         db = self.get_connection(self.hiveserver2_conn_id)  # type: ignore
 
-        auth_mechanism = db.extra_dejson.get('authMechanism', 'NONE')
+        if 'authMechanism' in db.extra_dejson:
+            warnings.warn(
+                "The 'authMechanism' option is deprecated. Please use 'auth_mechanism'.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            db.extra_dejson['auth_mechanism'] = db.extra_dejson['authMechanism']
+            del db.extra_dejson['authMechanism']
+
+        auth_mechanism = db.extra_dejson.get('auth_mechanism', 'NONE')
         if auth_mechanism == 'NONE' and db.login is None:
             # we need to give a username
             username = 'airflow'
         kerberos_service_name = None
         if conf.get('core', 'security') == 'kerberos':
-            auth_mechanism = db.extra_dejson.get('authMechanism', 'KERBEROS')
+            auth_mechanism = db.extra_dejson.get('auth_mechanism', 'KERBEROS')
             kerberos_service_name = db.extra_dejson.get('kerberos_service_name', 'hive')
 
         # pyhive uses GSSAPI instead of KERBEROS as a auth_mechanism identifier
         if auth_mechanism == 'GSSAPI':
             self.log.warning(
-                "Detected deprecated 'GSSAPI' for authMechanism for %s. Please use 'KERBEROS' instead",
+                "Detected deprecated 'GSSAPI' for auth_mechanism for %s. Please use 'KERBEROS' instead",
                 self.hiveserver2_conn_id,  # type: ignore
             )
             auth_mechanism = 'KERBEROS'
diff --git a/tests/providers/apache/hive/hooks/test_hive.py b/tests/providers/apache/hive/hooks/test_hive.py
index 76892d6e8b372..49e26863be8d1 100644
--- a/tests/providers/apache/hive/hooks/test_hive.py
+++ b/tests/providers/apache/hive/hooks/test_hive.py
@@ -625,7 +625,7 @@ def test_get_conn_with_password(self, mock_connect):
 
         with mock.patch.dict(
             'os.environ',
-            {conn_env: "jdbc+hive2://conn_id:conn_pass@localhost:10000/default?authMechanism=LDAP"},
+            {conn_env: "jdbc+hive2://conn_id:conn_pass@localhost:10000/default?auth_mechanism=LDAP"},
         ):
             HiveServer2Hook(hiveserver2_conn_id=conn_id).get_conn()
             mock_connect.assert_called_once_with(