fix(hive): support metastore partition key extraction (#27029)

Author: SaaiAravindhRaja

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

@@ -185,21 +185,48 @@ def get_table_partition_details(
         partition_keys: List[str] = []
         in_partition_section = False
         try:
-            with self.engine.connect() as conn:
-                rows = conn.execute(
-                    text(f"DESCRIBE FORMATTED `{schema_name}`.`{table_name}`")
+            drivername = getattr(getattr(self.engine, "url", None), "drivername", "")
+            if drivername in {"hive+mysql", "hive+postgres"}:
+                query = (
+                    """
+                    SELECT pk.PKEY_NAME
+                    FROM PARTITION_KEYS pk
+                    JOIN TBLS tbl ON pk.TBL_ID = tbl.TBL_ID
+                    JOIN DBS db ON tbl.DB_ID = db.DB_ID
+                    WHERE db.NAME = :schema AND tbl.TBL_NAME = :table_name
+                    ORDER BY pk.INTEGER_IDX
+                    """
+                    if drivername == "hive+mysql"
+                    else """
+                    SELECT pk."PKEY_NAME"
+                    FROM "PARTITION_KEYS" pk
+                    JOIN "TBLS" tbl ON pk."TBL_ID" = tbl."TBL_ID"
+                    JOIN "DBS" db ON tbl."DB_ID" = db."DB_ID"
+                    WHERE db."NAME" = :schema AND tbl."TBL_NAME" = :table_name
+                    ORDER BY pk."INTEGER_IDX"
+                    """
                 )
-                for row in rows:
-                    col_name = row[0].strip() if row[0] else ""
-                    if col_name == "# Partition Information":
-                        in_partition_section = True
-                        continue
-                    if in_partition_section:
-                        if not col_name or col_name.startswith("# Detailed"):
-                            break
-                        if col_name.startswith("#"):
+                rows = self.connection.execute(
+                    text(query),
+                    {"table_name": table_name, "schema": schema_name},
+                ).fetchall()
+                partition_keys = [row[0] for row in rows if row and row[0]]
+            else:
+                with self.engine.connect() as conn:
+                    rows = conn.execute(
+                        text(f"DESCRIBE FORMATTED `{schema_name}`.`{table_name}`")
+                    )
+                    for row in rows:
+                        col_name = row[0].strip() if row[0] else ""
+                        if col_name == "# Partition Information":
+                            in_partition_section = True
                             continue
-                        partition_keys.append(col_name)
+                        if in_partition_section:
+                            if not col_name or col_name.startswith("# Detailed"):
+                                break
+                            if col_name.startswith("#"):
+                                continue
+                            partition_keys.append(col_name)
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.warning(

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

@@ -360,15 +360,15 @@ def __init__(
         self.thread_id = self.hive.context.get_current_thread_id()
         self.hive._inspector_map[self.thread_id] = types.SimpleNamespace()
 
-        self.hive._inspector_map[
-            self.thread_id
-        ].get_pk_constraint = lambda table_name, schema_name: []
-        self.hive._inspector_map[
-            self.thread_id
-        ].get_unique_constraints = lambda table_name, schema_name: []
-        self.hive._inspector_map[
-            self.thread_id
-        ].get_foreign_keys = lambda table_name, schema_name: []
+        self.hive._inspector_map[self.thread_id].get_pk_constraint = (
+            lambda table_name, schema_name: []
+        )
+        self.hive._inspector_map[self.thread_id].get_unique_constraints = (
+            lambda table_name, schema_name: []
+        )
+        self.hive._inspector_map[self.thread_id].get_foreign_keys = (
+            lambda table_name, schema_name: []
+        )
 
     def test_yield_database(self):
         assert EXPECTED_DATABASE == [
@@ -518,6 +518,41 @@ def test_get_columns_deduplicates_partition_column_with_sentinel(self):
         )
         self.assertEqual(col_names, ["id", "name", "dt"])
 
+    def test_get_table_partition_details_from_metastore_mysql(self):
+        self.hive.engine = types.SimpleNamespace(
+            url=types.SimpleNamespace(drivername="hive+mysql")
+        )
+        mock_connection = Mock()
+        mock_connection.execute.return_value.fetchall.return_value = [
+            ("dt",),
+            ("region",),
+        ]
+        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()
+        )
+
+        self.assertTrue(is_partitioned)
+        self.assertIsNotNone(partition)
+        self.assertEqual([c.columnName for c in partition.columns], ["dt", "region"])
+
+    def test_get_table_partition_details_from_metastore_postgres(self):
+        self.hive.engine = types.SimpleNamespace(
+            url=types.SimpleNamespace(drivername="hive+postgres")
+        )
+        mock_connection = Mock()
+        mock_connection.execute.return_value.fetchall.return_value = [("dt",)]
+        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()
+        )
+
+        self.assertTrue(is_partitioned)
+        self.assertIsNotNone(partition)
+        self.assertEqual([c.columnName for c in partition.columns], ["dt"])
+
     def test_ssl_connection_configuration(self):
         """
         Test SSL configuration in Hive connection