This is an automated email from the ASF dual-hosted git repository.
stankiewicz pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 10d3b8d316d Fix temp table cleanup for continuous read in BigQueryIO
(#37535)
10d3b8d316d is described below
commit 10d3b8d316d76bc19b8f19508063388dd3896728
Author: Radosław Stankiewicz <[email protected]>
AuthorDate: Mon Oct 19 10:00:47 2026 +0200
Fix temp table cleanup for continuous read in BigQueryIO (#37535)
* Fix temp table cleanup for continuous read in BigQueryIO
* Fix missing coder for CleanupOperationMessage side outputs
* Fix BQ streaming input tasks: move Javadoc, fix empty table edge case,
and add test case.
---
.../beam/sdk/io/gcp/bigquery/BigQueryIO.java | 242 +++++++++++++++++++--
.../gcp/bigquery/BigQueryStorageStreamSource.java | 21 +-
.../io/gcp/bigquery/CleanupTempTableDoFnTest.java | 153 +++++++++++++
3 files changed, 396 insertions(+), 20 deletions(-)
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
index 0cefea4639f..07fa1a320e9 100644
---
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
@@ -23,6 +23,7 @@ import static
org.apache.beam.sdk.transforms.errorhandling.BadRecordRouter.BAD_R
import static
org.apache.beam.sdk.transforms.errorhandling.BadRecordRouter.RECORDING_ROUTER;
import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull;
import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
+import static
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects.firstNonNull;
import static
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
import static
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
@@ -117,9 +118,14 @@ import
org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
import org.apache.beam.sdk.schemas.FieldAccessDescriptor;
import org.apache.beam.sdk.schemas.ProjectionProducer;
import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.ValueState;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.DoFn.MultiOutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.StateId;
+import org.apache.beam.sdk.transforms.Flatten;
import org.apache.beam.sdk.transforms.MapElements;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
@@ -141,6 +147,7 @@ import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollection.IsBounded;
+import org.apache.beam.sdk.values.PCollectionList;
import org.apache.beam.sdk.values.PCollectionTuple;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.Row;
@@ -957,11 +964,18 @@ public class BigQueryIO {
DynamicRead() {}
class CreateBoundedSourceForTable
- extends DoFn<KV<String, BigQueryDynamicReadDescriptor>,
BigQueryStorageStreamSource<T>> {
+ extends DoFn<
+ KV<String, BigQueryDynamicReadDescriptor>, KV<String,
BigQueryStorageStreamSource<T>>> {
+
+ private final TupleTag<KV<String, CleanupOperationMessage>>
cleanupInfoTag;
+
+ CreateBoundedSourceForTable(TupleTag<KV<String,
CleanupOperationMessage>> cleanupInfoTag) {
+ this.cleanupInfoTag = cleanupInfoTag;
+ }
@ProcessElement
public void processElement(
- OutputReceiver<BigQueryStorageStreamSource<T>> receiver,
+ MultiOutputReceiver receiver,
@Element KV<String, BigQueryDynamicReadDescriptor> kv,
PipelineOptions options)
throws Exception {
@@ -986,7 +1000,14 @@ public class BigQueryIO {
// shards
long desiredChunkSize = getDesiredChunkSize(options, output);
List<BigQueryStorageStreamSource<T>> split =
output.split(desiredChunkSize, options);
- split.stream().forEach(source -> receiver.output(source));
+ split.stream()
+ .forEach(
+ source ->
+ receiver
+ .get(
+ new TupleTag<KV<String,
BigQueryStorageStreamSource<T>>>(
+ "mainOutput"))
+ .output(KV.of(kv.getKey(), source)));
} else {
// run query
BigQueryStorageQuerySource<T> querySource =
@@ -1023,7 +1044,21 @@ public class BigQueryIO {
// shards
long desiredChunkSize = getDesiredChunkSize(options, output);
List<BigQueryStorageStreamSource<T>> split =
output.split(desiredChunkSize, options);
- split.stream().forEach(source -> receiver.output(source));
+ split.stream()
+ .forEach(
+ source ->
+ receiver
+ .get(
+ new TupleTag<KV<String,
BigQueryStorageStreamSource<T>>>(
+ "mainOutput"))
+ .output(KV.of(kv.getKey(), source.withFromQuery())));
+ boolean datasetCreatedByBeam = getQueryTempDataset() == null;
+ CleanupInfo cleanupInfo =
+ new CleanupInfo(
+ queryResultTable.getTableReference(), datasetCreatedByBeam,
split.size());
+ receiver
+ .get(cleanupInfoTag)
+ .output(KV.of(kv.getKey(),
CleanupOperationMessage.initialize(cleanupInfo)));
}
}
@@ -1037,6 +1072,9 @@ public class BigQueryIO {
@Override
public PCollection<T> expand(PCollection<BigQueryDynamicReadDescriptor>
input) {
TupleTag<T> rowTag = new TupleTag<>();
+ TupleTag<KV<String, BigQueryStorageStreamSource<T>>> streamTag = new
TupleTag<>("mainOutput");
+ TupleTag<KV<String, CleanupOperationMessage>> cleanupInfoTag = new
TupleTag<>();
+
PCollection<KV<String, BigQueryDynamicReadDescriptor>> addJobId =
input
.apply(
@@ -1051,21 +1089,193 @@ public class BigQueryIO {
.apply("Checkpoint", Redistribute.byKey());
PCollectionTuple resultTuple =
- addJobId
- .apply("Create streams", ParDo.of(new
CreateBoundedSourceForTable()))
+ addJobId.apply(
+ "Create streams",
+ ParDo.of(new CreateBoundedSourceForTable(cleanupInfoTag))
+ .withOutputTags(streamTag, TupleTagList.of(cleanupInfoTag)));
+
+ PCollection<KV<String, BigQueryStorageStreamSource<T>>> streams =
+ resultTuple
+ .get(streamTag)
.setCoder(
- SerializableCoder.of(new
TypeDescriptor<BigQueryStorageStreamSource<T>>() {}))
- .apply("Redistribute", Redistribute.arbitrarily())
- .apply(
- "Read Streams with storage read api",
- ParDo.of(
- new TypedRead.ReadTableSource<T>(
- rowTag, getParseFn(), getBadRecordRouter()))
- .withOutputTags(rowTag,
TupleTagList.of(BAD_RECORD_TAG)));
+ KvCoder.of(
+ StringUtf8Coder.of(),
+ SerializableCoder.of(
+ new TypeDescriptor<BigQueryStorageStreamSource<T>>()
{})))
+ .apply("Redistribute", Redistribute.arbitrarily());
+
+ PCollectionTuple readResultTuple =
+ streams.apply(
+ "Read Streams with storage read api",
+ ParDo.of(
+ new ReadDynamicStreamSource<T>(
+ rowTag, getParseFn(), getBadRecordRouter(),
cleanupInfoTag))
+ .withOutputTags(rowTag,
TupleTagList.of(BAD_RECORD_TAG).and(cleanupInfoTag)));
+
+ PCollection<KV<String, CleanupOperationMessage>> cleanupMessages1 =
+ resultTuple
+ .get(cleanupInfoTag)
+ .setCoder(
+ KvCoder.of(
+ StringUtf8Coder.of(),
SerializableCoder.of(CleanupOperationMessage.class)));
+
+ PCollection<KV<String, CleanupOperationMessage>> cleanupMessages2 =
+ readResultTuple
+ .get(cleanupInfoTag)
+ .setCoder(
+ KvCoder.of(
+ StringUtf8Coder.of(),
SerializableCoder.of(CleanupOperationMessage.class)));
+
+ PCollectionList.of(cleanupMessages1)
+ .and(cleanupMessages2)
+ .apply(Flatten.pCollections())
+ .apply("CleanupTempTables", ParDo.of(new
CleanupTempTableDoFn(getBigQueryServices())));
+
getBadRecordErrorHandler()
.addErrorCollection(
-
resultTuple.get(BAD_RECORD_TAG).setCoder(BadRecord.getCoder(input.getPipeline())));
- return resultTuple.get(rowTag).setCoder(getOutputCoder());
+ readResultTuple
+ .get(BAD_RECORD_TAG)
+ .setCoder(BadRecord.getCoder(input.getPipeline())));
+ return readResultTuple.get(rowTag).setCoder(getOutputCoder());
+ }
+ }
+
+ static class CleanupInfo implements Serializable {
+ private final String projectId;
+ private final String datasetId;
+ private final String tableId;
+ private final boolean datasetCreatedByBeam;
+ private final int totalStreams;
+
+ public CleanupInfo(TableReference tableRef, boolean datasetCreatedByBeam,
int totalStreams) {
+ this.projectId = tableRef.getProjectId();
+ this.datasetId = tableRef.getDatasetId();
+ this.tableId = tableRef.getTableId();
+ this.datasetCreatedByBeam = datasetCreatedByBeam;
+ this.totalStreams = totalStreams;
+ }
+
+ public TableReference getTableReference() {
+ return new TableReference()
+ .setProjectId(projectId)
+ .setDatasetId(datasetId)
+ .setTableId(tableId);
+ }
+
+ public boolean isDatasetCreatedByBeam() {
+ return datasetCreatedByBeam;
+ }
+
+ public int getTotalStreams() {
+ return totalStreams;
+ }
+ }
+
+ static class CleanupOperationMessage implements Serializable {
+ private final @Nullable CleanupInfo cleanupInfo;
+ private final boolean isStreamCompletion;
+
+ private CleanupOperationMessage(@Nullable CleanupInfo cleanupInfo, boolean
isStreamCompletion) {
+ this.cleanupInfo = cleanupInfo;
+ this.isStreamCompletion = isStreamCompletion;
+ }
+
+ public static CleanupOperationMessage streamComplete() {
+ return new CleanupOperationMessage(null, true);
+ }
+
+ public static CleanupOperationMessage initialize(CleanupInfo cleanupInfo) {
+ return new CleanupOperationMessage(cleanupInfo, false);
+ }
+
+ public @Nullable CleanupInfo getCleanupInfo() {
+ return cleanupInfo;
+ }
+
+ public boolean isStreamCompletion() {
+ return isStreamCompletion;
+ }
+ }
+
+ static class CleanupTempTableDoFn extends DoFn<KV<String,
CleanupOperationMessage>, Void> {
+ private final BigQueryServices bqServices;
+ private static final Logger LOG =
LoggerFactory.getLogger(CleanupTempTableDoFn.class);
+
+ @StateId("cleanupInfo")
+ private final StateSpec<ValueState<CleanupInfo>> cleanupInfoSpec =
StateSpecs.value();
+
+ @StateId("completedStreams")
+ private final StateSpec<ValueState<Integer>> completedStreamsSpec =
StateSpecs.value();
+
+ CleanupTempTableDoFn(BigQueryServices bqServices) {
+ this.bqServices = bqServices;
+ }
+
+ @ProcessElement
+ public void processElement(
+ @Element KV<String, CleanupOperationMessage> element,
+ @StateId("cleanupInfo") ValueState<CleanupInfo> cleanupInfoState,
+ @StateId("completedStreams") ValueState<Integer> completedStreamsState,
+ PipelineOptions options)
+ throws Exception {
+
+ CleanupOperationMessage msg = element.getValue();
+ CleanupInfo cleanupInfo = cleanupInfoState.read();
+ int completed = firstNonNull(completedStreamsState.read(), 0);
+
+ if (msg.isStreamCompletion()) {
+ completed += 1;
+ completedStreamsState.write(completed);
+ } else {
+
cleanupInfoState.write(Preconditions.checkNotNull(msg.getCleanupInfo()));
+ cleanupInfo = msg.getCleanupInfo();
+ }
+
+ if (cleanupInfo != null && completed == cleanupInfo.getTotalStreams()) {
+ TableReference tempTable = cleanupInfo.getTableReference();
+ try (DatasetService datasetService =
+ bqServices.getDatasetService(options.as(BigQueryOptions.class))) {
+ LOG.info("Deleting temporary table with query results {}",
tempTable);
+ datasetService.deleteTable(tempTable);
+ if (cleanupInfo.isDatasetCreatedByBeam()) {
+ LOG.info("Deleting temporary dataset with query results {}",
tempTable.getDatasetId());
+ datasetService.deleteDataset(tempTable.getProjectId(),
tempTable.getDatasetId());
+ }
+ } catch (Exception e) {
+ LOG.warn("Failed to delete temporary BigQuery table {}", tempTable,
e);
+ }
+ cleanupInfoState.clear();
+ completedStreamsState.clear();
+ }
+ }
+ }
+
+ private static class ReadDynamicStreamSource<T>
+ extends DoFn<KV<String, BigQueryStorageStreamSource<T>>, T> {
+ private final TypedRead.ReadTableSource<T> readTableSource;
+ private final TupleTag<KV<String, CleanupOperationMessage>>
streamCompletionTag;
+
+ ReadDynamicStreamSource(
+ TupleTag<T> rowTag,
+ SerializableFunction<SchemaAndRecord, T> parseFn,
+ BadRecordRouter badRecordRouter,
+ TupleTag<KV<String, CleanupOperationMessage>> streamCompletionTag) {
+ this.readTableSource = new TypedRead.ReadTableSource<>(rowTag, parseFn,
badRecordRouter);
+ this.streamCompletionTag = streamCompletionTag;
+ }
+
+ @ProcessElement
+ public void processElement(
+ @Element KV<String, BigQueryStorageStreamSource<T>> element,
+ MultiOutputReceiver receiver,
+ PipelineOptions options)
+ throws Exception {
+ readTableSource.processElement(element.getValue(), receiver, options);
+ if (element.getValue().getFromQuery()) {
+ receiver
+ .get(streamCompletionTag)
+ .output(KV.of(element.getKey(),
CleanupOperationMessage.streamComplete()));
+ }
}
}
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageStreamSource.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageStreamSource.java
index 1bc0b2895d9..6376af7d816 100644
---
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageStreamSource.java
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageStreamSource.java
@@ -77,7 +77,13 @@ class BigQueryStorageStreamSource<T> extends
BoundedSource<T> {
toJsonString(Preconditions.checkArgumentNotNull(tableSchema,
"tableSchema")),
parseFn,
outputCoder,
- bqServices);
+ bqServices,
+ false);
+ }
+
+ public BigQueryStorageStreamSource<T> withFromQuery() {
+ return new BigQueryStorageStreamSource<>(
+ readSession, readStream, jsonTableSchema, parseFn, outputCoder,
bqServices, true);
}
@Override
@@ -106,13 +112,13 @@ class BigQueryStorageStreamSource<T> extends
BoundedSource<T> {
*/
public BigQueryStorageStreamSource<T> fromExisting(ReadStream newReadStream)
{
return new BigQueryStorageStreamSource<>(
- readSession, newReadStream, jsonTableSchema, parseFn, outputCoder,
bqServices);
+ readSession, newReadStream, jsonTableSchema, parseFn, outputCoder,
bqServices, fromQuery);
}
public BigQueryStorageStreamSource<T> fromExisting(
SerializableFunction<SchemaAndRecord, T> parseFn) {
return new BigQueryStorageStreamSource<>(
- readSession, readStream, jsonTableSchema, parseFn, outputCoder,
bqServices);
+ readSession, readStream, jsonTableSchema, parseFn, outputCoder,
bqServices, fromQuery);
}
private final ReadSession readSession;
@@ -121,6 +127,7 @@ class BigQueryStorageStreamSource<T> extends
BoundedSource<T> {
private final SerializableFunction<SchemaAndRecord, T> parseFn;
private final Coder<T> outputCoder;
private final BigQueryServices bqServices;
+ private final boolean fromQuery;
private BigQueryStorageStreamSource(
ReadSession readSession,
@@ -128,13 +135,19 @@ class BigQueryStorageStreamSource<T> extends
BoundedSource<T> {
String jsonTableSchema,
SerializableFunction<SchemaAndRecord, T> parseFn,
Coder<T> outputCoder,
- BigQueryServices bqServices) {
+ BigQueryServices bqServices,
+ boolean fromQuery) {
this.readSession = Preconditions.checkArgumentNotNull(readSession,
"readSession");
this.readStream = Preconditions.checkArgumentNotNull(readStream, "stream");
this.jsonTableSchema = Preconditions.checkArgumentNotNull(jsonTableSchema,
"jsonTableSchema");
this.parseFn = Preconditions.checkArgumentNotNull(parseFn, "parseFn");
this.outputCoder = Preconditions.checkArgumentNotNull(outputCoder,
"outputCoder");
this.bqServices = Preconditions.checkArgumentNotNull(bqServices,
"bqServices");
+ this.fromQuery = fromQuery;
+ }
+
+ public boolean getFromQuery() {
+ return fromQuery;
}
@Override
diff --git
a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/CleanupTempTableDoFnTest.java
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/CleanupTempTableDoFnTest.java
new file mode 100644
index 00000000000..b1fa3dc8e57
--- /dev/null
+++
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/CleanupTempTableDoFnTest.java
@@ -0,0 +1,153 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.gcp.bigquery;
+
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+import com.google.api.services.bigquery.model.TableReference;
+import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.CleanupInfo;
+import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.CleanupOperationMessage;
+import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.CleanupTempTableDoFn;
+import org.apache.beam.sdk.io.gcp.testing.FakeBigQueryServices;
+import org.apache.beam.sdk.io.gcp.testing.FakeDatasetService;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class CleanupTempTableDoFnTest {
+
+ @Rule public transient TestPipeline p = TestPipeline.create();
+
+ @Before
+ public void setUp() {
+ FakeDatasetService.setUp();
+ }
+
+ @Test
+ public void testCleanupTempTableDoFn() throws Exception {
+ FakeDatasetService fakeDatasetService = new FakeDatasetService();
+ FakeBigQueryServices fakeBqServices =
+ new FakeBigQueryServices().withDatasetService(fakeDatasetService);
+
+ TableReference tableRef =
+ new
TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
+
+ fakeDatasetService.createDataset(
+ tableRef.getProjectId(), tableRef.getDatasetId(), "", "", null);
+ fakeDatasetService.createTable(
+ new
com.google.api.services.bigquery.model.Table().setTableReference(tableRef));
+
+ assertTrue(
+ fakeDatasetService.getDataset(tableRef.getProjectId(),
tableRef.getDatasetId()) != null);
+ assertTrue(fakeDatasetService.getTable(tableRef) != null);
+
+ CleanupInfo cleanupInfo = new CleanupInfo(tableRef, true, 2);
+
+ PCollection<KV<String, CleanupOperationMessage>> input =
+ p.apply(
+ Create.of(
+ KV.of("job1", CleanupOperationMessage.initialize(cleanupInfo)),
+ KV.of("job1", CleanupOperationMessage.streamComplete()),
+ KV.of("job1", CleanupOperationMessage.streamComplete())));
+
+ input.apply(ParDo.of(new CleanupTempTableDoFn(fakeBqServices)));
+
+ p.run().waitUntilFinish();
+
+ // The dataset is deleted, so getDataset and getTable should throw a 404
Exception or return
+ // null
+ try {
+ fakeDatasetService.getDataset(tableRef.getProjectId(),
tableRef.getDatasetId());
+ fail("Dataset should have been deleted");
+ } catch (Exception e) {
+ assertTrue(
+ e.getMessage().contains("Tried to get a dataset")
+ || e.getMessage().contains("Not Found"));
+ }
+
+ try {
+ TableReference deletedRef =
+ new
TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
+ Object table = fakeDatasetService.getTable(deletedRef);
+ assertTrue(table == null);
+ } catch (Exception e) {
+ assertTrue(
+ e.getMessage().contains("Tried to get a dataset")
+ || e.getMessage().contains("Not Found"));
+ }
+ }
+
+ @Test
+ public void testCleanupTempTableDoFnWithZeroStreams() throws Exception {
+ FakeDatasetService fakeDatasetService = new FakeDatasetService();
+ FakeBigQueryServices fakeBqServices =
+ new FakeBigQueryServices().withDatasetService(fakeDatasetService);
+
+ TableReference tableRef =
+ new
TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
+
+ fakeDatasetService.createDataset(
+ tableRef.getProjectId(), tableRef.getDatasetId(), "", "", null);
+ fakeDatasetService.createTable(
+ new
com.google.api.services.bigquery.model.Table().setTableReference(tableRef));
+
+ assertTrue(
+ fakeDatasetService.getDataset(tableRef.getProjectId(),
tableRef.getDatasetId()) != null);
+ assertTrue(fakeDatasetService.getTable(tableRef) != null);
+
+ CleanupInfo cleanupInfo = new CleanupInfo(tableRef, true, 0);
+
+ PCollection<KV<String, CleanupOperationMessage>> input =
+ p.apply(
+ "CreateInputWithZeroStreams",
+ Create.of(KV.of("job1",
CleanupOperationMessage.initialize(cleanupInfo))));
+
+ input.apply("CleanupWithZeroStreams", ParDo.of(new
CleanupTempTableDoFn(fakeBqServices)));
+
+ p.run().waitUntilFinish();
+
+ try {
+ fakeDatasetService.getDataset(tableRef.getProjectId(),
tableRef.getDatasetId());
+ fail("Dataset should have been deleted");
+ } catch (Exception e) {
+ assertTrue(
+ e.getMessage().contains("Tried to get a dataset")
+ || e.getMessage().contains("Not Found"));
+ }
+
+ try {
+ TableReference deletedRef =
+ new
TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
+ Object table = fakeDatasetService.getTable(deletedRef);
+ assertTrue(table == null);
+ } catch (Exception e) {
+ assertTrue(
+ e.getMessage().contains("Tried to get a dataset")
+ || e.getMessage().contains("Not Found"));
+ }
+ }
+}
