Integration tests with require_partition_filter = true for DPO time p…

…artitioned tables
GoogleCloudDataproc · Dec 5, 2024 · fc0ffb8 · fc0ffb8
1 parent b842be6
commit fc0ffb8
Showing 1 changed file with 18 additions and 13 deletions.
diff --git a/...n/src/test/java/com/google/cloud/spark/bigquery/integration/WriteIntegrationTestBase.java b/...n/src/test/java/com/google/cloud/spark/bigquery/integration/WriteIntegrationTestBase.java
@@ -1591,6 +1591,11 @@ protected Dataset<Row> writeAndLoadDatasetOverwriteDynamicPartition(Dataset<Row>
         .option("temporaryGcsBucket", TestConstants.TEMPORARY_GCS_BUCKET)
         .save();
 
+    IntegrationTestUtils.runQuery(
+        String.format(
+            "ALTER TABLE %s.%s SET OPTIONS (require_partition_filter = false)",
+            testDataset, testTable));
+
     return spark
         .read()
         .format("bigquery")
@@ -1607,9 +1612,9 @@ public void testOverwriteDynamicPartition_partitionTimestampByHour() {
     IntegrationTestUtils.runQuery(
         String.format(
             "CREATE TABLE `%s.%s` (%s INTEGER, %s TIMESTAMP) "
-                + "PARTITION BY timestamp_trunc(order_date_time, HOUR) "
+                + "PARTITION BY timestamp_trunc(order_date_time, HOUR) OPTIONS (require_partition_filter = true)"
                 + "AS SELECT * FROM UNNEST([(1, TIMESTAMP '2023-09-28 1:00:00 UTC'), "
-                + "(2, TIMESTAMP '2023-09-28 10:00:00 UTC'), (3, TIMESTAMP '2023-09-28 10:30:00 UTC')])",
+                + "(2, TIMESTAMP '2023-09-28 10:00:00 UTC'), (3, TIMESTAMP '2023-09-28 10:30:00 UTC')]) ",
             testDataset, testTable, orderId, orderDateTime));
 
     Dataset<Row> df =
@@ -1651,7 +1656,7 @@ public void testOverwriteDynamicPartition_partitionTimestampByDay() {
     IntegrationTestUtils.runQuery(
         String.format(
             "CREATE TABLE `%s.%s` (%s INTEGER, %s TIMESTAMP) "
-                + "PARTITION BY DATE(order_date_time) "
+                + "PARTITION BY DATE(order_date_time) OPTIONS (require_partition_filter = true)"
                 + "AS SELECT * FROM UNNEST([(1, TIMESTAMP '2023-09-28 1:00:00 UTC'), "
                 + "(2, TIMESTAMP '2023-09-29 10:00:00 UTC'), (3, TIMESTAMP '2023-09-29 17:00:00 UTC')])",
             testDataset, testTable, orderId, orderDateTime));
@@ -1695,7 +1700,7 @@ public void testOverwriteDynamicPartition_partitionTimestampByMonth() {
     IntegrationTestUtils.runQuery(
         String.format(
             "CREATE TABLE `%s.%s` (%s INTEGER, %s TIMESTAMP) "
-                + "PARTITION BY timestamp_trunc(order_date_time, MONTH) "
+                + "PARTITION BY timestamp_trunc(order_date_time, MONTH) OPTIONS (require_partition_filter = true)"
                 + "AS SELECT * FROM UNNEST([(1, TIMESTAMP '2023-09-28 1:00:00 UTC'), "
                 + "(2, TIMESTAMP '2023-10-20 10:00:00 UTC'), (3, TIMESTAMP '2023-10-25 12:00:00 UTC')])",
             testDataset, testTable, orderId, orderDateTime));
@@ -1739,7 +1744,7 @@ public void testOverwriteDynamicPartition_partitionTimestampByYear() {
     IntegrationTestUtils.runQuery(
         String.format(
             "CREATE TABLE `%s.%s` (%s INTEGER, %s TIMESTAMP) "
-                + "PARTITION BY timestamp_trunc(order_date_time, YEAR) "
+                + "PARTITION BY timestamp_trunc(order_date_time, YEAR) OPTIONS (require_partition_filter = true)"
                 + "AS SELECT * FROM UNNEST([(1, TIMESTAMP '2022-09-28 1:00:00 UTC'), "
                 + "(2, TIMESTAMP '2023-10-20 10:00:00 UTC'), (2, TIMESTAMP '2023-10-25 12:00:00 UTC')])",
             testDataset, testTable, orderId, orderDateTime));
@@ -1783,7 +1788,7 @@ public void testOverwriteDynamicPartition_partitionDateByDay() {
     IntegrationTestUtils.runQuery(
         String.format(
             "CREATE TABLE `%s.%s` (%s INTEGER, %s DATE) "
-                + "PARTITION BY order_date "
+                + "PARTITION BY order_date OPTIONS (require_partition_filter = true)"
                 + "AS SELECT * FROM UNNEST([(1, DATE('2023-09-28')), (2, DATE('2023-09-29'))])",
             testDataset, testTable, orderId, orderDate));
 
@@ -1823,7 +1828,7 @@ public void testOverwriteDynamicPartition_partitionDateByMonth() {
     IntegrationTestUtils.runQuery(
         String.format(
             "CREATE TABLE `%s.%s` (%s INTEGER, %s DATE) "
-                + "PARTITION BY DATE_TRUNC(order_date, MONTH) "
+                + "PARTITION BY DATE_TRUNC(order_date, MONTH) OPTIONS (require_partition_filter = true)"
                 + "AS SELECT * FROM UNNEST([(1, DATE('2023-09-28')), "
                 + "(2, DATE('2023-10-29')), (2, DATE('2023-10-28'))])",
             testDataset, testTable, orderId, orderDate));
@@ -1864,7 +1869,7 @@ public void testOverwriteDynamicPartition_partitionDateByYear() {
     IntegrationTestUtils.runQuery(
         String.format(
             "CREATE TABLE `%s.%s` (%s INTEGER, %s DATE) "
-                + "PARTITION BY DATE_TRUNC(order_date, YEAR) "
+                + "PARTITION BY DATE_TRUNC(order_date, YEAR) OPTIONS (require_partition_filter = true)"
                 + "AS SELECT * FROM UNNEST([(1, DATE('2022-09-28')), "
                 + "(2, DATE('2023-10-29')), (2, DATE('2023-11-28'))])",
             testDataset, testTable, orderId, orderDate));
@@ -1905,7 +1910,7 @@ public void testOverwriteDynamicPartition_partitionDateTimeByHour() {
     TimeZone.setDefault(TimeZone.getTimeZone("UTC"));
     IntegrationTestUtils.runQuery(
         String.format(
-            "CREATE TABLE `%s.%s` (%s INTEGER, %s DATETIME) "
+            "CREATE TABLE `%s.%s` (%s INTEGER, %s DATETIME) OPTIONS (require_partition_filter = true)"
                 + "PARTITION BY timestamp_trunc(order_date_time, HOUR) "
                 + "AS SELECT * FROM UNNEST([(1, DATETIME '2023-09-28 1:00:00'), "
                 + "(2, DATETIME '2023-09-28 10:00:00'), (3, DATETIME '2023-09-28 10:30:00')])",
@@ -1948,7 +1953,7 @@ public void testOverwriteDynamicPartition_partitionDateTimeByDay() {
     IntegrationTestUtils.runQuery(
         String.format(
             "CREATE TABLE `%s.%s` (%s INTEGER, %s DATETIME) "
-                + "PARTITION BY timestamp_trunc(order_date_time, DAY) "
+                + "PARTITION BY timestamp_trunc(order_date_time, DAY) OPTIONS (require_partition_filter = true)"
                 + "AS SELECT * FROM UNNEST([(1, DATETIME '2023-09-28 1:00:00'), "
                 + "(2, DATETIME '2023-09-29 10:00:00'), (3, DATETIME '2023-09-29 17:30:00')])",
             testDataset, testTable, orderId, orderDateTime));
@@ -1990,7 +1995,7 @@ public void testOverwriteDynamicPartition_partitionDateTimeByMonth() {
     IntegrationTestUtils.runQuery(
         String.format(
             "CREATE TABLE `%s.%s` (%s INTEGER, %s DATETIME) "
-                + "PARTITION BY timestamp_trunc(order_date_time, MONTH) "
+                + "PARTITION BY timestamp_trunc(order_date_time, MONTH) OPTIONS (require_partition_filter = true)"
                 + "AS SELECT * FROM UNNEST([(1, DATETIME '2023-09-28 1:00:00'), "
                 + "(2, DATETIME '2023-10-29 10:00:00'), (3, DATETIME '2023-10-29 17:30:00')])",
             testDataset, testTable, orderId, orderDateTime));
@@ -2032,7 +2037,7 @@ public void testOverwriteDynamicPartition_partitionDateTimeByYear() {
     IntegrationTestUtils.runQuery(
         String.format(
             "CREATE TABLE `%s.%s` (%s INTEGER, %s DATETIME) "
-                + "PARTITION BY timestamp_trunc(order_date_time, YEAR) "
+                + "PARTITION BY timestamp_trunc(order_date_time, YEAR) OPTIONS (require_partition_filter = true)"
                 + "AS SELECT * FROM UNNEST([(1, DATETIME '2022-09-28 1:00:00'), "
                 + "(2, DATETIME '2023-10-29 10:00:00'), (3, DATETIME '2023-11-29 17:30:00')])",
             testDataset, testTable, orderId, orderDateTime));
@@ -2111,7 +2116,7 @@ public void testOverwriteDynamicPartition_rangePartitioned() {
     IntegrationTestUtils.runQuery(
         String.format(
             "CREATE TABLE `%s.%s` (%s INTEGER, %s INTEGER) "
-                + "PARTITION BY RANGE_BUCKET(order_id, GENERATE_ARRAY(1, 100, 10)) "
+                + "PARTITION BY RANGE_BUCKET(order_id, GENERATE_ARRAY(1, 100, 10))"
                 + "AS SELECT * FROM UNNEST([(1, 1000), "
                 + "(8, 1005), ( 21, 1010), (83, 1020)])",
             testDataset, testTable, orderId, orderCount));