diff --git a/weather_mv/loader_pipeline/bq.py b/weather_mv/loader_pipeline/bq.py index 10f5aeb..6328a15 100644 --- a/weather_mv/loader_pipeline/bq.py +++ b/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 (..). + 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 (..). ") + 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) ) diff --git a/weather_mv/loader_pipeline/pipeline.py b/weather_mv/loader_pipeline/pipeline.py index e8d02cc..85d9230 100644 --- a/weather_mv/loader_pipeline/pipeline.py +++ b/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: