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/integration/common/beam/__init__.py b/tests/integration/common/beam/__init__.py new file mode 100644 index 000000000..e69de29bb 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()