From 20dc0178575c7f06d0fdb8a61b3863ceb171d2e9 Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Thu, 10 Feb 2022 13:48:06 -0800 Subject: [PATCH] Set larger limit get_partitions_by_filter in HiveMetastoreHook --- airflow/providers/apache/hive/hooks/hive.py | 4 +++- tests/providers/apache/hive/hooks/test_hive.py | 8 ++++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/airflow/providers/apache/hive/hooks/hive.py b/airflow/providers/apache/hive/hooks/hive.py index 6f76e236dce97..bccd279f77cf0 100644 --- a/airflow/providers/apache/hive/hooks/hive.py +++ b/airflow/providers/apache/hive/hooks/hive.py @@ -556,7 +556,9 @@ def check_for_partition(self, schema: str, table: str, partition: str) -> bool: True """ with self.metastore as client: - partitions = client.get_partitions_by_filter(schema, table, partition, 1) + partitions = client.get_partitions_by_filter( + schema, table, partition, HiveMetastoreHook.MAX_PART_COUNT + ) return bool(partitions) diff --git a/tests/providers/apache/hive/hooks/test_hive.py b/tests/providers/apache/hive/hooks/test_hive.py index 3171f207403b9..8587cb60f0b14 100644 --- a/tests/providers/apache/hive/hooks/test_hive.py +++ b/tests/providers/apache/hive/hooks/test_hive.py @@ -456,7 +456,9 @@ def test_check_for_partition(self): assert self.hook.check_for_partition(self.database, self.table, partition) - metastore.get_partitions_by_filter(self.database, self.table, partition, 1) + metastore.get_partitions_by_filter( + self.database, self.table, partition, HiveMetastoreHook.MAX_PART_COUNT + ) # Check for non-existent partition. missing_partition = f"{self.partition_by}='{self.next_day}'" @@ -464,7 +466,9 @@ def test_check_for_partition(self): assert not self.hook.check_for_partition(self.database, self.table, missing_partition) - metastore.get_partitions_by_filter.assert_called_with(self.database, self.table, missing_partition, 1) + metastore.get_partitions_by_filter.assert_called_with( + self.database, self.table, missing_partition, HiveMetastoreHook.MAX_PART_COUNT + ) def test_check_for_named_partition(self):