From 95f3ccf2b31a3bdf72757b6c6a1a02d8439ece37 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Tue, 14 Jul 2026 21:48:24 +0000 Subject: [PATCH 01/11] Add optional quota_project_id to BigQuery export helpers 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 --- gigl/common/data/export.py | 24 ++++++++- tests/unit/common/data/export_test.py | 74 ++++++++++++++++++++++++++- 2 files changed, 94 insertions(+), 4 deletions(-) diff --git a/gigl/common/data/export.py b/gigl/common/data/export.py index e4ebcca2a..28cc4fbf1 100644 --- a/gigl/common/data/export.py +++ b/gigl/common/data/export.py @@ -17,6 +17,7 @@ import fastavro.types import requests import torch +from google.api_core.client_options import ClientOptions from google.cloud import bigquery from google.cloud.bigquery.job import LoadJob from google.cloud.exceptions import GoogleCloudError @@ -348,6 +349,7 @@ def _load_records_to_bigquery( table_id: str, schema: Sequence[bigquery.SchemaField], should_run_async: bool = False, + quota_project_id: Optional[str] = None, ) -> LoadJob: """ Loads multiple Avro files containing GNN records from GCS into BigQuery. @@ -359,6 +361,9 @@ def _load_records_to_bigquery( table_id (str): The BigQuery table ID. schema (Sequence[bigquery.SchemaField]): The BigQuery schema for the records. should_run_async (bool): Whether loading to bigquery step should happen asynchronously. Defaults to False. + quota_project_id (Optional[str]): The GCP project to bill for BigQuery quota / usage. Scopes only this + BigQuery client, leaving other Google clients unaffected. When None (the default), no quota project + is forced and google-auth's default resolution applies. Returns: LoadJob: A BigQuery LoadJob object representing the load operation, which allows @@ -368,8 +373,13 @@ def _load_records_to_bigquery( """ start = time.perf_counter() logger.info(f"Loading records from {gcs_folder} to BigQuery.") - # Initialize the BigQuery client - bigquery_client = bigquery.Client(project=project_id) + # Initialize the BigQuery client. When quota_project_id is provided, scope the + # quota / billing project to this client only (leaving other Google clients on + # their default quota project). + client_options = ( + ClientOptions(quota_project_id=quota_project_id) if quota_project_id else None + ) + bigquery_client = bigquery.Client(project=project_id, client_options=client_options) # Construct dataset and table references dataset_ref = bigquery_client.dataset(dataset_id) @@ -407,6 +417,7 @@ def load_embeddings_to_bigquery( dataset_id: str, table_id: str, should_run_async: bool = False, + quota_project_id: Optional[str] = None, ) -> LoadJob: """ Loads multiple Avro files containing GNN embeddings from GCS into BigQuery. @@ -425,6 +436,9 @@ def load_embeddings_to_bigquery( dataset_id (str): The BigQuery dataset ID. table_id (str): The BigQuery table ID. should_run_async (bool): Whether loading to bigquery step should happen asynchronously. Defaults to False. + quota_project_id (Optional[str]): The GCP project to bill for BigQuery quota / usage. Scopes only this + BigQuery client, leaving other Google clients unaffected. When None (the default), no quota project + is forced and google-auth's default resolution applies. Returns: LoadJob: A BigQuery LoadJob object representing the load operation, which allows @@ -439,6 +453,7 @@ def load_embeddings_to_bigquery( table_id, EMBEDDING_BIGQUERY_SCHEMA, should_run_async, + quota_project_id=quota_project_id, ) @@ -448,6 +463,7 @@ def load_predictions_to_bigquery( dataset_id: str, table_id: str, should_run_async: bool = False, + quota_project_id: Optional[str] = None, ) -> LoadJob: """ Loads multiple Avro files containing GNN predictions from GCS into BigQuery. @@ -466,6 +482,9 @@ def load_predictions_to_bigquery( dataset_id (str): The BigQuery dataset ID. table_id (str): The BigQuery table ID. should_run_async (bool): Whether loading to bigquery step should happen asynchronously. Defaults to False. + quota_project_id (Optional[str]): The GCP project to bill for BigQuery quota / usage. Scopes only this + BigQuery client, leaving other Google clients unaffected. When None (the default), no quota project + is forced and google-auth's default resolution applies. Returns: LoadJob: A BigQuery LoadJob object representing the load operation, which allows @@ -480,4 +499,5 @@ def load_predictions_to_bigquery( table_id, PREDICTION_BIGQUERY_SCHEMA, should_run_async, + quota_project_id=quota_project_id, ) diff --git a/tests/unit/common/data/export_test.py b/tests/unit/common/data/export_test.py index 8f1f300b9..8a60fbbea 100644 --- a/tests/unit/common/data/export_test.py +++ b/tests/unit/common/data/export_test.py @@ -419,7 +419,10 @@ def test_load_embedding_to_bigquery( ) # Assertions - mock_bigquery_client.assert_called_once_with(project=project_id) + # When quota_project_id is omitted, no quota project is forced (client_options=None). + mock_bigquery_client.assert_called_once_with( + project=project_id, client_options=None + ) mock_client.load_table_from_uri.assert_called_once_with( source_uris=f"{gcs_folder.uri}/*.avro", destination=mock_client.dataset.return_value.table.return_value, @@ -427,6 +430,38 @@ def test_load_embedding_to_bigquery( ) self.assertEqual(load_job.output_rows, 1000) + @patch("gigl.common.data.export.bigquery.Client") + def test_load_embedding_to_bigquery_with_quota_project(self, mock_bigquery_client): + # Mock inputs + gcs_folder = GcsUri("gs://test-bucket/test-folder") + project_id = "test-project" + dataset_id = "test-dataset" + table_id = "test-table" + quota_project_id = "test-quota-project" + + # Mock BigQuery client and load job + mock_client = MagicMock() + mock_client.load_table_from_uri.return_value.output_rows = 1000 + mock_bigquery_client.return_value = mock_client + + # Call the function + load_embeddings_to_bigquery( + gcs_folder, + project_id, + dataset_id, + table_id, + quota_project_id=quota_project_id, + ) + + # The BigQuery client should be constructed with client_options whose + # quota_project_id is scoped to the provided value. + mock_bigquery_client.assert_called_once() + _, call_kwargs = mock_bigquery_client.call_args + self.assertEqual(call_kwargs["project"], project_id) + self.assertEqual( + call_kwargs["client_options"].quota_project_id, quota_project_id + ) + class TestPredictionsExporter(TestCase): def setUp(self): @@ -771,7 +806,10 @@ def test_load_prediction_to_bigquery( ) # Assertions - mock_bigquery_client.assert_called_once_with(project=project_id) + # When quota_project_id is omitted, no quota project is forced (client_options=None). + mock_bigquery_client.assert_called_once_with( + project=project_id, client_options=None + ) mock_client.load_table_from_uri.assert_called_once_with( source_uris=f"{gcs_folder.uri}/*.avro", destination=mock_client.dataset.return_value.table.return_value, @@ -779,6 +817,38 @@ def test_load_prediction_to_bigquery( ) self.assertEqual(load_job.output_rows, 1000) + @patch("gigl.common.data.export.bigquery.Client") + def test_load_prediction_to_bigquery_with_quota_project(self, mock_bigquery_client): + # Mock inputs + gcs_folder = GcsUri("gs://test-bucket/test-folder") + project_id = "test-project" + dataset_id = "test-dataset" + table_id = "test-table" + quota_project_id = "test-quota-project" + + # Mock BigQuery client and load job + mock_client = MagicMock() + mock_client.load_table_from_uri.return_value.output_rows = 1000 + mock_bigquery_client.return_value = mock_client + + # Call the function + load_predictions_to_bigquery( + gcs_folder, + project_id, + dataset_id, + table_id, + quota_project_id=quota_project_id, + ) + + # The BigQuery client should be constructed with client_options whose + # quota_project_id is scoped to the provided value. + mock_bigquery_client.assert_called_once() + _, call_kwargs = mock_bigquery_client.call_args + self.assertEqual(call_kwargs["project"], project_id) + self.assertEqual( + call_kwargs["client_options"].quota_project_id, quota_project_id + ) + if __name__ == "__main__": absltest.main() From aaddb79f7d28450189e9eb75c67b5437e34082c3 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Wed, 15 Jul 2026 18:00:06 +0000 Subject: [PATCH 02/11] BqUtils: honor GIGL_BIGQUERY_QUOTA_PROJECT and add load_files_to_bq 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 --- gigl/env/constants.py | 7 ++ gigl/src/common/utils/bq.py | 96 ++++++++++++++++++++++++-- tests/unit/src/common/utils/bq_test.py | 86 +++++++++++++++++++++++ 3 files changed, 184 insertions(+), 5 deletions(-) diff --git a/gigl/env/constants.py b/gigl/env/constants.py index 78a8c5505..7d84ea0de 100644 --- a/gigl/env/constants.py +++ b/gigl/env/constants.py @@ -26,6 +26,13 @@ GIGL_CUDA_DOCKER_URI_ENV_KEY: Final[str] = "GIGL_CUDA_DOCKER_URI" GIGL_COMPONENT_ENV_KEY: Final[str] = "GIGL_COMPONENT" +# Optional override for the quota / billing project used by GiGL's BigQuery +# clients. When set to a non-empty value, every ``BqUtils`` BigQuery client +# bills quota to this project. The override is scoped to BigQuery clients only, +# leaving other Google clients (e.g. Cloud Storage) on their default quota +# project. Unset (or empty) means no override. +GIGL_BIGQUERY_QUOTA_PROJECT_ENV_KEY: Final[str] = "GIGL_BIGQUERY_QUOTA_PROJECT" + _LEGACY_RESOURCE_CONFIG_ENV_KEY: Final[str] = "RESOURCE_CONFIG_PATH" _legacy_resource_config_env_warned: bool = False _logger = Logger() diff --git a/gigl/src/common/utils/bq.py b/gigl/src/common/utils/bq.py index dfdec91b9..ab24d8c36 100644 --- a/gigl/src/common/utils/bq.py +++ b/gigl/src/common/utils/bq.py @@ -1,18 +1,22 @@ import datetime import itertools +import os import re -from typing import Iterable, Optional, Tuple, Union +import time +from typing import Iterable, Optional, Sequence, Tuple, Union import google.api_core.retry import google.cloud.bigquery as bigquery +from google.api_core.client_options import ClientOptions from google.api_core.exceptions import NotFound from google.cloud.bigquery._helpers import _record_field_to_json -from google.cloud.bigquery.job import _AsyncJob +from google.cloud.bigquery.job import LoadJob, _AsyncJob from google.cloud.bigquery.table import RowIterator from gigl.common import GcsUri, LocalUri, Uri from gigl.common.logger import Logger from gigl.common.utils.retry import retry +from gigl.env.constants import GIGL_BIGQUERY_QUOTA_PROJECT_ENV_KEY from gigl.src.common.constants.time import DEFAULT_DATE_FORMAT from gigl.src.common.utils.time import convert_days_to_ms, current_datetime @@ -58,9 +62,44 @@ def _load_file_to_bq_with_retry( class BqUtils: - def __init__(self, project: Optional[str] = None) -> None: - logger.info(f"BqUtils initialized with project: {project}") - self.__bq_client = bigquery.Client(project=project) + def __init__( + self, + project: Optional[str] = None, + quota_project_id: Optional[str] = None, + ) -> None: + """Initializes a BqUtils instance wrapping a single BigQuery client. + + The quota / billing project for the underlying client is resolved as: + the explicit ``quota_project_id`` argument if provided, else the + ``GIGL_BIGQUERY_QUOTA_PROJECT`` environment variable if set to a + non-empty value, else no override (google-auth's default resolution). + + The override scopes only this BigQuery client; other Google clients in + the process are unaffected. + + Args: + project (Optional[str]): The GCP project the client acts on + (default project for datasets, tables, and jobs). + quota_project_id (Optional[str]): The GCP project billed for + BigQuery quota / usage for this client. + """ + # An empty string (from either the argument or the env var) is treated + # as "unset" so callers can opt out with an empty value. + resolved_quota_project_id = quota_project_id or os.environ.get( + GIGL_BIGQUERY_QUOTA_PROJECT_ENV_KEY + ) + client_options = ( + ClientOptions(quota_project_id=resolved_quota_project_id) + if resolved_quota_project_id + else None + ) + logger.info( + f"BqUtils initialized with project: {project}, " + f"quota project: {resolved_quota_project_id}" + ) + self.__bq_client = bigquery.Client( + project=project, client_options=client_options + ) def create_bq_dataset(self, dataset_id, exists_ok=True) -> None: dataset = bigquery.Dataset(dataset_id) @@ -412,6 +451,53 @@ def fetch_bq_table_schema(self, bq_table: str) -> dict[str, bigquery.SchemaField schema_dict = {field.name: field for field in bq_schema} return schema_dict + def load_files_to_bq( + self, + source_uris: Union[str, Sequence[str]], + bq_path: str, + job_config: bigquery.LoadJobConfig, + should_run_async: bool = False, + ) -> LoadJob: + """Loads one or more GCS files into a BigQuery table with a single load job. + + ``source_uris`` may contain a trailing wildcard (e.g. + ``gs://bucket/folder/*.avro``) to load every matching file. + + Args: + source_uris (Union[str, Sequence[str]]): GCS URI(s), optionally with + a ``*`` wildcard, of the file(s) to load. + bq_path (str): Destination table path in ``project.dataset.table`` + format. + job_config (bigquery.LoadJobConfig): The load job configuration + (source format, write disposition, schema, ...). + should_run_async (bool): If True, return immediately after the load + job is created without waiting for completion. Defaults to False. + + Returns: + LoadJob: The BigQuery load job. Completed if + ``should_run_async=False``; possibly still running if + ``should_run_async=True``. + """ + # Timing is measured around the load job here (rather than at the call + # site) so all callers get consistent load-duration logging. + start = time.perf_counter() + load_job = self.__bq_client.load_table_from_uri( + source_uris=source_uris, + destination=bq_path, + job_config=job_config, + ) + if should_run_async: + logger.info( + f"Started load job {load_job.job_id} for {bq_path}, running asynchronously." + ) + else: + load_job.result() # Waits for the job to complete. + logger.info( + f"Loaded {load_job.output_rows:,} rows into {bq_path} in " + f"{time.perf_counter() - start:.2f} seconds." + ) + return load_job + def load_file_to_bq( self, source_path: Uri, diff --git a/tests/unit/src/common/utils/bq_test.py b/tests/unit/src/common/utils/bq_test.py index efa02e342..c710d7780 100644 --- a/tests/unit/src/common/utils/bq_test.py +++ b/tests/unit/src/common/utils/bq_test.py @@ -1,8 +1,11 @@ +import os from typing import Sequence from unittest.mock import MagicMock, patch +import google.cloud.bigquery as bigquery from parameterized import param, parameterized +from gigl.env.constants import GIGL_BIGQUERY_QUOTA_PROJECT_ENV_KEY from gigl.src.common.utils.bq import BqUtils from tests.test_assets.test_case import TestCase @@ -141,3 +144,86 @@ def test_returns_latest_table_with_hourly_suffix( bq_table_path_prefix=self.PREFIX, table_partition_suffix="YYYYMMDDHH" ) self.assertEqual(result, "myproject.mydataset.events_2025010112") + + +@patch("gigl.src.common.utils.bq.bigquery.Client") +class BqClientConstructionTest(TestCase): + def test_no_quota_project_when_env_unset_and_no_param( + self, mock_client_cls: MagicMock + ) -> None: + with patch.dict(os.environ, {}, clear=True): + BqUtils(project="my-project") + _, call_kwargs = mock_client_cls.call_args + self.assertEqual(call_kwargs["project"], "my-project") + self.assertIsNone(call_kwargs["client_options"]) + + def test_env_var_sets_quota_project(self, mock_client_cls: MagicMock) -> None: + with patch.dict( + os.environ, {GIGL_BIGQUERY_QUOTA_PROJECT_ENV_KEY: "env-quota-project"} + ): + BqUtils(project="my-project") + _, call_kwargs = mock_client_cls.call_args + self.assertEqual(call_kwargs["project"], "my-project") + self.assertEqual( + call_kwargs["client_options"].quota_project_id, "env-quota-project" + ) + + def test_explicit_param_overrides_env_var(self, mock_client_cls: MagicMock) -> None: + with patch.dict( + os.environ, {GIGL_BIGQUERY_QUOTA_PROJECT_ENV_KEY: "env-quota-project"} + ): + BqUtils(project="my-project", quota_project_id="param-quota-project") + _, call_kwargs = mock_client_cls.call_args + self.assertEqual( + call_kwargs["client_options"].quota_project_id, "param-quota-project" + ) + + def test_empty_env_var_is_treated_as_unset( + self, mock_client_cls: MagicMock + ) -> None: + with patch.dict(os.environ, {GIGL_BIGQUERY_QUOTA_PROJECT_ENV_KEY: ""}): + BqUtils(project="my-project") + _, call_kwargs = mock_client_cls.call_args + self.assertIsNone(call_kwargs["client_options"]) + + +@patch("gigl.src.common.utils.bq.bigquery.Client") +class LoadFilesToBqTest(TestCase): + def test_sync_load_waits_for_completion_and_returns_job( + self, mock_client_cls: MagicMock + ) -> None: + mock_load_job = mock_client_cls.return_value.load_table_from_uri.return_value + mock_load_job.output_rows = 1000 + bq_utils = BqUtils(project="my-project") + job_config = bigquery.LoadJobConfig(source_format=bigquery.SourceFormat.AVRO) + + load_job = bq_utils.load_files_to_bq( + source_uris="gs://bucket/folder/*.avro", + bq_path="my-project.my-dataset.my-table", + job_config=job_config, + should_run_async=False, + ) + + mock_client_cls.return_value.load_table_from_uri.assert_called_once_with( + source_uris="gs://bucket/folder/*.avro", + destination="my-project.my-dataset.my-table", + job_config=job_config, + ) + mock_load_job.result.assert_called_once() + self.assertIs(load_job, mock_load_job) + + def test_async_load_returns_without_waiting( + self, mock_client_cls: MagicMock + ) -> None: + mock_load_job = mock_client_cls.return_value.load_table_from_uri.return_value + bq_utils = BqUtils(project="my-project") + + load_job = bq_utils.load_files_to_bq( + source_uris="gs://bucket/folder/*.avro", + bq_path="my-project.my-dataset.my-table", + job_config=bigquery.LoadJobConfig(), + should_run_async=True, + ) + + mock_load_job.result.assert_not_called() + self.assertIs(load_job, mock_load_job) From e066e2481e01f87912f68010fa753a6a84fd2638 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Wed, 15 Jul 2026 18:02:33 +0000 Subject: [PATCH 03/11] Route BigQuery Avro export through BqUtils.load_files_to_bq 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 --- gigl/common/data/export.py | 47 ++----- tests/unit/common/data/export_test.py | 192 ++++++++++++++------------ 2 files changed, 115 insertions(+), 124 deletions(-) diff --git a/gigl/common/data/export.py b/gigl/common/data/export.py index 28cc4fbf1..2bd0716d2 100644 --- a/gigl/common/data/export.py +++ b/gigl/common/data/export.py @@ -17,7 +17,6 @@ import fastavro.types import requests import torch -from google.api_core.client_options import ClientOptions from google.cloud import bigquery from google.cloud.bigquery.job import LoadJob from google.cloud.exceptions import GoogleCloudError @@ -26,6 +25,7 @@ from gigl.common import GcsUri, LocalUri, Uri from gigl.common.logger import Logger from gigl.common.utils.retry import retry +from gigl.src.common.utils.bq import BqUtils from gigl.src.common.utils.file_loader import FileLoader logger = Logger() @@ -341,7 +341,6 @@ def add_prediction( self.add_record(batched_records) -# TODO(kmonte): We should migrate this over to `BqUtils.load_files_to_bq` once that is implemented. def _load_records_to_bigquery( gcs_folder: GcsUri, project_id: str, @@ -362,28 +361,18 @@ def _load_records_to_bigquery( schema (Sequence[bigquery.SchemaField]): The BigQuery schema for the records. should_run_async (bool): Whether loading to bigquery step should happen asynchronously. Defaults to False. quota_project_id (Optional[str]): The GCP project to bill for BigQuery quota / usage. Scopes only this - BigQuery client, leaving other Google clients unaffected. When None (the default), no quota project - is forced and google-auth's default resolution applies. + BigQuery client, leaving other Google clients unaffected. When None (the default), the + ``GIGL_BIGQUERY_QUOTA_PROJECT`` environment variable is consulted; if that is also unset, no quota + project is forced and google-auth's default resolution applies. Returns: LoadJob: A BigQuery LoadJob object representing the load operation, which allows user to monitor and retrieve details about the job status and result. The returned job will be done if `should_run_async=False` and will be returned immediately after creation (not necessarily complete) if - `should_run_asnyc=True`. + `should_run_async=True`. """ - start = time.perf_counter() logger.info(f"Loading records from {gcs_folder} to BigQuery.") - # Initialize the BigQuery client. When quota_project_id is provided, scope the - # quota / billing project to this client only (leaving other Google clients on - # their default quota project). - client_options = ( - ClientOptions(quota_project_id=quota_project_id) if quota_project_id else None - ) - bigquery_client = bigquery.Client(project=project_id, client_options=client_options) - - # Construct dataset and table references - dataset_ref = bigquery_client.dataset(dataset_id) - table_ref = dataset_ref.table(table_id) + bq_utils = BqUtils(project=project_id, quota_project_id=quota_project_id) # Configure the load job job_config = bigquery.LoadJobConfig( @@ -392,24 +381,16 @@ def _load_records_to_bigquery( schema=schema, ) - load_job = bigquery_client.load_table_from_uri( + # Build the destination path directly rather than via BqUtils.join_path, + # which asserts on dot count and would reject domain-scoped project IDs + # (e.g. "example.com:project"). + return bq_utils.load_files_to_bq( source_uris=os.path.join(gcs_folder.uri, "*.avro"), - destination=table_ref, + bq_path=f"{project_id}.{dataset_id}.{table_id}", job_config=job_config, + should_run_async=should_run_async, ) - if should_run_async: - logger.info( - f"Started loading process for {dataset_id}:{table_id} with job id {load_job.job_id}, running asynchronously" - ) - else: - load_job.result() # Wait for the job to complete. - logger.info( - f"Loading {load_job.output_rows:,} rows into {dataset_id}:{table_id} in {time.perf_counter() - start:.2f} seconds." - ) - - return load_job - def load_embeddings_to_bigquery( gcs_folder: GcsUri, @@ -444,7 +425,7 @@ def load_embeddings_to_bigquery( LoadJob: A BigQuery LoadJob object representing the load operation, which allows user to monitor and retrieve details about the job status and result. The returned job will be done if `should_run_async=False` and will be returned immediately after creation (not necessarily complete) if - `should_run_asnyc=True`. + `should_run_async=True`. """ return _load_records_to_bigquery( gcs_folder, @@ -490,7 +471,7 @@ def load_predictions_to_bigquery( LoadJob: A BigQuery LoadJob object representing the load operation, which allows user to monitor and retrieve details about the job status and result. The returned job will be done if `should_run_async=False` and will be returned immediately after creation (not necessarily complete) if - `should_run_asnyc=True`. + `should_run_async=True`. """ return _load_records_to_bigquery( gcs_folder, diff --git a/tests/unit/common/data/export_test.py b/tests/unit/common/data/export_test.py index 8a60fbbea..9a839036d 100644 --- a/tests/unit/common/data/export_test.py +++ b/tests/unit/common/data/export_test.py @@ -1,14 +1,16 @@ import io +import os import tempfile from pathlib import Path from typing import Optional -from unittest.mock import ANY, MagicMock, patch +from unittest.mock import MagicMock, patch from uuid import uuid4 import fastavro import requests import torch from absl.testing import absltest +from google.cloud import bigquery from google.cloud.exceptions import GoogleCloudError from parameterized import param, parameterized @@ -18,12 +20,15 @@ _NODE_ID_KEY, _NODE_TYPE_KEY, _PREDICTION_KEY, + EMBEDDING_BIGQUERY_SCHEMA, + PREDICTION_BIGQUERY_SCHEMA, EmbeddingExporter, PredictionExporter, load_embeddings_to_bigquery, load_predictions_to_bigquery, ) from gigl.common.utils.retry import RetriesFailedException +from gigl.env.constants import GIGL_BIGQUERY_QUOTA_PROJECT_ENV_KEY from tests.test_assets.test_case import TestCase TEST_NODE_TYPE = "test_type" @@ -394,72 +399,86 @@ def test_skips_flush_if_empty(self, mock_gcs_utils_class): ), ] ) - @patch("gigl.common.data.export.bigquery.Client") + @patch("gigl.common.data.export.BqUtils") def test_load_embedding_to_bigquery( - self, _, mock_bigquery_client, should_run_async: bool + self, _, mock_bq_utils_cls, should_run_async: bool ): - # Mock inputs gcs_folder = GcsUri("gs://test-bucket/test-folder") project_id = "test-project" dataset_id = "test-dataset" table_id = "test-table" - # Mock BigQuery client and load job - mock_client = MagicMock() - mock_client.load_table_from_uri.return_value.output_rows = 1000 - mock_bigquery_client.return_value = mock_client + # Clear the env so the "no quota project" arg is not perturbed by a set + # GIGL_BIGQUERY_QUOTA_PROJECT in the developer's shell. + with patch.dict(os.environ, {}, clear=True): + load_job = load_embeddings_to_bigquery( + gcs_folder, + project_id, + dataset_id, + table_id, + should_run_async=should_run_async, + ) - # Call the function - load_job = load_embeddings_to_bigquery( - gcs_folder, - project_id, - dataset_id, - table_id, - should_run_async=should_run_async, + # The export delegates to BqUtils: no quota project is forced when the + # param is omitted (BqUtils itself resolves the env var). + mock_bq_utils_cls.assert_called_once_with( + project=project_id, quota_project_id=None ) - - # Assertions - # When quota_project_id is omitted, no quota project is forced (client_options=None). - mock_bigquery_client.assert_called_once_with( - project=project_id, client_options=None + mock_load = mock_bq_utils_cls.return_value.load_files_to_bq + mock_load.assert_called_once() + _, load_kwargs = mock_load.call_args + self.assertEqual(load_kwargs["source_uris"], f"{gcs_folder.uri}/*.avro") + self.assertEqual( + load_kwargs["bq_path"], f"{project_id}.{dataset_id}.{table_id}" ) - mock_client.load_table_from_uri.assert_called_once_with( - source_uris=f"{gcs_folder.uri}/*.avro", - destination=mock_client.dataset.return_value.table.return_value, - job_config=ANY, + self.assertEqual(load_kwargs["should_run_async"], should_run_async) + job_config = load_kwargs["job_config"] + self.assertEqual(job_config.source_format, bigquery.SourceFormat.AVRO) + self.assertEqual( + job_config.write_disposition, bigquery.WriteDisposition.WRITE_APPEND ) - self.assertEqual(load_job.output_rows, 1000) + self.assertEqual(list(job_config.schema), list(EMBEDDING_BIGQUERY_SCHEMA)) + self.assertIs(load_job, mock_load.return_value) - @patch("gigl.common.data.export.bigquery.Client") - def test_load_embedding_to_bigquery_with_quota_project(self, mock_bigquery_client): - # Mock inputs + @patch("gigl.common.data.export.BqUtils") + def test_load_embedding_to_bigquery_with_quota_project(self, mock_bq_utils_cls): gcs_folder = GcsUri("gs://test-bucket/test-folder") - project_id = "test-project" - dataset_id = "test-dataset" - table_id = "test-table" quota_project_id = "test-quota-project" - # Mock BigQuery client and load job - mock_client = MagicMock() - mock_client.load_table_from_uri.return_value.output_rows = 1000 - mock_bigquery_client.return_value = mock_client - - # Call the function load_embeddings_to_bigquery( gcs_folder, - project_id, - dataset_id, - table_id, + "test-project", + "test-dataset", + "test-table", quota_project_id=quota_project_id, ) - # The BigQuery client should be constructed with client_options whose - # quota_project_id is scoped to the provided value. - mock_bigquery_client.assert_called_once() + # The explicit quota project override is forwarded to BqUtils. + mock_bq_utils_cls.assert_called_once_with( + project="test-project", quota_project_id=quota_project_id + ) + + @patch("gigl.src.common.utils.bq.bigquery.Client") + def test_load_embedding_to_bigquery_reads_quota_project_from_env( + self, mock_bigquery_client + ): + # End-to-end through the real BqUtils: the env var must reach the + # underlying BigQuery client's quota project. + gcs_folder = GcsUri("gs://test-bucket/test-folder") + mock_client = MagicMock() + mock_client.load_table_from_uri.return_value.output_rows = 1000 + mock_bigquery_client.return_value = mock_client + + with patch.dict( + os.environ, {GIGL_BIGQUERY_QUOTA_PROJECT_ENV_KEY: "env-quota-project"} + ): + load_embeddings_to_bigquery( + gcs_folder, "test-project", "test-dataset", "test-table" + ) + _, call_kwargs = mock_bigquery_client.call_args - self.assertEqual(call_kwargs["project"], project_id) self.assertEqual( - call_kwargs["client_options"].quota_project_id, quota_project_id + call_kwargs["client_options"].quota_project_id, "env-quota-project" ) @@ -781,72 +800,63 @@ def test_skips_flush_if_empty(self, mock_gcs_utils_class): ), ] ) - @patch("gigl.common.data.export.bigquery.Client") + @patch("gigl.common.data.export.BqUtils") def test_load_prediction_to_bigquery( - self, _, mock_bigquery_client, should_run_async: bool + self, _, mock_bq_utils_cls, should_run_async: bool ): - # Mock inputs gcs_folder = GcsUri("gs://test-bucket/test-folder") project_id = "test-project" dataset_id = "test-dataset" table_id = "test-table" - # Mock BigQuery client and load job - mock_client = MagicMock() - mock_client.load_table_from_uri.return_value.output_rows = 1000 - mock_bigquery_client.return_value = mock_client + # Clear the env so the "no quota project" arg is not perturbed by a set + # GIGL_BIGQUERY_QUOTA_PROJECT in the developer's shell. + with patch.dict(os.environ, {}, clear=True): + load_job = load_predictions_to_bigquery( + gcs_folder, + project_id, + dataset_id, + table_id, + should_run_async=should_run_async, + ) - # Call the function - load_job = load_predictions_to_bigquery( - gcs_folder, - project_id, - dataset_id, - table_id, - should_run_async=should_run_async, + # The export delegates to BqUtils: no quota project is forced when the + # param is omitted (BqUtils itself resolves the env var). + mock_bq_utils_cls.assert_called_once_with( + project=project_id, quota_project_id=None ) - - # Assertions - # When quota_project_id is omitted, no quota project is forced (client_options=None). - mock_bigquery_client.assert_called_once_with( - project=project_id, client_options=None + mock_load = mock_bq_utils_cls.return_value.load_files_to_bq + mock_load.assert_called_once() + _, load_kwargs = mock_load.call_args + self.assertEqual(load_kwargs["source_uris"], f"{gcs_folder.uri}/*.avro") + self.assertEqual( + load_kwargs["bq_path"], f"{project_id}.{dataset_id}.{table_id}" ) - mock_client.load_table_from_uri.assert_called_once_with( - source_uris=f"{gcs_folder.uri}/*.avro", - destination=mock_client.dataset.return_value.table.return_value, - job_config=ANY, + self.assertEqual(load_kwargs["should_run_async"], should_run_async) + job_config = load_kwargs["job_config"] + self.assertEqual(job_config.source_format, bigquery.SourceFormat.AVRO) + self.assertEqual( + job_config.write_disposition, bigquery.WriteDisposition.WRITE_APPEND ) - self.assertEqual(load_job.output_rows, 1000) + self.assertEqual(list(job_config.schema), list(PREDICTION_BIGQUERY_SCHEMA)) + self.assertIs(load_job, mock_load.return_value) - @patch("gigl.common.data.export.bigquery.Client") - def test_load_prediction_to_bigquery_with_quota_project(self, mock_bigquery_client): - # Mock inputs + @patch("gigl.common.data.export.BqUtils") + def test_load_prediction_to_bigquery_with_quota_project(self, mock_bq_utils_cls): gcs_folder = GcsUri("gs://test-bucket/test-folder") - project_id = "test-project" - dataset_id = "test-dataset" - table_id = "test-table" quota_project_id = "test-quota-project" - # Mock BigQuery client and load job - mock_client = MagicMock() - mock_client.load_table_from_uri.return_value.output_rows = 1000 - mock_bigquery_client.return_value = mock_client - - # Call the function load_predictions_to_bigquery( gcs_folder, - project_id, - dataset_id, - table_id, + "test-project", + "test-dataset", + "test-table", quota_project_id=quota_project_id, ) - # The BigQuery client should be constructed with client_options whose - # quota_project_id is scoped to the provided value. - mock_bigquery_client.assert_called_once() - _, call_kwargs = mock_bigquery_client.call_args - self.assertEqual(call_kwargs["project"], project_id) - self.assertEqual( - call_kwargs["client_options"].quota_project_id, quota_project_id + # The explicit quota project override is forwarded to BqUtils. + mock_bq_utils_cls.assert_called_once_with( + project="test-project", quota_project_id=quota_project_id ) From 9617bbb9425c472de1c1fab75127234a3f16329a Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Wed, 15 Jul 2026 18:03:28 +0000 Subject: [PATCH 04/11] Route sharded_read shard-key validation through BqUtils 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 --- gigl/common/beam/sharded_read.py | 12 +++--- tests/unit/common/beam/__init__.py | 0 tests/unit/common/beam/sharded_read_test.py | 48 +++++++++++++++++++++ 3 files changed, 54 insertions(+), 6 deletions(-) create mode 100644 tests/unit/common/beam/__init__.py create mode 100644 tests/unit/common/beam/sharded_read_test.py diff --git a/gigl/common/beam/sharded_read.py b/gigl/common/beam/sharded_read.py index affff00fa..e14608f43 100644 --- a/gigl/common/beam/sharded_read.py +++ b/gigl/common/beam/sharded_read.py @@ -4,9 +4,9 @@ from apache_beam.io.gcp.bigquery import BigQueryQueryPriority from apache_beam.io.gcp.internal.clients.bigquery import DatasetReference from apache_beam.pvalue import PBegin -from google.cloud import bigquery from gigl.common.logger import Logger +from gigl.src.common.utils.bq import BqUtils logger = Logger() @@ -31,11 +31,11 @@ def _assert_shard_key_in_table(table_name: str, shard_key: str) -> None: """ Validate that the shard key is a valid column in the BigQuery table. """ - client = bigquery.Client() - table_ref = bigquery.TableReference.from_string(table_name) - - table = client.get_table(table_ref) - column_names = [field.name for field in table.schema] + # BqUtils.fetch_bq_table_schema resolves the table via BigQuery's own path + # parsing, so it accepts the same `project.dataset.table` strings (and, as a + # superset, `project:dataset.table`) as the previous TableReference lookup. + bq_utils = BqUtils() + column_names = list(bq_utils.fetch_bq_table_schema(bq_table=table_name).keys()) if shard_key not in column_names: raise ValueError( diff --git a/tests/unit/common/beam/__init__.py b/tests/unit/common/beam/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit/common/beam/sharded_read_test.py b/tests/unit/common/beam/sharded_read_test.py new file mode 100644 index 000000000..cc8c48055 --- /dev/null +++ b/tests/unit/common/beam/sharded_read_test.py @@ -0,0 +1,48 @@ +from unittest.mock import MagicMock, patch + +from absl.testing import absltest +from google.api_core.exceptions import NotFound + +from gigl.common.beam.sharded_read import _assert_shard_key_in_table +from tests.test_assets.test_case import TestCase + + +@patch("gigl.common.beam.sharded_read.BqUtils") +class AssertShardKeyInTableTest(TestCase): + def test_passes_when_shard_key_is_a_column( + self, mock_bq_utils_cls: MagicMock + ) -> None: + mock_bq_utils_cls.return_value.fetch_bq_table_schema.return_value = { + "user_id": MagicMock(), + "timestamp": MagicMock(), + } + _assert_shard_key_in_table( + table_name="my-project.my-dataset.my-table", shard_key="user_id" + ) + mock_bq_utils_cls.return_value.fetch_bq_table_schema.assert_called_once_with( + bq_table="my-project.my-dataset.my-table" + ) + + def test_raises_when_shard_key_missing(self, mock_bq_utils_cls: MagicMock) -> None: + mock_bq_utils_cls.return_value.fetch_bq_table_schema.return_value = { + "user_id": MagicMock(), + } + with self.assertRaises(ValueError): + _assert_shard_key_in_table( + table_name="my-project.my-dataset.my-table", shard_key="not_a_column" + ) + + def test_propagates_not_found_for_missing_table( + self, mock_bq_utils_cls: MagicMock + ) -> None: + mock_bq_utils_cls.return_value.fetch_bq_table_schema.side_effect = NotFound( + "Table not found" + ) + with self.assertRaises(NotFound): + _assert_shard_key_in_table( + table_name="my-project.my-dataset.missing-table", shard_key="user_id" + ) + + +if __name__ == "__main__": + absltest.main() From 131e971b67753eec836f06fc9438bd6383aa1966 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Wed, 15 Jul 2026 20:01:07 +0000 Subject: [PATCH 05/11] Resolve export BigQuery quota project from env only 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 --- gigl/common/data/export.py | 29 +++---- tests/unit/common/data/export_test.py | 109 +++++--------------------- 2 files changed, 31 insertions(+), 107 deletions(-) diff --git a/gigl/common/data/export.py b/gigl/common/data/export.py index 2bd0716d2..de55ee0c6 100644 --- a/gigl/common/data/export.py +++ b/gigl/common/data/export.py @@ -348,7 +348,6 @@ def _load_records_to_bigquery( table_id: str, schema: Sequence[bigquery.SchemaField], should_run_async: bool = False, - quota_project_id: Optional[str] = None, ) -> LoadJob: """ Loads multiple Avro files containing GNN records from GCS into BigQuery. @@ -360,10 +359,10 @@ def _load_records_to_bigquery( table_id (str): The BigQuery table ID. schema (Sequence[bigquery.SchemaField]): The BigQuery schema for the records. should_run_async (bool): Whether loading to bigquery step should happen asynchronously. Defaults to False. - quota_project_id (Optional[str]): The GCP project to bill for BigQuery quota / usage. Scopes only this - BigQuery client, leaving other Google clients unaffected. When None (the default), the - ``GIGL_BIGQUERY_QUOTA_PROJECT`` environment variable is consulted; if that is also unset, no quota - project is forced and google-auth's default resolution applies. + + The BigQuery quota / billing project may be overridden via the + ``GIGL_BIGQUERY_QUOTA_PROJECT`` environment variable, which is resolved by + :class:`~gigl.src.common.utils.bq.BqUtils`. Returns: LoadJob: A BigQuery LoadJob object representing the load operation, which allows @@ -372,7 +371,7 @@ def _load_records_to_bigquery( `should_run_async=True`. """ logger.info(f"Loading records from {gcs_folder} to BigQuery.") - bq_utils = BqUtils(project=project_id, quota_project_id=quota_project_id) + bq_utils = BqUtils(project=project_id) # Configure the load job job_config = bigquery.LoadJobConfig( @@ -398,7 +397,6 @@ def load_embeddings_to_bigquery( dataset_id: str, table_id: str, should_run_async: bool = False, - quota_project_id: Optional[str] = None, ) -> LoadJob: """ Loads multiple Avro files containing GNN embeddings from GCS into BigQuery. @@ -417,9 +415,10 @@ def load_embeddings_to_bigquery( dataset_id (str): The BigQuery dataset ID. table_id (str): The BigQuery table ID. should_run_async (bool): Whether loading to bigquery step should happen asynchronously. Defaults to False. - quota_project_id (Optional[str]): The GCP project to bill for BigQuery quota / usage. Scopes only this - BigQuery client, leaving other Google clients unaffected. When None (the default), no quota project - is forced and google-auth's default resolution applies. + + The BigQuery quota / billing project may be overridden via the + ``GIGL_BIGQUERY_QUOTA_PROJECT`` environment variable, which is resolved by + :class:`~gigl.src.common.utils.bq.BqUtils`. Returns: LoadJob: A BigQuery LoadJob object representing the load operation, which allows @@ -434,7 +433,6 @@ def load_embeddings_to_bigquery( table_id, EMBEDDING_BIGQUERY_SCHEMA, should_run_async, - quota_project_id=quota_project_id, ) @@ -444,7 +442,6 @@ def load_predictions_to_bigquery( dataset_id: str, table_id: str, should_run_async: bool = False, - quota_project_id: Optional[str] = None, ) -> LoadJob: """ Loads multiple Avro files containing GNN predictions from GCS into BigQuery. @@ -463,9 +460,10 @@ def load_predictions_to_bigquery( dataset_id (str): The BigQuery dataset ID. table_id (str): The BigQuery table ID. should_run_async (bool): Whether loading to bigquery step should happen asynchronously. Defaults to False. - quota_project_id (Optional[str]): The GCP project to bill for BigQuery quota / usage. Scopes only this - BigQuery client, leaving other Google clients unaffected. When None (the default), no quota project - is forced and google-auth's default resolution applies. + + The BigQuery quota / billing project may be overridden via the + ``GIGL_BIGQUERY_QUOTA_PROJECT`` environment variable, which is resolved by + :class:`~gigl.src.common.utils.bq.BqUtils`. Returns: LoadJob: A BigQuery LoadJob object representing the load operation, which allows @@ -480,5 +478,4 @@ def load_predictions_to_bigquery( table_id, PREDICTION_BIGQUERY_SCHEMA, should_run_async, - quota_project_id=quota_project_id, ) diff --git a/tests/unit/common/data/export_test.py b/tests/unit/common/data/export_test.py index 9a839036d..aa4c5327f 100644 --- a/tests/unit/common/data/export_test.py +++ b/tests/unit/common/data/export_test.py @@ -1,5 +1,4 @@ import io -import os import tempfile from pathlib import Path from typing import Optional @@ -28,7 +27,6 @@ load_predictions_to_bigquery, ) from gigl.common.utils.retry import RetriesFailedException -from gigl.env.constants import GIGL_BIGQUERY_QUOTA_PROJECT_ENV_KEY from tests.test_assets.test_case import TestCase TEST_NODE_TYPE = "test_type" @@ -408,22 +406,16 @@ def test_load_embedding_to_bigquery( dataset_id = "test-dataset" table_id = "test-table" - # Clear the env so the "no quota project" arg is not perturbed by a set - # GIGL_BIGQUERY_QUOTA_PROJECT in the developer's shell. - with patch.dict(os.environ, {}, clear=True): - load_job = load_embeddings_to_bigquery( - gcs_folder, - project_id, - dataset_id, - table_id, - should_run_async=should_run_async, - ) - - # The export delegates to BqUtils: no quota project is forced when the - # param is omitted (BqUtils itself resolves the env var). - mock_bq_utils_cls.assert_called_once_with( - project=project_id, quota_project_id=None + load_job = load_embeddings_to_bigquery( + gcs_folder, + project_id, + dataset_id, + table_id, + should_run_async=should_run_async, ) + + # The export delegates to BqUtils, which owns quota-project resolution. + mock_bq_utils_cls.assert_called_once_with(project=project_id) mock_load = mock_bq_utils_cls.return_value.load_files_to_bq mock_load.assert_called_once() _, load_kwargs = mock_load.call_args @@ -440,47 +432,6 @@ def test_load_embedding_to_bigquery( self.assertEqual(list(job_config.schema), list(EMBEDDING_BIGQUERY_SCHEMA)) self.assertIs(load_job, mock_load.return_value) - @patch("gigl.common.data.export.BqUtils") - def test_load_embedding_to_bigquery_with_quota_project(self, mock_bq_utils_cls): - gcs_folder = GcsUri("gs://test-bucket/test-folder") - quota_project_id = "test-quota-project" - - load_embeddings_to_bigquery( - gcs_folder, - "test-project", - "test-dataset", - "test-table", - quota_project_id=quota_project_id, - ) - - # The explicit quota project override is forwarded to BqUtils. - mock_bq_utils_cls.assert_called_once_with( - project="test-project", quota_project_id=quota_project_id - ) - - @patch("gigl.src.common.utils.bq.bigquery.Client") - def test_load_embedding_to_bigquery_reads_quota_project_from_env( - self, mock_bigquery_client - ): - # End-to-end through the real BqUtils: the env var must reach the - # underlying BigQuery client's quota project. - gcs_folder = GcsUri("gs://test-bucket/test-folder") - mock_client = MagicMock() - mock_client.load_table_from_uri.return_value.output_rows = 1000 - mock_bigquery_client.return_value = mock_client - - with patch.dict( - os.environ, {GIGL_BIGQUERY_QUOTA_PROJECT_ENV_KEY: "env-quota-project"} - ): - load_embeddings_to_bigquery( - gcs_folder, "test-project", "test-dataset", "test-table" - ) - - _, call_kwargs = mock_bigquery_client.call_args - self.assertEqual( - call_kwargs["client_options"].quota_project_id, "env-quota-project" - ) - class TestPredictionsExporter(TestCase): def setUp(self): @@ -809,22 +760,16 @@ def test_load_prediction_to_bigquery( dataset_id = "test-dataset" table_id = "test-table" - # Clear the env so the "no quota project" arg is not perturbed by a set - # GIGL_BIGQUERY_QUOTA_PROJECT in the developer's shell. - with patch.dict(os.environ, {}, clear=True): - load_job = load_predictions_to_bigquery( - gcs_folder, - project_id, - dataset_id, - table_id, - should_run_async=should_run_async, - ) - - # The export delegates to BqUtils: no quota project is forced when the - # param is omitted (BqUtils itself resolves the env var). - mock_bq_utils_cls.assert_called_once_with( - project=project_id, quota_project_id=None + load_job = load_predictions_to_bigquery( + gcs_folder, + project_id, + dataset_id, + table_id, + should_run_async=should_run_async, ) + + # The export delegates to BqUtils, which owns quota-project resolution. + mock_bq_utils_cls.assert_called_once_with(project=project_id) mock_load = mock_bq_utils_cls.return_value.load_files_to_bq mock_load.assert_called_once() _, load_kwargs = mock_load.call_args @@ -841,24 +786,6 @@ def test_load_prediction_to_bigquery( self.assertEqual(list(job_config.schema), list(PREDICTION_BIGQUERY_SCHEMA)) self.assertIs(load_job, mock_load.return_value) - @patch("gigl.common.data.export.BqUtils") - def test_load_prediction_to_bigquery_with_quota_project(self, mock_bq_utils_cls): - gcs_folder = GcsUri("gs://test-bucket/test-folder") - quota_project_id = "test-quota-project" - - load_predictions_to_bigquery( - gcs_folder, - "test-project", - "test-dataset", - "test-table", - quota_project_id=quota_project_id, - ) - - # The explicit quota project override is forwarded to BqUtils. - mock_bq_utils_cls.assert_called_once_with( - project="test-project", quota_project_id=quota_project_id - ) - if __name__ == "__main__": absltest.main() From a1912a42e8917fe67913d2921566355b058d0f2f Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Wed, 15 Jul 2026 20:01:54 +0000 Subject: [PATCH 06/11] Add prediction-load BigQuery integration test 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 --- tests/integration/common/data/export_test.py | 62 +++++++++++++++++++- 1 file changed, 61 insertions(+), 1 deletion(-) diff --git a/tests/integration/common/data/export_test.py b/tests/integration/common/data/export_test.py index cb54344b5..f83b3fa9e 100644 --- a/tests/integration/common/data/export_test.py +++ b/tests/integration/common/data/export_test.py @@ -5,7 +5,12 @@ from parameterized import param, parameterized from gigl.common import GcsUri -from gigl.common.data.export import EmbeddingExporter, load_embeddings_to_bigquery +from gigl.common.data.export import ( + EmbeddingExporter, + PredictionExporter, + load_embeddings_to_bigquery, + load_predictions_to_bigquery, +) from gigl.common.logger import Logger from gigl.common.utils.gcs import GcsUtils from gigl.env.pipelines_config import get_resource_config @@ -105,3 +110,58 @@ def test_embedding_export(self, _, should_run_async: bool): bq_client.count_number_of_rows_in_bq_table(bq_export_table_path), num_nodes * 2, ) + + +class PredictionExportIntegrationTest(TestCase): + def setUp(self): + resource_config = get_resource_config() + test_unique_name = f"GiGL-Integration-Prediction-Exporter-{uuid.uuid4().hex}" + self.prediction_output_dir = GcsUri.join( + resource_config.temp_assets_regional_bucket_path, + test_unique_name, + "predictions", + ) + self.prediction_output_bq_project = resource_config.project + self.prediction_output_bq_dataset = resource_config.temp_assets_bq_dataset_name + self.prediction_output_bq_table = test_unique_name + + def tearDown(self): + gcs_utils = GcsUtils() + gcs_utils.delete_files_in_bucket_dir(self.prediction_output_dir) + bq_client = BqUtils() + bq_export_table_path = bq_client.join_path( + self.prediction_output_bq_project, + self.prediction_output_bq_dataset, + self.prediction_output_bq_table, + ) + bq_client.delete_bq_table_if_exist( + bq_table_path=bq_export_table_path, + ) + + def test_prediction_export(self): + num_nodes = 100 + with PredictionExporter(export_dir=self.prediction_output_dir) as exporter: + for i in torch.arange(num_nodes): + exporter.add_prediction(torch.tensor([i]), torch.ones(1) * i, "node") + + bq_client = BqUtils() + bq_export_table_path = bq_client.join_path( + self.prediction_output_bq_project, + self.prediction_output_bq_dataset, + self.prediction_output_bq_table, + ) + logger.info( + f"Will try exporting {self.prediction_output_dir} to BQ: {bq_export_table_path}" + ) + load_job = load_predictions_to_bigquery( + gcs_folder=self.prediction_output_dir, + project_id=self.prediction_output_bq_project, + dataset_id=self.prediction_output_bq_dataset, + table_id=self.prediction_output_bq_table, + ) + + self.assertEqual(load_job.output_rows, num_nodes) + self.assertEqual( + bq_client.count_number_of_rows_in_bq_table(bq_export_table_path), + num_nodes, + ) From 4d5fadbe0503f7ac38e76f0425cc0525d8525e9e Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Wed, 15 Jul 2026 20:03:38 +0000 Subject: [PATCH 07/11] Test sharded_read shard-key validation against real BigQuery 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 --- .../common/beam/__init__.py | 0 .../common/beam/sharded_read_test.py | 95 +++++++++++++++++++ tests/unit/common/beam/sharded_read_test.py | 48 ---------- 3 files changed, 95 insertions(+), 48 deletions(-) rename tests/{unit => integration}/common/beam/__init__.py (100%) create mode 100644 tests/integration/common/beam/sharded_read_test.py delete mode 100644 tests/unit/common/beam/sharded_read_test.py diff --git a/tests/unit/common/beam/__init__.py b/tests/integration/common/beam/__init__.py similarity index 100% rename from tests/unit/common/beam/__init__.py rename to tests/integration/common/beam/__init__.py diff --git a/tests/integration/common/beam/sharded_read_test.py b/tests/integration/common/beam/sharded_read_test.py new file mode 100644 index 000000000..d00db76fd --- /dev/null +++ b/tests/integration/common/beam/sharded_read_test.py @@ -0,0 +1,95 @@ +import uuid + +from absl.testing import absltest +from google.api_core.exceptions import NotFound + +from gigl.common.beam.sharded_read import ( + BigQueryShardedReadConfig, + ShardedExportRead, +) +from gigl.env.pipelines_config import get_resource_config +from gigl.src.common.utils.bq import BqUtils +from tests.test_assets.test_case import TestCase + + +class ShardedReadIntegrationTest(TestCase): + """Integration tests for ShardedExportRead's shard-key validation against real BigQuery.""" + + def setUp(self): + resource_config = get_resource_config() + test_unique_name = f"gigl_sharded_read_test_{uuid.uuid4().hex}" + + self.bq_utils = BqUtils(project=resource_config.project) + self.bq_project = resource_config.project + self.bq_dataset = resource_config.temp_assets_bq_dataset_name + self.table_path = self.bq_utils.join_path( + self.bq_project, + self.bq_dataset, + test_unique_name, + ) + + table_path = self.bq_utils.format_bq_path(self.table_path) + create_query = f""" + CREATE OR REPLACE TABLE `{table_path}` AS + SELECT + id as user_id, + CONCAT('test_name_', CAST(id AS STRING)) as name + FROM UNNEST(GENERATE_ARRAY(0, 19)) as id + """ + self.bq_utils.run_query(query=create_query, labels={}) + + def tearDown(self): + self.bq_utils.delete_bq_table_if_exist( + bq_table_path=self.table_path, not_found_ok=True + ) + + def test_constructs_when_shard_key_is_a_column(self): + sharded_read_config = BigQueryShardedReadConfig( + shard_key="user_id", + project_id=self.bq_project, + temp_dataset_name=self.bq_dataset, + num_shards=2, + ) + + # Construction validates the shard key against the real table schema. + ShardedExportRead( + table_name=self.table_path, + sharded_read_info=sharded_read_config, + ) + + def test_raises_when_shard_key_is_not_a_column(self): + sharded_read_config = BigQueryShardedReadConfig( + shard_key="not_a_column", + project_id=self.bq_project, + temp_dataset_name=self.bq_dataset, + num_shards=2, + ) + + with self.assertRaises(ValueError): + ShardedExportRead( + table_name=self.table_path, + sharded_read_info=sharded_read_config, + ) + + def test_raises_when_table_does_not_exist(self): + missing_table_path = self.bq_utils.join_path( + self.bq_project, + self.bq_dataset, + f"gigl_sharded_read_missing_{uuid.uuid4().hex}", + ) + sharded_read_config = BigQueryShardedReadConfig( + shard_key="user_id", + project_id=self.bq_project, + temp_dataset_name=self.bq_dataset, + num_shards=2, + ) + + with self.assertRaises(NotFound): + ShardedExportRead( + table_name=missing_table_path, + sharded_read_info=sharded_read_config, + ) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/unit/common/beam/sharded_read_test.py b/tests/unit/common/beam/sharded_read_test.py deleted file mode 100644 index cc8c48055..000000000 --- a/tests/unit/common/beam/sharded_read_test.py +++ /dev/null @@ -1,48 +0,0 @@ -from unittest.mock import MagicMock, patch - -from absl.testing import absltest -from google.api_core.exceptions import NotFound - -from gigl.common.beam.sharded_read import _assert_shard_key_in_table -from tests.test_assets.test_case import TestCase - - -@patch("gigl.common.beam.sharded_read.BqUtils") -class AssertShardKeyInTableTest(TestCase): - def test_passes_when_shard_key_is_a_column( - self, mock_bq_utils_cls: MagicMock - ) -> None: - mock_bq_utils_cls.return_value.fetch_bq_table_schema.return_value = { - "user_id": MagicMock(), - "timestamp": MagicMock(), - } - _assert_shard_key_in_table( - table_name="my-project.my-dataset.my-table", shard_key="user_id" - ) - mock_bq_utils_cls.return_value.fetch_bq_table_schema.assert_called_once_with( - bq_table="my-project.my-dataset.my-table" - ) - - def test_raises_when_shard_key_missing(self, mock_bq_utils_cls: MagicMock) -> None: - mock_bq_utils_cls.return_value.fetch_bq_table_schema.return_value = { - "user_id": MagicMock(), - } - with self.assertRaises(ValueError): - _assert_shard_key_in_table( - table_name="my-project.my-dataset.my-table", shard_key="not_a_column" - ) - - def test_propagates_not_found_for_missing_table( - self, mock_bq_utils_cls: MagicMock - ) -> None: - mock_bq_utils_cls.return_value.fetch_bq_table_schema.side_effect = NotFound( - "Table not found" - ) - with self.assertRaises(NotFound): - _assert_shard_key_in_table( - table_name="my-project.my-dataset.missing-table", shard_key="user_id" - ) - - -if __name__ == "__main__": - absltest.main() From ee77f88a5c9070068415631f5aed71a9065bc1fe Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Wed, 15 Jul 2026 20:26:10 +0000 Subject: [PATCH 08/11] Run sharded read end-to-end in integration test 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 --- .../common/beam/sharded_read_test.py | 34 ++++++++++++++++--- 1 file changed, 30 insertions(+), 4 deletions(-) diff --git a/tests/integration/common/beam/sharded_read_test.py b/tests/integration/common/beam/sharded_read_test.py index d00db76fd..a51069d0b 100644 --- a/tests/integration/common/beam/sharded_read_test.py +++ b/tests/integration/common/beam/sharded_read_test.py @@ -1,6 +1,9 @@ import uuid +import apache_beam as beam from absl.testing import absltest +from apache_beam.options.pipeline_options import GoogleCloudOptions, PipelineOptions +from apache_beam.testing.util import assert_that, equal_to from google.api_core.exceptions import NotFound from gigl.common.beam.sharded_read import ( @@ -11,9 +14,11 @@ from gigl.src.common.utils.bq import BqUtils from tests.test_assets.test_case import TestCase +NUM_ROWS = 20 + class ShardedReadIntegrationTest(TestCase): - """Integration tests for ShardedExportRead's shard-key validation against real BigQuery.""" + """Integration tests for ShardedExportRead against real BigQuery.""" def setUp(self): resource_config = get_resource_config() @@ -22,6 +27,10 @@ def setUp(self): self.bq_utils = BqUtils(project=resource_config.project) self.bq_project = resource_config.project self.bq_dataset = resource_config.temp_assets_bq_dataset_name + self.temp_location = ( + f"{resource_config.temp_assets_regional_bucket_path}/" + f"sharded_read_test/{test_unique_name}" + ) self.table_path = self.bq_utils.join_path( self.bq_project, self.bq_dataset, @@ -34,7 +43,7 @@ def setUp(self): SELECT id as user_id, CONCAT('test_name_', CAST(id AS STRING)) as name - FROM UNNEST(GENERATE_ARRAY(0, 19)) as id + FROM UNNEST(GENERATE_ARRAY(0, {NUM_ROWS - 1})) as id """ self.bq_utils.run_query(query=create_query, labels={}) @@ -43,7 +52,7 @@ def tearDown(self): bq_table_path=self.table_path, not_found_ok=True ) - def test_constructs_when_shard_key_is_a_column(self): + def test_reads_all_rows_across_shards(self): sharded_read_config = BigQueryShardedReadConfig( shard_key="user_id", project_id=self.bq_project, @@ -52,11 +61,28 @@ def test_constructs_when_shard_key_is_a_column(self): ) # Construction validates the shard key against the real table schema. - ShardedExportRead( + sharded_read = ShardedExportRead( table_name=self.table_path, sharded_read_info=sharded_read_config, ) + # ReadFromBigQuery's EXPORT method stages query results in GCS, so the + # DirectRunner pipeline needs a project and temp_location to run against. + options = PipelineOptions() + google_cloud_options = options.view_as(GoogleCloudOptions) + google_cloud_options.project = self.bq_project + google_cloud_options.temp_location = self.temp_location + + # ShardedExportRead flattens every shard, so the read yields one dict per + # row with the table's column names as keys. We assert that all NUM_ROWS + # rows are recovered across the shards. + expected_rows = [ + {"user_id": i, "name": f"test_name_{i}"} for i in range(NUM_ROWS) + ] + with beam.Pipeline(options=options) as pipeline: + rows = pipeline | sharded_read + assert_that(rows, equal_to(expected_rows)) + def test_raises_when_shard_key_is_not_a_column(self): sharded_read_config = BigQueryShardedReadConfig( shard_key="not_a_column", From 2eb82b0fd29d9997cc43db423aaa387e957578b5 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Wed, 15 Jul 2026 20:26:16 +0000 Subject: [PATCH 09/11] Drop mock-only load-to-bigquery export unit tests 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 --- tests/unit/common/data/export_test.py | 99 --------------------------- 1 file changed, 99 deletions(-) diff --git a/tests/unit/common/data/export_test.py b/tests/unit/common/data/export_test.py index aa4c5327f..e17441746 100644 --- a/tests/unit/common/data/export_test.py +++ b/tests/unit/common/data/export_test.py @@ -9,7 +9,6 @@ import requests import torch from absl.testing import absltest -from google.cloud import bigquery from google.cloud.exceptions import GoogleCloudError from parameterized import param, parameterized @@ -19,12 +18,8 @@ _NODE_ID_KEY, _NODE_TYPE_KEY, _PREDICTION_KEY, - EMBEDDING_BIGQUERY_SCHEMA, - PREDICTION_BIGQUERY_SCHEMA, EmbeddingExporter, PredictionExporter, - load_embeddings_to_bigquery, - load_predictions_to_bigquery, ) from gigl.common.utils.retry import RetriesFailedException from tests.test_assets.test_case import TestCase @@ -385,53 +380,6 @@ def test_skips_flush_if_empty(self, mock_gcs_utils_class): exporter = EmbeddingExporter(export_dir=gcs_base_uri) exporter.flush_records() - @parameterized.expand( - [ - param( - "Test if we can load embeddings synchronously", - should_run_async=False, - ), - param( - "Test if we can load embeddings asynchronously", - should_run_async=True, - ), - ] - ) - @patch("gigl.common.data.export.BqUtils") - def test_load_embedding_to_bigquery( - self, _, mock_bq_utils_cls, should_run_async: bool - ): - gcs_folder = GcsUri("gs://test-bucket/test-folder") - project_id = "test-project" - dataset_id = "test-dataset" - table_id = "test-table" - - load_job = load_embeddings_to_bigquery( - gcs_folder, - project_id, - dataset_id, - table_id, - should_run_async=should_run_async, - ) - - # The export delegates to BqUtils, which owns quota-project resolution. - mock_bq_utils_cls.assert_called_once_with(project=project_id) - mock_load = mock_bq_utils_cls.return_value.load_files_to_bq - mock_load.assert_called_once() - _, load_kwargs = mock_load.call_args - self.assertEqual(load_kwargs["source_uris"], f"{gcs_folder.uri}/*.avro") - self.assertEqual( - load_kwargs["bq_path"], f"{project_id}.{dataset_id}.{table_id}" - ) - self.assertEqual(load_kwargs["should_run_async"], should_run_async) - job_config = load_kwargs["job_config"] - self.assertEqual(job_config.source_format, bigquery.SourceFormat.AVRO) - self.assertEqual( - job_config.write_disposition, bigquery.WriteDisposition.WRITE_APPEND - ) - self.assertEqual(list(job_config.schema), list(EMBEDDING_BIGQUERY_SCHEMA)) - self.assertIs(load_job, mock_load.return_value) - class TestPredictionsExporter(TestCase): def setUp(self): @@ -739,53 +687,6 @@ def test_skips_flush_if_empty(self, mock_gcs_utils_class): exporter = PredictionExporter(export_dir=gcs_base_uri) exporter.flush_records() - @parameterized.expand( - [ - param( - "Test if we can load predictions synchronously", - should_run_async=False, - ), - param( - "Test if we can load predictions asynchronously", - should_run_async=True, - ), - ] - ) - @patch("gigl.common.data.export.BqUtils") - def test_load_prediction_to_bigquery( - self, _, mock_bq_utils_cls, should_run_async: bool - ): - gcs_folder = GcsUri("gs://test-bucket/test-folder") - project_id = "test-project" - dataset_id = "test-dataset" - table_id = "test-table" - - load_job = load_predictions_to_bigquery( - gcs_folder, - project_id, - dataset_id, - table_id, - should_run_async=should_run_async, - ) - - # The export delegates to BqUtils, which owns quota-project resolution. - mock_bq_utils_cls.assert_called_once_with(project=project_id) - mock_load = mock_bq_utils_cls.return_value.load_files_to_bq - mock_load.assert_called_once() - _, load_kwargs = mock_load.call_args - self.assertEqual(load_kwargs["source_uris"], f"{gcs_folder.uri}/*.avro") - self.assertEqual( - load_kwargs["bq_path"], f"{project_id}.{dataset_id}.{table_id}" - ) - self.assertEqual(load_kwargs["should_run_async"], should_run_async) - job_config = load_kwargs["job_config"] - self.assertEqual(job_config.source_format, bigquery.SourceFormat.AVRO) - self.assertEqual( - job_config.write_disposition, bigquery.WriteDisposition.WRITE_APPEND - ) - self.assertEqual(list(job_config.schema), list(PREDICTION_BIGQUERY_SCHEMA)) - self.assertIs(load_job, mock_load.return_value) - if __name__ == "__main__": absltest.main() From 90b8802b8808fd5674043aaf6c8eff9ec5a70bc1 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Wed, 15 Jul 2026 20:26:23 +0000 Subject: [PATCH 10/11] Slim down bigquery client util unit tests 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 --- tests/unit/src/common/utils/bq_test.py | 44 -------------------------- 1 file changed, 44 deletions(-) diff --git a/tests/unit/src/common/utils/bq_test.py b/tests/unit/src/common/utils/bq_test.py index c710d7780..f4c6df008 100644 --- a/tests/unit/src/common/utils/bq_test.py +++ b/tests/unit/src/common/utils/bq_test.py @@ -2,7 +2,6 @@ from typing import Sequence from unittest.mock import MagicMock, patch -import google.cloud.bigquery as bigquery from parameterized import param, parameterized from gigl.env.constants import GIGL_BIGQUERY_QUOTA_PROJECT_ENV_KEY @@ -163,7 +162,6 @@ def test_env_var_sets_quota_project(self, mock_client_cls: MagicMock) -> None: ): BqUtils(project="my-project") _, call_kwargs = mock_client_cls.call_args - self.assertEqual(call_kwargs["project"], "my-project") self.assertEqual( call_kwargs["client_options"].quota_project_id, "env-quota-project" ) @@ -185,45 +183,3 @@ def test_empty_env_var_is_treated_as_unset( BqUtils(project="my-project") _, call_kwargs = mock_client_cls.call_args self.assertIsNone(call_kwargs["client_options"]) - - -@patch("gigl.src.common.utils.bq.bigquery.Client") -class LoadFilesToBqTest(TestCase): - def test_sync_load_waits_for_completion_and_returns_job( - self, mock_client_cls: MagicMock - ) -> None: - mock_load_job = mock_client_cls.return_value.load_table_from_uri.return_value - mock_load_job.output_rows = 1000 - bq_utils = BqUtils(project="my-project") - job_config = bigquery.LoadJobConfig(source_format=bigquery.SourceFormat.AVRO) - - load_job = bq_utils.load_files_to_bq( - source_uris="gs://bucket/folder/*.avro", - bq_path="my-project.my-dataset.my-table", - job_config=job_config, - should_run_async=False, - ) - - mock_client_cls.return_value.load_table_from_uri.assert_called_once_with( - source_uris="gs://bucket/folder/*.avro", - destination="my-project.my-dataset.my-table", - job_config=job_config, - ) - mock_load_job.result.assert_called_once() - self.assertIs(load_job, mock_load_job) - - def test_async_load_returns_without_waiting( - self, mock_client_cls: MagicMock - ) -> None: - mock_load_job = mock_client_cls.return_value.load_table_from_uri.return_value - bq_utils = BqUtils(project="my-project") - - load_job = bq_utils.load_files_to_bq( - source_uris="gs://bucket/folder/*.avro", - bq_path="my-project.my-dataset.my-table", - job_config=bigquery.LoadJobConfig(), - should_run_async=True, - ) - - mock_load_job.result.assert_not_called() - self.assertIs(load_job, mock_load_job) From 7f3922014dbdbe2bceac053fb4d5e3f6b32ed561 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Wed, 15 Jul 2026 20:36:36 +0000 Subject: [PATCH 11/11] Use BqUtils.join_path for the export destination path Co-Authored-By: Claude Opus 4.8 --- gigl/common/data/export.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/gigl/common/data/export.py b/gigl/common/data/export.py index de55ee0c6..874be4b99 100644 --- a/gigl/common/data/export.py +++ b/gigl/common/data/export.py @@ -380,12 +380,9 @@ def _load_records_to_bigquery( schema=schema, ) - # Build the destination path directly rather than via BqUtils.join_path, - # which asserts on dot count and would reject domain-scoped project IDs - # (e.g. "example.com:project"). return bq_utils.load_files_to_bq( source_uris=os.path.join(gcs_folder.uri, "*.avro"), - bq_path=f"{project_id}.{dataset_id}.{table_id}", + bq_path=bq_utils.join_path(project_id, dataset_id, table_id), job_config=job_config, should_run_async=should_run_async, )