[JeremyMcCormick]

Summary

Replaces the Dataflow job that staged replica chunks with BigQuery load jobs run from a new ChunkStager. The stage_chunk Cloud Run function uses the stager (lsst-dm/ppdb-cloud-functions#38).

Changes

• ChunkStager (bigquery/chunk_stager.py): stages one chunk at a time with stage_chunk(logger, chunk_id, folder)
  • Accepts only a chunk whose status is uploaded or retrying, and raises ChunkStatusError otherwise (or ChunkStagingError if the chunk doesn't exist); it then sets the chunk to staging
  • Reads the chunk's manifest and checks that it is for the same chunk
  • Checks that the chunk's rows are not already in the staging tables (_check_not_already_staged)
  • Loads each Parquet file into its own temporary landing table, with the staging table's schema (_copy_to_landing)
  • Copies the landing tables into the staging tables in one transaction, setting geo_point and checking each table's row count against the manifest with an ASSERT (_copy_to_staging)
  • Records the ID of the copy job, then sets the chunk to staged
  • No raw errors leave stage_chunk; it raises only ChunkStatusError, ChunkStagingError or ChunkStagingRetryableError
    • A database error before the chunk is set to staging leaves its status unchanged; it raises ChunkStagingRetryableError if handle_database_error classifies it as temporary, and ChunkStagingError otherwise
    • While staging, temporary Google Cloud and network errors, and BigQuery quota and rate-limit errors (quotaExceeded, rateLimitExceeded), set the chunk to retrying and raise ChunkStagingRetryableError; any other error, including an unexpected one, sets it to failed and raises ChunkStagingError
    • If waiting for the staging transaction fails with a temporary error, the job is cancelled and the chunk set to failed, so a retry can't start a second transaction and stage the chunk twice
    • A database error while setting the final status raises ChunkStagingError; the chunk is left in staging
  • Each BigQuery API request is retried after a temporary error for up to retry_timeout, with the client's default delays; failed jobs are never resubmitted (job_retry=None), so the recorded job ID is the one that ran
  • Each BigQuery job has a time limit (check_timeout, job_timeout)
  • Created once with the PPDB, a BigQuery client and the limits, and reused for every chunk
• Chunk status (bigquery/ppdb_replica_chunk_extended.py)
  • New retrying status, and new ChunkStatusError
  • New status_time, error_message and staging_job_id fields
  • New with_new_error_message and with_new_staging_job_id, like the existing with_new_status and with_new_gcs_uri
  • UpdatableField moved here from ppdb_bigquery.py
• Chunk table (PpdbBigQuery)
  • New status_time, error_message and staging_job_id columns
  • update_chunks sets status_time when it changes the status
• Datasets
  • New landing dataset (DatasetType.LANDING, TableRefs.landing). It has no tables of its own, so it is created by Terraform, not by the dataset builder
  • The staging DIA tables get the geo_point column and are clustered by apdb_replica_chunk
  • ChunkPromoter copies geo_point from staging instead of computing it
• Uploads (ChunkUploader, PpdbBigQueryConfig)
  • Chunks are uploaded to the top of the chunk bucket; the object_prefix setting and --object-prefix option are removed
  • The unused dataset field is dropped from the stage chunk message
• Database errors (gcp/error_handling.py)
  • New DATABASE_ERRORS: SQLAlchemyError plus the authentication, aiohttp and network errors that the Cloud SQL Python Connector can raise when opening a connection, which SQLAlchemy doesn't wrap
  • handle_database_error classifies all of them: dropped connections, network errors, token endpoint failures and 429/5xx Admin API responses are retryable; other errors are not
  • aiohttp is now a listed dependency
• Other
  • Removed stale references to the Dataflow staging job
• Tests
  • tests/test_chunk_stager.py: integration test for ChunkStager, and pickling of its exceptions
  • Test datasets are cleaned up even if their creation fails partway, and the dataset builder test no longer risks deleting the real landing dataset

Deployment notes

• Existing chunk tables need the status_time, error_message and staging_job_id columns
• Existing staging tables need the geo_point column
• A chunk left in staging by a staging attempt that died has to be reset by hand

Related

• Jira: https://rubinobs.atlassian.net/browse/DM-56864
• ppdb-cloud-functions: DM-56864: Stage chunks with BigQuery load jobs instead of Dataflow lsst-dm/ppdb-cloud-functions#38

🤖 Generated with Claude Code

https://claude.ai/code/session_01VJ51nAbGJ4BBgwpjau1VX7

[codecov]

Codecov Report

❌ Patch coverage is 79.27928% with 46 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.90%. Comparing base (788f62c) to head (131243d).
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
python/lsst/dax/ppdb/bigquery/chunk_stager.py	77.85%	28 Missing and 3 partials ⚠️
python/lsst/dax/ppdb/gcp/error_handling.py	40.90%	13 Missing ⚠️
python/lsst/dax/ppdb/bigquery/ppdb_bigquery.py	81.81%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #69      +/-   ##
==========================================
+ Coverage   80.83%   80.90%   +0.06%     
==========================================
  Files          54       55       +1     
  Lines        2812     3000     +188     
  Branches      258      284      +26     
==========================================
+ Hits         2273     2427     +154     
- Misses        448      479      +31     
- Partials       91       94       +3     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

[JeremyMcCormick]

Review

Replacing the Dataflow job with load jobs and a single transaction is a clear improvement. The design holds up well. The conditional claim stops concurrent deliveries, and the ASSERT row count check inside the transaction rejects a duplicate insert even if the pre-check is bypassed. Job resubmission is also turned off throughout. The error handling has been tightened a lot in response to earlier comments. The main remaining risk is the "unknown outcome" logic around the staging transaction: one path still treats a committed transaction as a rolled-back one. A second risk is that the classification branches are complex and have almost no direct test coverage.

Findings

1. A committed transaction can be recorded as retrying and then failed (python/lsst/dax/ppdb/bigquery/chunk_stager.py:506). job.done(reload=False) is true whenever the job state is DONE, including a successful job. In google-cloud-bigquery 3.38, QueryJob.result() still calls _reload_query_results (jobs.getQueryResults) after jobs.get reports DONE with no error, and that call can fail with a network error, a 5xx or a RetryError once retry_timeout runs out. The error is then re-raised as "the job's own", the outer handler sets the chunk to retrying, and the next delivery finds the rows in _check_not_already_staged and sets the chunk to failed. Data isn't duplicated, but a correctly staged chunk ends up failed and is never promoted. Suggested fix: when the job is done, check job.error_result is None (or job.exception() is None). If there is no job error, treat the transaction as committed and return job.job_id instead of re-raising.

2. Rate-limit exhaustion at submission fails the chunk instead of retrying it (python/lsst/dax/ppdb/bigquery/chunk_stager.py:481). DEFAULT_RETRY retries rateLimitExceeded and backendError on jobs.insert. Once retry_timeout runs out it raises RetryError, which is in NETWORK_ERRORS. A burst of 429/rateLimitExceeded responses, where BigQuery refused the request and created no job, therefore ends as a permanent failed chunk. That conflicts with the stated intent that 429 and quota errors are still retried. Suggested fix: inspect e.cause on a RetryError. If the last error was a 429 or a rateLimitExceeded/quotaExceeded reason, let it reach the retryable path. Keep failed only for 5xx or network causes.

3. Cloud SQL Connector exceptions escape stage_chunk unwrapped (python/lsst/dax/ppdb/gcp/error_handling.py:46, python/lsst/dax/ppdb/bigquery/chunk_stager.py:190, :243). google.cloud.sql.connector.exceptions defines DnsResolutionError, RefreshNotValidError, CloudSQLIPTypeError, ClosedConnectorError and others, all subclassing Exception directly. SQLAlchemy doesn't wrap errors raised by a creator function, so these aren't in DATABASE_ERRORS. If one is raised while claiming the chunk or recording its final status, it breaks the "only ChunkStatusError, ChunkStagingError or ChunkStagingRetryableError leave stage_chunk" contract. The caller then decides retry or ack on an unclassified error, and a transient DNS failure isn't retried as intended. Suggested fix: add the connector's exception classes to DATABASE_ERRORS, treat DnsResolutionError and RefreshNotValidError as retryable in handle_database_error, and treat the rest as not retryable.

4. A stale staging_job_id remains after a new claim (python/lsst/dax/ppdb/bigquery/chunk_stager.py:180). The claim clears error_message but not staging_job_id. A retrying chunk whose next attempt fails before the transaction is submitted (bad manifest, load job error) keeps the job ID from the earlier attempt. That contradicts the field's documented meaning, "the last staging attempt", and could mislead anyone checking job history during a manual reset. Suggested fix: also clear it in the claim with .with_new_staging_job_id(None) and UpdatableField.STAGING_JOB_ID.

5. The riskiest error paths have no tests (tests/test_chunk_stager.py, python/lsst/dax/ppdb/gcp/error_handling.py). The only stager test is a GCP integration test covering the success path, a duplicate claim, rows already staged, a missing manifest and a failed ASSERT. Nothing covers the cancel-on-unknown-outcome branch, an ambiguous submission, quota and rate-limit classification, _record_failure raising, or the new handle_database_error branches (OSError, GoogleAuthError.retryable, aiohttp.ClientResponseError status codes). These branches are what prevent double staging, and finding 1 is the kind of regression such tests would catch. Suggested fix: add unit tests with a mocked bigquery.Client, QueryJob and PpdbBigQuery that drive each branch and check the resulting chunk status and exception type. Add a table-driven test for handle_database_error.

🤖 Generated with Claude Code

[JeremyMcCormick]

@andy-slac @fajpunk I resolved the issues raised in the Copilot review.

I'm still working on improving the tests in tests/test_chunk_stager.py so they're more understandable and cover the important control and error paths.

[JeremyMcCormick]

Reply to the review above:

1. Committed transaction recorded as retrying, then failed. Fixed in 019e22a. If waiting for the staging job fails after the job finished with no error_result, the transaction committed, so the job ID is returned and the chunk is set to staged. The error is only logged (staging_result_not_fetched).
2. Rate-limit exhaustion at submission fails the chunk. Fixed in 6002e50. If the RetryError from client.query was caused by a 429 or rateLimitExceeded, that cause is raised, so the chunk is set to retrying. Any other RetryError, 5xx or network error still fails the chunk. An earlier attempt could in principle have created the job with its response lost. The next attempt's check for rows already in staging catches a job that has committed, and the docstring notes this.
3. Cloud SQL Connector errors unclassified. Fixed in a7fb473. The connector's exceptions are in DATABASE_ERRORS. DnsResolutionError and RefreshNotValidError are retryable, and the configuration and closed-connector errors are not.
4. Stale staging_job_id after a new claim. Fixed in 256f98c. The claim clears staging_job_id along with error_message.
5. Error paths have no tests. Partly addressed in f87831a and the fixups after it. ChunkStagerErrorTestCase adds three tests that need no credentials and check a chunk's final status: a staging job that may still be running (failed, job cancelled), a job that committed but whose result couldn't be read (staged), and a temporary error before the transaction (retrying).

Notes on the tests

• These error paths can't be triggered on purpose against real BigQuery: a lost response, or a 503 at a particular point in staging. The new tests therefore use a mocked BigQuery client and jobs, and a mocked PPDB. They check ChunkStager's own decisions, but they rely on assumptions about how the BigQuery client behaves, so they can't catch a wrong assumption about the library.
• handle_database_error still has no direct tests.
• ChunkStagerTestCase.test_stage_chunks runs several scenarios in sequence: staging every chunk, re-staging a staged chunk, rows already in staging, a missing manifest, and a manifest row count that fails the ASSERT. Each scenario needs the BigQuery datasets in a different state, and each separate test would rerun the full setup (filling the APDB, replicating it, creating the datasets, uploading the chunks and creating the staging tables), so splitting it would make the test too expensive. 5516a8b notes this on the method.

Possible follow-up work

• Move the scenarios that don't depend on BigQuery (re-staging a staged chunk, rows already in staging, a missing manifest) into the mocked tests, and split the remaining integration test into independent tests.
• Share the expensive BigQuery setup across integration tests, so they can be split without repeating it.
• Add direct tests of handle_database_error.

🤖 Generated with Claude Code

[JeremyMcCormick]

@andy-slac @fajpunk I think the tests are good enough for now after I made a few improvements. I won't make any other changes until there is a review.

There's still a lot of uncovered code, but so much mocking would be required to test those paths that I don't think it is really worth it.

[fajpunk]

This is great, I love getting rid of the dataflow and associated Google Cloud resources.

I admit that I get lost in the error handling

[JeremyMcCormick]

This is great, I love getting rid of the dataflow and associated Google Cloud resources.

I admit that I get lost in the error handling

I've made my best effort in the error handling to:

• Distinguish between re-tryable and fatal errors and rethrow a specific module-level exception for each
• Catch all the different error types that can be thrown in the database or GCP calls and handle them in logical groupings (mainly considering re-tryable vs fatal but also logging information)
• Emit clear error messages when errors occur
• Simplify error-handling for the caller, e.g., the cloud function, so it can determine if the request should be retried.

But I know that it is starting to look a lot like Java code now. :)

[andy-slac]

Looks OK to me. Exception handling is very dense and hard to follow. General question - in this new reimplementation with retries is there a possibility that chunks are ingested out of order? I worry is that chunks with update records implicitly depend on all previous chunks to be already ingested.

[JeremyMcCormick]

Looks OK to me. Exception handling is very dense and hard to follow. General question - in this new reimplementation with retries is there a possibility that chunks are ingested out of order? I worry is that chunks with update records implicitly depend on all previous chunks to be already ingested.

Chunks can be staged out of order as the function is not blocking and multiple copies can be running, but the promotion process should ensure that all prior chunks have already been promoted, including application of their updates. Updates from the whole set of chunks being promoted are applied at once to the new prod table. So it should not be possible for an update to be applied out of order or when a record it depends on is missing.

[JeremyMcCormick]

@andy-slac

Exception handling is very dense and hard to follow.

Agreed, the error handling got pretty byzantine, though much of that complexity is necessary for trapping and categorizing various types of errors. (We don't want to permanently fail a chunk if there is a brief connectivity issue within Google Cloud, for instance.)

Some of this should be cleaned up on DM-56938 where I plan to have the ppdb_bigquery handle database errors internally and then rethrow a more generic type.