Refactoring table creation in Big Query Pipeline

weather_mv/loader_pipeline/bq.py

@@ -17,6 +17,7 @@
 import json
 import logging
 import os
+import re
 import typing as t
 from pprint import pformat
 
@@ -130,6 +131,12 @@ def validate_arguments(cls, known_args: argparse.Namespace, pipeline_args: t.Lis
         pipeline_options = PipelineOptions(pipeline_args)
         pipeline_options_dict = pipeline_options.get_all_options()
 
+        if known_args.output_table:
+            # checking if the output table is in format (<project>.<dataset>.<table>).
+            output_table_pattern = r'^[\w-]+\.[\w-]+\.[\w-]+$'
+            if not bool(re.match(output_table_pattern, known_args.output_table)):
+                raise RuntimeError("output_table is not in correct format (<project>.<dataset_id>.<table_id>). ")
+
         if known_args.area:
             assert len(known_args.area) == 4, 'Must specify exactly 4 lat/long values for area: N, W, S, E boundaries.'
 
@@ -149,10 +156,20 @@ def validate_arguments(cls, known_args: argparse.Namespace, pipeline_args: t.Lis
             logger.info('Region validation completed successfully.')
 
     def __post_init__(self):
-        """Initializes Sink by creating a BigQuery table based on user input."""
+        """Initializes BigQuery table based on user input."""
+        self.project, self.dataset_id, self.table_id = self.output_table.split('.')
+        self.table = None
+
         if self.zarr:
             self.xarray_open_dataset_kwargs = self.zarr_kwargs
-        with open_dataset(self.first_uri, self.xarray_open_dataset_kwargs,
+
+    def create_bq_table(self, uri: str) -> str:
+        """Create a big query table for the first uri. After table is created, subsequent uris are returned."""
+        # Skip table creation.
+        if self.table:
+            return uri
+
+        with open_dataset(uri, self.xarray_open_dataset_kwargs,
                           self.disable_grib_schema_normalization, self.tif_metadata_for_datetime,
                           is_zarr=self.zarr) as open_ds:
             # Define table from user input
@@ -170,12 +187,13 @@ def __post_init__(self):
         if self.dry_run:
             logger.debug('Created the BigQuery table with schema...')
             logger.debug(f'\n{pformat(table_schema)}')
-            return
+            return uri
 
         # Create the table in BigQuery
         try:
             table = bigquery.Table(self.output_table, schema=table_schema)
             self.table = bigquery.Client().create_table(table, exists_ok=True)
+            return uri
         except Exception as e:
             logger.error(f'Unable to create table in BigQuery: {e}')
             raise
@@ -243,6 +261,7 @@ def expand(self, paths):
         """Extract rows of variables from data paths into a BigQuery table."""
         extracted_rows = (
                 paths
+                | 'CreateTable' >> beam.Map(self.create_bq_table)
                 | 'PrepareCoordinates' >> beam.FlatMap(self.prepare_coordinates)
                 | beam.Reshuffle()
                 | 'ExtractRows' >> beam.FlatMapTuple(self.extract_rows)
@@ -252,9 +271,9 @@ def expand(self, paths):
             (
                     extracted_rows
                     | 'WriteToBigQuery' >> WriteToBigQuery(
-                        project=self.table.project,
-                        dataset=self.table.dataset_id,
-                        table=self.table.table_id,
+                        project=self.project,
+                        dataset=self.dataset_id,
+                        table=self.table_id,
                         write_disposition=BigQueryDisposition.WRITE_APPEND,
                         create_disposition=BigQueryDisposition.CREATE_NEVER)
             )

weather_mv/loader_pipeline/pipeline.py

@@ -47,11 +47,15 @@ def pattern_to_uris(match_pattern: str, is_zarr: bool = False) -> t.Iterable[str
 
 def pipeline(known_args: argparse.Namespace, pipeline_args: t.List[str]) -> None:
     all_uris = list(pattern_to_uris(known_args.uris, known_args.zarr))
-    if not all_uris:
+    if not all_uris and not known_args.topic:
         raise FileNotFoundError(f"File pattern '{known_args.uris}' matched no objects")
 
     # First URI is useful to get an example data shard. It also can be a Zarr path.
-    known_args.first_uri = next(iter(all_uris))
+    if all_uris:
+        known_args.first_uri = next(iter(all_uris))
+    else:
+        # If it's a streaming pipeline, it will allow first_uri to be empty.
+        known_args.first_uri = None
 
     with beam.Pipeline(argv=pipeline_args) as p:
         if known_args.topic or known_args.subscription: