From 125c73980ab5c2aeab5219ffffafc3127a89dcee Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Tue, 21 Jul 2026 16:03:34 +0000 Subject: [PATCH 1/2] Use BqUtils for sharded read schema validation --- gigl/common/beam/sharded_read.py | 9 ++-- tests/unit/common/beam/__init__.py | 0 tests/unit/common/beam/sharded_read_test.py | 53 +++++++++++++++++++++ 3 files changed, 56 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..877725171 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,8 @@ 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] + bq_utils = BqUtils() + column_names = list(bq_utils.fetch_bq_table_schema(bq_table=table_name)) 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..f04cf7e41 --- /dev/null +++ b/tests/unit/common/beam/sharded_read_test.py @@ -0,0 +1,53 @@ +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.assert_called_once_with() + 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 ad44a913c73ce652326cb4cdfe70309bc704a191 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Tue, 21 Jul 2026 16:17:34 +0000 Subject: [PATCH 2/2] Test sharded reads against BigQuery --- .../common/beam/__init__.py | 0 .../common/beam/sharded_read_test.py | 121 ++++++++++++++++++ tests/unit/common/beam/sharded_read_test.py | 53 -------- 3 files changed, 121 insertions(+), 53 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..a51069d0b --- /dev/null +++ b/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() 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 f04cf7e41..000000000 --- a/tests/unit/common/beam/sharded_read_test.py +++ /dev/null @@ -1,53 +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.assert_called_once_with() - 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()