New add_load_date_suffix import option. (GoogleCloudPlatform#681)

With the new import option , the pipeline will now create
a new BigQuery table for each new load date. This helps handle API incompatible
schema changes better as each day's imports only need to be consistent rather
then all prior imports.

Co-authored-by: Abdel SGHIOUAR <abdelfettah@google.com>

tools/asset-inventory/asset_inventory/import_pipeline.py

@@ -32,8 +32,8 @@
 from datetime import datetime
 import json
 import logging
-import random
 import pprint
+import random
 
 import apache_beam as beam
 from apache_beam.io import ReadFromText
@@ -294,10 +294,17 @@ def finish_bundle(self):
 class BigQueryDoFn(beam.DoFn):
     """Superclass for a DoFn that requires BigQuery dataset information."""
 
-    def __init__(self, dataset):
+    def __init__(self, dataset, add_load_date_suffix, load_time):
         if isinstance(dataset, string_types):
             dataset = StaticValueProvider(str, dataset)
         self.dataset = dataset
+        if isinstance(add_load_date_suffix, string_types):
+            add_load_date_suffix = StaticValueProvider(
+                str, add_load_date_suffix)
+        self.add_load_date_suffix = add_load_date_suffix
+        if isinstance(load_time, string_types):
+            load_time = StaticValueProvider(str, load_time)
+        self.load_time = load_time
         self.bigquery_client = None
         self.dataset_location = None
         self.load_jobs = {}
@@ -316,7 +323,12 @@ def get_dataset_location(self):
         return None
 
     def asset_type_to_table_name(self, asset_type):
-        return asset_type.replace('.', '_').replace('/', '_')
+        suffix = ''
+        add_load_date_suffix = self.add_load_date_suffix.get()
+        if (add_load_date_suffix and
+            add_load_date_suffix.lower() in ('yes', 'true', 't', '1')):
+            suffix = '_' + self.load_time.get()[0:10].replace('-', '')
+        return asset_type.replace('.', '_').replace('/', '_') + suffix
 
     def start_bundle(self):
         if not self.bigquery_client:
@@ -332,11 +344,12 @@ class DeleteDataSetTables(BigQueryDoFn):
     dataset before loading so that no old asset types remain.
     """
 
-    def __init__(self, dataset, write_disposition):
+    def __init__(self, dataset, add_load_date_suffix, load_time,
+                 write_disposition):
         # Can't use super().
         # https://issues.apache.org/jira/browse/BEAM-6158?focusedCommentId=16919945
         # super(DeleteDataSetTables, self).__init__(dataset)
-        BigQueryDoFn.__init__(self, dataset)
+        BigQueryDoFn.__init__(self, dataset, add_load_date_suffix, load_time)
         if isinstance(write_disposition, string_types):
             write_disposition = StaticValueProvider(str, write_disposition)
         self.write_disposition = write_disposition
@@ -362,14 +375,11 @@ class LoadToBigQuery(BigQueryDoFn):
     this must be done within the workers.
     """
 
-    def __init__(self, dataset, load_time):
+    def __init__(self, dataset, add_load_date_suffix, load_time):
         # Can't use super().
         # https://issues.apache.org/jira/browse/BEAM-6158?focusedCommentId=16919945
         # super(LoadToBigQuery, self).__init__(dataset)
-        BigQueryDoFn.__init__(self, dataset)
-        if isinstance(load_time, string_types):
-            load_time = StaticValueProvider(str, load_time)
-        self.load_time = load_time
+        BigQueryDoFn.__init__(self, dataset, add_load_date_suffix, load_time)
 
     def to_bigquery_schema(self, fields):
         """Convert list of dicts into `bigquery.SchemaFields`."""
@@ -463,6 +473,11 @@ def _add_argparse_args(cls, parser):
             default=datetime.now().isoformat(),
             help='Load time of the data (YYYY-MM-DD[HH:MM:SS])).')
 
+        parser.add_value_provider_argument(
+            '--add_load_date_suffix',
+            default='False',
+            help='If the load date [YYYYMMDD] is added as a table suffix.')
+
         parser.add_value_provider_argument(
             '--dataset', help='BigQuery dataset to load to.')
 
@@ -508,9 +523,12 @@ def run(argv=None):
          WriteToGCS(options.stage, options.load_time))
      | 'group_written_objects_by_key' >> beam.GroupByKey()
      | 'delete_tables' >> beam.ParDo(
-         DeleteDataSetTables(options.dataset, options.write_disposition))
+         DeleteDataSetTables(options.dataset, options.add_load_date_suffix,
+                             options.load_time,
+                             options.write_disposition))
      | 'load_to_bigquery' >> beam.ParDo(
-         LoadToBigQuery(options.dataset, options.load_time),
+         LoadToBigQuery(options.dataset, options.add_load_date_suffix,
+                        options.load_time),
          beam.pvalue.AsDict(schemas)))
 
     return p.run()

tools/asset-inventory/asset_inventory/import_pipeline_metadata

@@ -31,6 +31,13 @@
     "help_text": "Either 'WRITE_APPEND' or WRITE_EMPTY.",
     "regexes": ["^(WRITE_APPEND|WRITE_EMPTY)$"]
   },
