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