[shunping]

Part of #40569

Summary

BigQueryIO.read() with the default EXPORT method runs a cleanup step after the read that looks up the extract job to find and delete the exported Avro files. That lookup is built without a job location. BigQuery requires the location to look up jobs in regional locations, so for tables in regional datasets the lookup fails. Depending on the endpoint, cleanup is either silently skipped (leaking export files) or fails repeatedly, which stalls streaming pipelines indefinitely.

This PR sets the job location on that lookup, the same way the export job itself does.

What happens today

The export job is created with a location (BigQuerySourceBase#executeExtract):

new JobReference().setProjectId(executingProject).setLocation(bqLocation).setJobId(jobId);

But the cleanup in TypedRead#expand rebuilds the reference without one:

JobReference jobRef = new JobReference().setProjectId(executingProject).setJobId(...); // no location
Job extractJob = getBigQueryServices().getJobService(bqOptions).getJob(jobRef);

The resulting jobs.get has no location parameter. For a job in a regional location:

• Google Cloud (default universe): BigQuery returns 404. getJob treats that as "not found" and returns null, so cleanup is silently skipped and the exported files under <tempLocation>/BigQueryExtractTemp/ are leaked. This is easy to miss, because datasets in the US/EU multi-regions work without a location.
• Trusted Partner Cloud (TPC): BigQuery returns 400 INVALID_ARGUMENT instead. getJob doesn't treat that as "not found", retries, and throws. In a streaming pipeline the failed bundle is retried forever, so the job never makes progress.
  Observed on TPC (streaming pipeline reading a table in u-us-prp1 as a side input):

java.io.IOException: Unable to find BigQuery job: GenericData{classInfo=[jobId, location, projectId],
  {jobId=beam_bq_job_EXPORT_..., projectId=tpczero-system:dataflow-tpc-test}}, aborting after 9 retries.
    at org.apache.beam.sdk.io.gcp.bigquery.BigQueryServicesImpl$JobServiceImpl.getJob(BigQueryServicesImpl.java:598)
    at org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO$TypedRead$4.cleanup(BigQueryIO.java:1940)
    at org.apache.beam.sdk.io.gcp.bigquery.PassThroughThenCleanup$1.processElement(PassThroughThenCleanup.java:73)
Caused by: com.google.api.client.googleapis.json.GoogleJsonResponseException: 400 Bad Request
GET https://bigquery.apis-tpczero.goog/bigquery/v2/projects/tpczero-system:dataflow-tpc-test/jobs/beam_bq_job_EXPORT_...
{ "code": 400, "message": "Error: 3848323", "status": "INVALID_ARGUMENT" }

The fix

When building the cleanup's job reference, set the location the export job ran in:

• Table reads: the source table's dataset location (BigQueryHelpers.getDatasetLocation), the same value BigQuerySourceBase uses for the export.
• Query reads: the withQueryLocation value.
  The location is only set when known. Query reads without an explicit location behave as before. Per the withQueryLocation Javadoc, a location must be specified for queries not executed in US/EU, and those multi-regions work without one.

[github-actions]

📋 Review Summary

This PR fixes a bug in BigQueryIO export reads where the post-read cleanup step looked up the extract job without specifying a location, causing either leaked temporary Avro files (on 404) or infinite retry loops (on 400 INVALID_ARGUMENT in regional endpoints). However, for query-based export reads, the cleanup step does not follow the same location resolution logic used when creating the query/temp dataset and extract job, leaving regional queries that infer location from withQueryTempDataset or dry-run referenced tables broken.

🔍 General Feedback

• The table read path in getExtractJobLocation accurately mirrors BigQuerySourceBase#split by looking up the source table's dataset location.
• For query reads (BigQueryQuerySourceDef), BigQueryQueryHelper.executeQuery resolves the location from three possible sources (configured queryTempDataset, dry-run referenced tables, or queryLocation) and creates the temporary dataset and destination table in that location. BigQuerySourceBase#split then launches the extract job in that temp dataset's location. Relying solely on getQueryLocation() in getExtractJobLocation misses the first two cases.

---

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.

[shunping]

/gemini review

[github-actions]

📋 Review Summary

This PR properly resolves and sets the BigQuery job location when looking up the extract job during post-read cleanup of BigQueryIO EXPORT reads. It handles both table reads (using the source table's dataset location) and query reads (via a shared BigQueryQueryHelper.resolveQueryLocation helper and tempDatasetId fallback), and includes comprehensive unit tests for both standard and template-compatible execution paths.

🔍 General Feedback

• The extracted BigQueryQueryHelper.resolveQueryLocation cleanly deduplicates location resolution between executeQuery and BigQueryQuerySourceDef.getQueryLocation without re-executing queries or leaking resources.
• Regression tests in BigQueryIOReadTest effectively cover both table and query export reads with and without template compatibility.

---

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.

[stankiewicz]

great find!

[github-actions]

Assigning reviewers:

R: @chamikaramj for label java.
Backup reviewer: Abacn (if the main reviewer is unresponsive, you can escalate to them by tagging their username)

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

• stop reviewer notifications - opt out of the automated review tooling
• remind me after tests pass - tag the comment author after tests pass
• waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).