BqUtils: support a BigQuery quota project override + consolidate client construction#706
Draft
kmontemayor2-sc wants to merge 11 commits into
Draft
BqUtils: support a BigQuery quota project override + consolidate client construction#706kmontemayor2-sc wants to merge 11 commits into
kmontemayor2-sc wants to merge 11 commits into
Conversation
Thread an optional `quota_project_id` through `_load_records_to_bigquery` and its public callers (`load_embeddings_to_bigquery`, `load_predictions_to_bigquery`). When provided, the BigQuery client is built with `ClientOptions(quota_project_id=...)` so only the BigQuery client's quota / billing project is scoped, leaving other Google clients (e.g. GCS) on their default quota project. The parameter defaults to None, so existing callers are unaffected: no quota project is forced and google-auth's default resolution applies. Extend the export unit tests to cover both the scoped and default cases. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Add quota-project resolution to `BqUtils.__init__` (new optional `quota_project_id` param): resolved as explicit param > the `GIGL_BIGQUERY_QUOTA_PROJECT` env var > none, where an empty value means "unset". When resolved, the underlying `bigquery.Client` is built with `ClientOptions(quota_project_id=...)` so the override scopes only this BigQuery client, leaving other Google clients unaffected. With neither the param nor the env var set, `client_options=None` preserves the exact prior behavior for all existing call sites. Also add `BqUtils.load_files_to_bq(...)`, a generic wildcard GCS-to-BQ load helper: the caller supplies the `LoadJobConfig`; sync waits for the job and logs the loaded row count, async returns the created job immediately. Add `GIGL_BIGQUERY_QUOTA_PROJECT_ENV_KEY` to `gigl/env/constants.py` and env-hermetic unit tests for both the quota resolution and the loader. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Migrate `_load_records_to_bigquery` off its own `bigquery.Client` onto `BqUtils` + `BqUtils.load_files_to_bq`, resolving the standing TODO. The public `load_embeddings_to_bigquery` / `load_predictions_to_bigquery` signatures and `LoadJob` return are unchanged; the `quota_project_id` param is now forwarded to `BqUtils` (which also honors the `GIGL_BIGQUERY_QUOTA_PROJECT` env var when the param is omitted). Behavior parity is preserved: Avro source format, WRITE_APPEND, schema, the `*.avro` glob, sync-waits / async-returns, and the returned job. The destination is built as a plain `project.dataset.table` string so domain-scoped project IDs are not rejected. Rewrite the export BQ tests to verify delegation to BqUtils (construction args + load_files_to_bq args) plus an end-to-end env-var test, and fix a `should_run_asnyc` docstring typo. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Replace the standalone `bigquery.Client()` in `_assert_shard_key_in_table` with `BqUtils.fetch_bq_table_schema`, so all GiGL BigQuery-client construction goes through `BqUtils` (and picks up the `GIGL_BIGQUERY_QUOTA_PROJECT` override). Behavior is unchanged: the schema lookup accepts the same table-path formats, and a missing table still propagates `NotFound`. Add the first unit tests for this module covering the pass, missing-key, and missing-table cases. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Drop the quota_project_id parameter from _load_records_to_bigquery, load_embeddings_to_bigquery, and load_predictions_to_bigquery. The export layer now gets the quota / billing project solely from the GIGL_BIGQUERY_QUOTA_PROJECT environment variable, resolved inside BqUtils. BqUtils itself keeps its explicit quota_project_id parameter for callers that need a per-client override. Trim the export unit tests accordingly: quota / billing resolution is BqUtils's responsibility and is covered in bq_test.py, so the export tests no longer assert on it. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Cover load_predictions_to_bigquery end-to-end against real GCS and BigQuery, mirroring the existing embedding export integration test. Real load behavior for the export helpers now lives in integration tests rather than mocked unit tests. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Replace the mocked sharded_read unit tests with an integration test that constructs ShardedExportRead against a real table, covering the valid shard-key path, the missing-column ValueError, and the missing table NotFound through the public constructor instead of the private helper. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Exercise ShardedExportRead against the real BigQuery table by running the read in a DirectRunner pipeline and asserting all rows are recovered across the shards, rather than only validating construction. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Remove test_load_embedding_to_bigquery and test_load_prediction_to_bigquery, which only verified mocks; the real load behavior is covered by the export integration tests. Prune imports left unused by the deletion. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Remove the load-files test class (covered by the export integration tests) and trim a redundant assertion from the client-construction precedence tests. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Collaborator
Author
|
/all_test |
Contributor
GiGL Automation@ 20:32:28UTC : 🔄 @ 20:40:24UTC : ✅ Workflow completed successfully. |
Contributor
GiGL Automation@ 20:32:29UTC : 🔄 @ 20:43:13UTC : ✅ Workflow completed successfully. |
Contributor
GiGL Automation@ 20:32:30UTC : 🔄 @ 20:34:21UTC : ✅ Workflow completed successfully. |
Contributor
GiGL Automation@ 20:32:31UTC : 🔄 @ 21:39:34UTC : ✅ Workflow completed successfully. |
Contributor
GiGL Automation@ 20:32:33UTC : 🔄 @ 22:45:01UTC : ✅ Workflow completed successfully. |
Contributor
GiGL Automation@ 20:32:33UTC : 🔄 @ 22:06:12UTC : ✅ Workflow completed successfully. |
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Why
The BigQuery client's quota project (used for API-enablement and per-API quota checks) defaults to the project associated with the active credentials. When a job runs under credentials whose default project differs from the project that owns the BigQuery resources — e.g. a workload that creates load jobs in a data project while its credentials belong to a different infrastructure project — the enablement/quota check runs against that default project, which may not have the BigQuery API enabled.
A per-client quota_project_id lets callers direct BigQuery's quota/billing to the appropriate project. It is intentionally BigQuery-scoped (set only on the BqUtils client), so it does not affect other Google Cloud clients (e.g. GCS) — unlike the process-global GOOGLE_CLOUD_QUOTA_PROJECT. Consolidating client construction in BqUtils means this (and future client config) applies uniformly to all callers without per-call-site plumbing.
Setting the quota project requires the caller to have serviceusage.services.use on that project (the roles/serviceusage.serviceUsageConsumer role).
Configuration