Fix MetastoreHivePartitionSensor pytests

apache · Dec 17, 2024 · 0e44a1f · 0e44a1f
1 parent 257adf0
commit 0e44a1f
Showing 1 changed file with 12 additions and 9 deletions.
diff --git a/providers/src/airflow/providers/google/cloud/hooks/dataproc_metastore.py b/providers/src/airflow/providers/google/cloud/hooks/dataproc_metastore.py
@@ -674,17 +674,20 @@ def list_hive_partitions(
         # because dictionaries are ordered since Python 3.7+
         _partitions = list(dict.fromkeys(partition_names)) if partition_names else []
 
-        query = f"""
-            SELECT PARTITIONS.*, TBLS.TBL_TYPE, TBLS.TBL_NAME
-            FROM PARTITIONS
-            INNER JOIN TBLS ON PARTITIONS.TBL_ID = TBLS.TBL_ID
-            WHERE TBLS.TBL_NAME = '{table}'"""
-
         if _partitions:
             partition_list = ", ".join(f"'{p}'" for p in _partitions)
-            query += f"""
-                AND PARTITIONS.PART_NAME IN ({partition_list})"""
-        query += ";"
+            query = f"""
+    SELECT PARTITIONS.*, TBLS.TBL_TYPE, TBLS.TBL_NAME
+    FROM PARTITIONS
+    INNER JOIN TBLS ON PARTITIONS.TBL_ID = TBLS.TBL_ID
+    WHERE TBLS.TBL_NAME = '{table}'
+        AND PARTITIONS.PART_NAME IN ({partition_list});"""
+        else:
+            query = f"""
+    SELECT PARTITIONS.*, TBLS.TBL_TYPE, TBLS.TBL_NAME
+    FROM PARTITIONS
+    INNER JOIN TBLS ON PARTITIONS.TBL_ID = TBLS.TBL_ID
+    WHERE TBLS.TBL_NAME = '{table}';"""
 
         request = {
             "service": f"projects/{project_id}/locations/{region}/services/{service_id}",