[Java] Set job location when cleaning up BigQueryIO export reads

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));

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<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 @@ 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<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

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.
+   *
+   * <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;

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<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