This is an automated email from the ASF dual-hosted git repository.
shunping 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 1e5058c24ee [Java] Set job location when cleaning up BigQueryIO export
reads (#40794)
1e5058c24ee is described below
commit 1e5058c24ee48aadbed23dcc9b7a96988470eefc
Author: Shunping Huang <[email protected]>
AuthorDate: Mon Oct 26 09:53:03 2026 -0400
[Java] Set job location when cleaning up BigQueryIO export reads (#40794)
The export-read cleanup looked up the extract job without a location,
which fails for regional datasets: cleanup is skipped (404) or fails
repeatedly, stalling streaming jobs (400, e.g. on TPC).
For table reads, use the source table's dataset location. For query
reads, resolve the location the same way executeQuery does (explicit
location, else the first referenced table's dataset), via a shared
helper that doesn't rerun the query. Fall back to a user-provided
temp dataset's location. Add regression tests.
---
.../beam/sdk/io/gcp/bigquery/BigQueryIO.java | 28 ++++
.../sdk/io/gcp/bigquery/BigQueryQueryHelper.java | 63 +++++---
.../io/gcp/bigquery/BigQueryQuerySourceDef.java | 34 ++++
.../sdk/io/gcp/bigquery/BigQueryIOReadTest.java | 179 +++++++++++++++++++++
4 files changed, 283 insertions(+), 21 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 07fa1a320e9..0b8263ecb1e 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
@@ -1936,6 +1936,12 @@ public class BigQueryIO {
.setJobId(
BigQueryResourceNaming.createJobIdPrefix(
bqOptions.getJobName(), jobUuid,
JobType.EXPORT));
+ // Job lookups must specify the location for regional locations;
without it the
+ // job isn't found, or some endpoints reject the request.
+ String location = getExtractJobLocation(bqOptions);
+ if (location != null) {
+ jobRef.setLocation(location);
+ }
Job extractJob =
getBigQueryServices().getJobService(bqOptions).getJob(jobRef);
@@ -1948,6 +1954,28 @@ public class BigQueryIO {
}
}
}
+
+ /**
+ * Returns the location the extract job ran in, which job lookups
must specify for
+ * regional locations: the source table's dataset location for
table reads, or the query
+ * location for query reads, resolved as {@link
BigQueryQueryHelper#executeQuery} does
+ * (explicit location, or inferred from the referenced tables).
+ */
+ private @Nullable String getExtractJobLocation(BigQueryOptions
bqOptions)
+ throws Exception {
+ if (sourceDef instanceof BigQueryTableSourceDef) {
+ TableReference tableRef =
+ ((BigQueryTableSourceDef)
sourceDef).getTableReference(bqOptions);
+ try (DatasetService datasetService =
+ getBigQueryServices().getDatasetService(bqOptions)) {
+ return BigQueryHelpers.getDatasetLocation(
+ datasetService, tableRef.getProjectId(),
tableRef.getDatasetId());
+ }
+ } else if (sourceDef instanceof BigQueryQuerySourceDef) {
+ return ((BigQueryQuerySourceDef)
sourceDef).getQueryLocation(bqOptions);
+ }
+ return null;
+ }
};
rows = rows.apply(new PassThroughThenCleanup<>(cleanupOperation,
jobIdTokenView));
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryQueryHelper.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryQueryHelper.java
index f5334a41cdc..b052f383f4b 100644
---
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryQueryHelper.java
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryQueryHelper.java
@@ -81,6 +81,38 @@ class BigQueryQueryHelper {
}
}
+ /**
+ * Returns the location the query runs in: the explicitly configured {@code
location} if set,
+ * otherwise the location of the first table referenced by the query
(determined via a dry run),
+ * or null if neither is available.
+ */
+ static @Nullable String resolveQueryLocation(
+ BigQueryServices bqServices,
+ DatasetService tableService,
+ BigQueryOptions options,
+ AtomicReference<@Nullable JobStatistics> dryRunJobStats,
+ String query,
+ Boolean flattenResults,
+ Boolean useLegacySql,
+ @Nullable String location)
+ throws InterruptedException, IOException {
+ if (location != null) {
+ return location;
+ }
+ List<TableReference> referencedTables =
+ dryRunQueryIfNeeded(
+ bqServices, options, dryRunJobStats, query, flattenResults,
useLegacySql, location)
+ .getQuery()
+ .getReferencedTables();
+ if (referencedTables != null && !referencedTables.isEmpty()) {
+ TableReference referencedTable = referencedTables.get(0);
+ return tableService
+ .getDataset(referencedTable.getProjectId(),
referencedTable.getDatasetId())
+ .getLocation();
+ }
+ return null;
+ }
+
public static TableReference executeQuery(
BigQueryServices bqServices,
BigQueryOptions options,
@@ -96,28 +128,17 @@ class BigQueryQueryHelper {
@Nullable String kmsKey)
throws InterruptedException, IOException {
// Step 1: Find the effective location of the query.
- @Nullable String effectiveLocation = location;
try (DatasetService tableService = bqServices.getDatasetService(options)) {
- if (effectiveLocation == null) {
- List<TableReference> referencedTables =
- dryRunQueryIfNeeded(
- bqServices,
- options,
- dryRunJobStats,
- query,
- flattenResults,
- useLegacySql,
- location)
- .getQuery()
- .getReferencedTables();
- if (referencedTables != null && !referencedTables.isEmpty()) {
- TableReference referencedTable = referencedTables.get(0);
- effectiveLocation =
- tableService
- .getDataset(referencedTable.getProjectId(),
referencedTable.getDatasetId())
- .getLocation();
- }
- }
+ @Nullable String effectiveLocation =
+ resolveQueryLocation(
+ bqServices,
+ tableService,
+ options,
+ dryRunJobStats,
+ query,
+ flattenResults,
+ useLegacySql,
+ location);
// Step 2: Create a temporary dataset in the query location only if the
user has not specified
// a temp dataset. The temp table name may be deterministic but the
query job ID has to be
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryQuerySourceDef.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryQuerySourceDef.java
index 29a38f004fc..591deef9f68 100644
---
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryQuerySourceDef.java
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryQuerySourceDef.java
@@ -136,6 +136,40 @@ class BigQueryQuerySourceDef implements BigQuerySourceDef {
kmsKey);
}
+ /**
+ * Returns the location of the query's result table, which is also where the
export job runs.
+ *
+ * <p>Uses the same resolution as {@link BigQueryQueryHelper#executeQuery}
(explicit location, or
+ * the location of the first referenced table). If that yields nothing and
the user supplied a
+ * temp dataset, falls back to that dataset's location. Does not run the
query or create any
+ * resources.
+ */
+ @Nullable String getQueryLocation(BigQueryOptions bqOptions) throws
Exception {
+ try (BigQueryServices.DatasetService datasetService =
bqServices.getDatasetService(bqOptions)) {
+ @Nullable String resolved =
+ BigQueryQueryHelper.resolveQueryLocation(
+ bqServices,
+ datasetService,
+ bqOptions,
+ dryRunJobStats,
+ query.get(),
+ flattenResults,
+ useLegacySql,
+ location);
+ if (resolved == null && tempDatasetId != null) {
+ String project = tempProjectId;
+ if (project == null) {
+ project =
+ bqOptions.getBigQueryProject() == null
+ ? bqOptions.getProject()
+ : bqOptions.getBigQueryProject();
+ }
+ resolved = datasetService.getDataset(project,
tempDatasetId).getLocation();
+ }
+ return resolved;
+ }
+ }
+
void cleanupTempResource(BigQueryOptions bqOptions, String stepUuid) throws
Exception {
Optional<String> queryTempDatasetOpt = Optional.ofNullable(tempDatasetId);
String project = tempProjectId;
diff --git
a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOReadTest.java
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOReadTest.java
index 58d6a7c5d88..47a83b85d49 100644
---
a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOReadTest.java
+++
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOReadTest.java
@@ -25,6 +25,9 @@ import static org.hamcrest.Matchers.containsInAnyOrder;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNull;
+import com.google.api.services.bigquery.model.Dataset;
+import com.google.api.services.bigquery.model.Job;
+import com.google.api.services.bigquery.model.JobReference;
import com.google.api.services.bigquery.model.JobStatistics;
import com.google.api.services.bigquery.model.JobStatistics2;
import com.google.api.services.bigquery.model.Streamingbuffer;
@@ -40,11 +43,15 @@ import java.io.Serializable;
import java.math.BigDecimal;
import java.math.BigInteger;
import java.nio.ByteBuffer;
+import java.nio.file.Files;
+import java.nio.file.Path;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.ExecutionException;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
import org.apache.avro.specific.SpecificDatumReader;
import org.apache.avro.specific.SpecificRecordBase;
import org.apache.beam.sdk.PipelineResult;
@@ -72,6 +79,7 @@ import org.apache.beam.sdk.testing.SourceTestUtils;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.MapElements;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.SerializableFunction;
@@ -81,6 +89,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.Row;
+import org.apache.beam.sdk.values.TypeDescriptors;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
@@ -591,6 +600,176 @@ public class BigQueryIOReadTest implements Serializable {
checkLineageSourceMetric(result,
"non-executing-project:somedataset.sometable");
}
+ private static final String REGIONAL_LOCATION = "us-east7";
+
+ /** A dataset service that reports every dataset in {@link
#REGIONAL_LOCATION}. */
+ private static class RegionalDatasetService extends FakeDatasetService {
+ @Override
+ public Dataset getDataset(String projectId, String datasetId)
+ throws IOException, InterruptedException {
+ return super.getDataset(projectId,
datasetId).setLocation(REGIONAL_LOCATION);
+ }
+ }
+
+ /**
+ * A job service whose job lookups fail unless they specify the job's
location ({@link
+ * #REGIONAL_LOCATION}), like BigQuery does for jobs in regional locations
on some endpoints.
+ */
+ private static class LocationRequiredJobService extends FakeJobService {
+ @Override
+ public Job getJob(JobReference jobRef) {
+ if (!REGIONAL_LOCATION.equals(jobRef.getLocation())) {
+ throw new RuntimeException(
+ String.format(
+ "Expected location %s to look up %s, got %s",
+ REGIONAL_LOCATION, jobRef.getJobId(), jobRef.getLocation()));
+ }
+ return super.getJob(jobRef);
+ }
+ }
+
+ @Test
+ public void
testExportReadCleanupWithRegionalDatasetWithoutTemplateCompatibility()
+ throws IOException, InterruptedException {
+ testExportReadCleanupWithRegionalDataset(false);
+ }
+
+ @Test
+ public void
testExportReadCleanupWithRegionalDatasetWithTemplateCompatibility()
+ throws IOException, InterruptedException {
+ testExportReadCleanupWithRegionalDataset(true);
+ }
+
+ private void testExportReadCleanupWithRegionalDataset(boolean
useTemplateCompatibility)
+ throws IOException, InterruptedException {
+ TableReference tableRef =
+ new TableReference()
+ .setProjectId("project-id")
+ .setDatasetId("regional_dataset")
+ .setTableId("sometable");
+ FakeDatasetService fakeDatasetService = new RegionalDatasetService();
+ fakeDatasetService.createDataset("project-id", "regional_dataset",
REGIONAL_LOCATION, "", null);
+ fakeDatasetService.createTable(
+ new Table()
+ .setTableReference(tableRef)
+ .setSchema(
+ new TableSchema()
+ .setFields(
+ ImmutableList.of(
+ new
TableFieldSchema().setName("name").setType("STRING")))));
+ fakeDatasetService.insertAll(
+ tableRef,
+ ImmutableList.of(new TableRow().set("name", "a"), new
TableRow().set("name", "b")),
+ null);
+
+ FakeBigQueryServices fakeBqServices =
+ new FakeBigQueryServices()
+ .withJobService(new LocationRequiredJobService())
+ .withDatasetService(fakeDatasetService);
+ BigQueryIO.TypedRead<TableRow> read =
+ BigQueryIO.readTableRows()
+ .from(tableRef)
+ .withTestServices(fakeBqServices)
+ .withoutValidation();
+ if (useTemplateCompatibility) {
+ read = read.withTemplateCompatibility();
+ }
+
+ PCollection<String> names =
+ p.apply(read)
+ .apply(
+ MapElements.into(TypeDescriptors.strings())
+ .via((TableRow row) -> (String) row.get("name")));
+ PAssert.that(names).containsInAnyOrder("a", "b");
+ p.run().waitUntilFinish();
+
+ assertNoExtractFilesRemain();
+ }
+
+ @Test
+ public void
testExportQueryReadCleanupWithInferredRegionalLocationWithoutTemplateCompatibility()
+ throws Exception {
+ testExportQueryReadCleanupWithInferredRegionalLocation(false);
+ }
+
+ @Test
+ public void
testExportQueryReadCleanupWithInferredRegionalLocationWithTemplateCompatibility()
+ throws Exception {
+ testExportQueryReadCleanupWithInferredRegionalLocation(true);
+ }
+
+ private void testExportQueryReadCleanupWithInferredRegionalLocation(
+ boolean useTemplateCompatibility) throws Exception {
+ TableReference sourceTableRef =
+ new TableReference()
+ .setProjectId("project-id")
+ .setDatasetId("regional_dataset")
+ .setTableId("sometable");
+ FakeDatasetService fakeDatasetService = new RegionalDatasetService();
+ fakeDatasetService.createDataset("project-id", "regional_dataset",
REGIONAL_LOCATION, "", null);
+ fakeDatasetService.createTable(new
Table().setTableReference(sourceTableRef));
+
+ Table queryResultTable =
+ new Table()
+ .setSchema(
+ new TableSchema()
+ .setFields(
+ ImmutableList.of(
+ new
TableFieldSchema().setName("name").setType("STRING"))));
+ String encodedQuery =
+ FakeBigQueryServices.encodeQueryResult(
+ queryResultTable,
+ ImmutableList.of(new TableRow().set("name", "a"), new
TableRow().set("name", "b")));
+
+ // No withQueryLocation(): the location is inferred from the referenced
table.
+ FakeJobService fakeJobService = new LocationRequiredJobService();
+ fakeJobService.expectDryRunQuery(
+ "project-id",
+ encodedQuery,
+ new JobStatistics()
+ .setQuery(
+ new JobStatistics2()
+ .setTotalBytesProcessed(100L)
+ .setReferencedTables(ImmutableList.of(sourceTableRef))));
+
+ FakeBigQueryServices fakeBqServices =
+ new FakeBigQueryServices()
+ .withJobService(fakeJobService)
+ .withDatasetService(fakeDatasetService);
+
+ BigQueryIO.TypedRead<TableRow> read =
+ BigQueryIO.readTableRows()
+ .fromQuery(encodedQuery)
+ .withTestServices(fakeBqServices)
+ .withoutValidation();
+ if (useTemplateCompatibility) {
+ read = read.withTemplateCompatibility();
+ }
+
+ PCollection<String> names =
+ p.apply(read)
+ .apply(
+ MapElements.into(TypeDescriptors.strings())
+ .via((TableRow row) -> (String) row.get("name")));
+ PAssert.that(names).containsInAnyOrder("a", "b");
+ p.run().waitUntilFinish();
+
+ assertNoExtractFilesRemain();
+ }
+
+ /** The export writes Avro files under the temp location; cleanup must have
deleted them. */
+ private void assertNoExtractFilesRemain() throws IOException {
+ List<String> remainingExtractFiles;
+ try (Stream<Path> paths = Files.walk(testFolder.getRoot().toPath())) {
+ remainingExtractFiles =
+ paths
+ .map(Path::toString)
+ .filter(path -> path.endsWith(".avro"))
+ .collect(Collectors.toList());
+ }
+ assertEquals(ImmutableList.of(), remainingExtractFiles);
+ }
+
@Test
public void testReadTableWithSchema() throws IOException,
InterruptedException {
// setup