+  {
+    "name": "add_load_date_suffix",
+    "label": "If the load date [YYYYMMDD] is added as a table suffix.",
+    "help_text": "Either 'true' or 'false'.",
+    "isOptional": true,
+    "regexes": ["^(true|false)$"]
+  },
   {
     "name": "dataset",
     "label": "BigQuery dataset to load to.",

tools/asset-inventory/asset_inventory/pipeline_runner.py

@@ -67,7 +67,8 @@ def wait_on_pipeline_job(df_service, pipeline_job):
 
 def run_pipeline_template(dataflow_project, template_region, template_location,
                           input_location, group_by, write_disposition, dataset,
-                          stage, load_time, num_shards, runtime_environment):
+                          stage, load_time, num_shards, add_load_date_suffix,
+                          runtime_environment):
     """Invoke the suplied pipeline template.
 
     Args:
@@ -81,6 +82,7 @@ def run_pipeline_template(dataflow_project, template_region, template_location,
         stage: GCS path to write BigQuery load files.
         load_time: Timestamp or date to load data with.
         num_shards: Shards for for each asset type.
+        add_load_date_suffix: If the load date is added as a table suffix.
         runtime_environment: Dict  suppling other runtime overrides.
     Returns:
         End state of the pipline and job object.
@@ -100,6 +102,7 @@ def run_pipeline_template(dataflow_project, template_region, template_location,
             'group_by': group_by,
             'write_disposition': write_disposition,
             'num_shards': num_shards,
+            'add_load_date_suffix': add_load_date_suffix,
             'dataset': dataset,
         },
         'environment': runtime_environment
@@ -118,7 +121,8 @@ def run_pipeline_template(dataflow_project, template_region, template_location,
 
 def run_pipeline_beam_runner(pipeline_runner, dataflow_project, input_location,
                              group_by, write_disposition, dataset, stage,
-                             load_time, num_shards, pipeline_arguments):
+                             load_time, num_shards, add_load_date_suffix,
+                             pipeline_arguments):
     """Invokes the pipeline with a beam runner.
 
     Only tested with the dataflow and direct runners.
@@ -133,6 +137,7 @@ def run_pipeline_beam_runner(pipeline_runner, dataflow_project, input_location,
         stage: GCS path to write BigQuery load files.
         load_time: Timestamp to add to data during during BigQuery load.
         num_shards: Shards for for each asset type.
+        add_load_date_suffix: If the load date is added as a table suffix.
         pipeline_arguments: List of additional runner arguments.
     Returns:
         The end state of the pipeline run (a string), and PipelineResult.
@@ -155,6 +160,7 @@ def run_pipeline_beam_runner(pipeline_runner, dataflow_project, input_location,
         '--group_by': group_by,
         '--write_disposition': write_disposition,
         '--num_shards': num_shards,
+        '--add_load_date_suffix': add_load_date_suffix,
         '--dataset': dataset,
         '--stage': stage,
         '--runner': pipeline_runner

tools/asset-inventory/gae/config.yaml

@@ -83,5 +83,8 @@ import_write_disposition: WRITE_APPEND
 # *=1,resource=100,google.cloud.bigquery.Table=100
 import_num_shards: "*=1"
 
+# If the load date [YYYYMMDD] is added as a table suffix.
+import_add_load_date_suffix: False
+
 # If we are running on App Engine and only want to be invoked by cron tasks for security reasons.
 restrict_to_cron_tasks: True

tools/asset-inventory/gae/main.py

@@ -73,6 +73,7 @@ def get_import_arguments():
             CONFIG.import_template_region, CONFIG.import_template_location,
             '{}/*.json'.format(CONFIG.gcs_destination), CONFIG.import_group_by,
             CONFIG.import_write_disposition, CONFIG.import_dataset,
+            CONFIG.import_add_load_date_suffix,
             CONFIG.import_stage,
             datetime.datetime.now().isoformat(),
             CONFIG.import_num_shards,
@@ -87,7 +88,8 @@ def run_import():
     import_arguments = get_import_arguments()
     logging.info('running import %s', import_arguments)
     (runner, dataflow_project, template_region, template_location,
-     input_location, group_by, write_disposition, dataset, stage, load_time,
+     input_location, group_by, write_disposition, dataset,
+     add_load_date_suffix, stage, load_time,
      num_shards, pipeline_arguments,
      pipeline_runtime_environment) = import_arguments
 
@@ -96,12 +98,13 @@ def run_import():
             dataflow_project, template_region,
             template_location, input_location,
             group_by, write_disposition, dataset, stage,
-            load_time, num_shards, pipeline_runtime_environment)
+            load_time, num_shards, add_load_date_suffix,
+            pipeline_runtime_environment)
     else:
         return pipeline_runner.run_pipeline_beam_runner(
             runner, dataflow_project, input_location,
             group_by, write_disposition, dataset, stage, load_time,
-            num_shards, pipeline_arguments)
+            num_shards, add_load_date_suffix, pipeline_arguments)
 
 
 @app.route('/export_import')