fix(hive): pass drivername to metastore helper and fix describe test mocks (#26712)

Author: SaaiAravindhRaja

ingestion/src/metadata/ingestion/source/database/hive/metadata.py

@@ -183,9 +183,8 @@ def _build_metastore_partition_query(self, drivername: str) -> str:
         )
 
     def _get_partition_keys_from_metastore(
-        self, table_name: str, schema_name: str
+        self, table_name: str, schema_name: str, drivername: str
     ) -> List[str]:
-        drivername = getattr(getattr(self.engine, "url", None), "drivername", "")
         query = self._build_metastore_partition_query(drivername)
         result = self.connection.execute(
             text(query),
@@ -235,7 +234,7 @@ def get_table_partition_details(  # pylint: disable=unused-argument
             drivername = getattr(getattr(self.engine, "url", None), "drivername", "")
             if drivername in {"hive+mysql", "hive+postgres"}:
                 partition_keys = self._get_partition_keys_from_metastore(
-                    table_name, schema_name
+                    table_name, schema_name, drivername
                 )
             else:
                 partition_keys = self._get_partition_keys_from_describe(

ingestion/tests/unit/topology/database/test_hive.py

@@ -566,26 +566,15 @@ def test_get_table_partition_details_non_partitioned_metastore(self):
         self.assertIsNone(partition)
 
     def test_get_table_partition_details_non_partitioned_describe(self):
-        class ConnectionContextManager:
-            def __init__(self, connection):
-                self._connection = connection
-
-            def __enter__(self):
-                return self._connection
-
-            def __exit__(self, exc_type, exc, tb):
-                return False
-
+        self.hive.engine = types.SimpleNamespace(
+            url=types.SimpleNamespace(drivername="hive")
+        )
         mock_connection = Mock()
         mock_connection.execute.return_value = [
             ("id", "int", None),
             ("name", "string", None),
         ]
-
-        self.hive.engine = types.SimpleNamespace(
-            url=types.SimpleNamespace(drivername="hive"),
-            connect=lambda: ConnectionContextManager(mock_connection),
-        )
+        self.hive._connection_map[self.thread_id] = mock_connection
 
         is_partitioned, partition = self.hive.get_table_partition_details(
             table_name="non_partitioned", schema_name="sample_schema", inspector=Mock()
@@ -595,16 +584,9 @@ def __exit__(self, exc_type, exc, tb):
         self.assertIsNone(partition)
 
     def test_get_table_partition_details_from_describe_formatted(self):
-        class ConnectionContextManager:
-            def __init__(self, connection):
-                self._connection = connection
-
-            def __enter__(self):
-                return self._connection
-
-            def __exit__(self, exc_type, exc, tb):
-                return False
-
+        self.hive.engine = types.SimpleNamespace(
+            url=types.SimpleNamespace(drivername="hive")
+        )
         mock_connection = Mock()
         mock_connection.execute.return_value = [
             ("id", "int", None),
@@ -616,11 +598,7 @@ def __exit__(self, exc_type, exc, tb):
             ("# Detailed Table Information", None, None),
             ("Database:", "default", None),
         ]
-
-        self.hive.engine = types.SimpleNamespace(
-            url=types.SimpleNamespace(drivername="hive"),
-            connect=lambda: ConnectionContextManager(mock_connection),
-        )
+        self.hive._connection_map[self.thread_id] = mock_connection
 
         is_partitioned, partition = self.hive.get_table_partition_details(
             table_name="sample_table", schema_name="sample_schema", inspector=Mock()