BqUtils: support a BigQuery quota project override + consolidate client construction

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(

gigl/common/data/export.py

@@ -25,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()
@@ -340,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,
@@ -360,20 +360,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.
 
+    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
         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
-    bigquery_client = bigquery.Client(project=project_id)
-
-    # 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)
 
     # Configure the load job
     job_config = bigquery.LoadJobConfig(
@@ -382,24 +380,13 @@ def _load_records_to_bigquery(
         schema=schema,
     )
 
-    load_job = bigquery_client.load_table_from_uri(
+    return bq_utils.load_files_to_bq(
         source_uris=os.path.join(gcs_folder.uri, "*.avro"),
-        destination=table_ref,
+        bq_path=bq_utils.join_path(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,
@@ -426,11 +413,15 @@ def load_embeddings_to_bigquery(
         table_id (str): The BigQuery table ID.
         should_run_async (bool): Whether loading to bigquery step should happen asynchronously. Defaults to False.
 
+    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
         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,
@@ -467,11 +458,15 @@ def load_predictions_to_bigquery(
         table_id (str): The BigQuery table ID.
         should_run_async (bool): Whether loading to bigquery step should happen asynchronously. Defaults to False.
 
+    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
         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,

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

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,

tests/integration/common/beam/sharded_read_test.py

@@ -0,0 +1,121 @@
+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 (
+    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
+
+NUM_ROWS = 20
+
+
+class ShardedReadIntegrationTest(TestCase):
+    """Integration tests for ShardedExportRead 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.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,
+            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, {NUM_ROWS - 1})) 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_reads_all_rows_across_shards(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.
+        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",
+            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()