Add temp location to batch file loads (#33115)

apache · Nov 14, 2024 · 667269a · 667269a
1 parent f38edb7
commit 667269a
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/...atform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryManagedIT.java b/...atform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryManagedIT.java
@@ -36,6 +36,7 @@
 import org.apache.beam.sdk.schemas.Schema;
 import org.apache.beam.sdk.testing.PAssert;
 import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.testing.TestPipelineOptions;
 import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.MapElements;
 import org.apache.beam.sdk.transforms.PeriodicImpulse;
@@ -98,6 +99,10 @@ public void testBatchFileLoadsWriteRead() {
         String.format("%s:%s.%s", PROJECT, BIG_QUERY_DATASET_ID, testName.getMethodName());
     Map<String, Object> config = ImmutableMap.of("table", table);
 
+    // file loads requires a GCS temp location
+    String tempLocation = writePipeline.getOptions().as(TestPipelineOptions.class).getTempRoot();
+    writePipeline.getOptions().setTempLocation(tempLocation);
+
     // batch write
     PCollectionRowTuple.of("input", getInput(writePipeline, false))
         .apply(Managed.write(Managed.BIGQUERY).withConfig(config));
@@ -146,6 +151,12 @@ public void testDynamicDestinations(boolean streaming) throws IOException, Inter
     Map<String, Object> config =
         ImmutableMap.of("table", destinationTemplate, "drop", Collections.singletonList("dest"));
 
+    if (!streaming) {
+      // file loads requires a GCS temp location
+      String tempLocation = writePipeline.getOptions().as(TestPipelineOptions.class).getTempRoot();
+      writePipeline.getOptions().setTempLocation(tempLocation);
+    }
+
     // write
     PCollectionRowTuple.of("input", getInput(writePipeline, streaming))
         .apply(Managed.write(Managed.BIGQUERY).withConfig(config));