From 83dba6210dfa048cbc8b3d8c4e100443fc13b411 Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Sat, 24 Oct 2026 23:27:50 -0400 Subject: [PATCH] [Java] Set job location when cleaning up BigQueryIO export reads 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 +++ .../io/gcp/bigquery/BigQueryQueryHelper.java | 63 ++++-- .../gcp/bigquery/BigQueryQuerySourceDef.java | 34 ++++ .../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 07fa1a320e92..0b8263ecb1e4 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 @@ void cleanup(PassThroughThenCleanup.ContextContainer c) throws Exception { .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 @@ void cleanup(PassThroughThenCleanup.ContextContainer c) throws Exception { } } } + + /** + * 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 f5334a41cdc1..b052f383f4ba 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 @@ public static JobStatistics dryRunQueryIfNeeded( } } + /** + * 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 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 @@ public static TableReference executeQuery( @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 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 29a38f004fce..591deef9f680 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 @@ TableReference getTableReference(BigQueryOptions bqOptions, String stepUuid) kmsKey); } + /** + * Returns the location of the query's result table, which is also where the export job runs. + * + *

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 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 58d6a7c5d88e..47a83b85d497 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.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.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.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.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 void processElement(ProcessContext c) throws Exception { 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 read = + BigQueryIO.readTableRows() + .from(tableRef) + .withTestServices(fakeBqServices) + .withoutValidation(); + if (useTemplateCompatibility) { + read = read.withTemplateCompatibility(); + } + + PCollection 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 read = + BigQueryIO.readTableRows() + .fromQuery(encodedQuery) + .withTestServices(fakeBqServices) + .withoutValidation(); + if (useTemplateCompatibility) { + read = read.withTemplateCompatibility(); + } + + PCollection 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 remainingExtractFiles; + try (Stream 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