Skip to content

Commit 667269a

Browse files
authored
Add temp location to batch file loads (#33115)
1 parent f38edb7 commit 667269a

File tree

1 file changed

+11
-0
lines changed
  • sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/providers

1 file changed

+11
-0
lines changed

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryManagedIT.java

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
import org.apache.beam.sdk.schemas.Schema;
3737
import org.apache.beam.sdk.testing.PAssert;
3838
import org.apache.beam.sdk.testing.TestPipeline;
39+
import org.apache.beam.sdk.testing.TestPipelineOptions;
3940
import org.apache.beam.sdk.transforms.Create;
4041
import org.apache.beam.sdk.transforms.MapElements;
4142
import org.apache.beam.sdk.transforms.PeriodicImpulse;
@@ -98,6 +99,10 @@ public void testBatchFileLoadsWriteRead() {
9899
String.format("%s:%s.%s", PROJECT, BIG_QUERY_DATASET_ID, testName.getMethodName());
99100
Map<String, Object> config = ImmutableMap.of("table", table);
100101

102+
// file loads requires a GCS temp location
103+
String tempLocation = writePipeline.getOptions().as(TestPipelineOptions.class).getTempRoot();
104+
writePipeline.getOptions().setTempLocation(tempLocation);
105+
101106
// batch write
102107
PCollectionRowTuple.of("input", getInput(writePipeline, false))
103108
.apply(Managed.write(Managed.BIGQUERY).withConfig(config));
@@ -146,6 +151,12 @@ public void testDynamicDestinations(boolean streaming) throws IOException, Inter
146151
Map<String, Object> config =
147152
ImmutableMap.of("table", destinationTemplate, "drop", Collections.singletonList("dest"));
148153

154+
if (!streaming) {
155+
// file loads requires a GCS temp location
156+
String tempLocation = writePipeline.getOptions().as(TestPipelineOptions.class).getTempRoot();
157+
writePipeline.getOptions().setTempLocation(tempLocation);
158+
}
159+
149160
// write
150161
PCollectionRowTuple.of("input", getInput(writePipeline, streaming))
151162
.apply(Managed.write(Managed.BIGQUERY).withConfig(config));

0 commit comments

Comments
 (0)