Pubsub to BQ load using beam

Question

I am trying to create beam code which will just load the data from pubsub to bq, when i am running the same code from colab its working perfectly but when i am creating template or dataflow as running i am getting issues not define etc.

Error = Invalid JSON payload received. Unknown name "json" at 'rows[0]': Proto field is not repeating, cannot start list. [while running 'Write Raw Data to BQ/_StreamToBigQuery/StreamInsertRows/ParDo(BigQueryWriteFn)-ptransform-60']

line 179, in process UnboundLocalError: local variable 'job' referenced before assignment [while running 'Write Raw Data to Big Query-ptransform-63'] passed through: ==> dist_proc/dax/workflow/worker/fnapi_service_impl.cc:922

Code :

            import argparse
            import json
            import os
            import logging
            import pandas as pd
            import apache_beam as beam
            from google.cloud import bigquery, pubsub_v1
            from google.cloud.bigquery import table
            import google.cloud.storage.client as gcs
            from apache_beam.options.pipeline_options import PipelineOptions, StandardOptions, GoogleCloudOptions, SetupOptions
            import apache_beam.transforms.window as window


            serviceAccount = '/content/test-project-0c3a740e58c8.json'
            os.environ["GOOGLE_APPLICATION_CREDENTIALS"]= serviceAccount

            logging.basicConfig(level=logging.INFO)
            logging.getLogger().setLevel(logging.INFO)

            project_id = "test-project"
            bigquery_dataset = "bq_load"
            bigquery_message_table = "message_table"
            bigquery_error_table = "Error_Table"
            bigquery_message_agg_table = "message_table_aggri"
            subscription_id = "topic_message-sub"

            INPUT_SUBSCRIPTION = pubsub_v1.SubscriberClient().subscription_path(project_id, subscription_id)
            WINDOW_SIZE = 100

            client = bigquery.Client(project = project_id)
            bq_dataset  = client.dataset(bigquery_dataset)
            bq_msg_table = bq_dataset.table(bigquery_message_table)
            bq_err_table = bq_dataset.table(bigquery_error_table)
            bq_agg_table = bq_dataset.table(bigquery_message_agg_table)

            message_table_schema = {
                      'name': 'MessageType',
                      'type': 'STRING',
                      'mode': 'NULLABLE'
                      }, {
                      'name': 'Location',
                      'type': 'STRING',
                      'mode': 'NULLABLE'
                      }, {
                      'name': 'Building',
                      'type': 'STRING',
                      'mode': 'NULLABLE'
                      }, {
                      'name': 'DateTime',
                      'type': 'DATETIME',
                      'mode': 'NULLABLE'
                      }, {
                      'name': 'ID',
                      'type': 'STRING',
                      'mode': 'NULLABLE'
                      }, {
                      'name': 'TID',
                      'type': 'STRING',
                      'mode': 'NULLABLE'
                      }

            error_table_schema = {
                      'name': 'Message',
                      'type': 'STRING',
                      'mode': 'NULLABLE'
                      }, {
                      'name': 'Error',
                      'type': 'STRING',
                      'mode': 'NULLABLE'
                      }

            table_schema_agri = {
                      'name': 'MessageType',
                      'type': 'STRING',
                      'mode': 'NULLABLE'
                      }, {
                      'name': 'Location',
                      'type': 'STRING',
                      'mode': 'NULLABLE'
                      }, {
                      'name': 'Building',
                      'type': 'STRING',
                      'mode': 'NULLABLE'
                      }, {
                      'name': 'ID',
                      'type': 'STRING',
                      'mode': 'NULLABLE'
                      }, {
                      'name': 'TID',
                      'type': 'STRING',
                      'mode': 'NULLABLE'
                      }, {
                      'name': 'no_of_rows',
                      'type': 'NUMERIC',
                      'mode': 'NULLABLE'
                      }

            def format_schema(schema):
              formatted_schema = []
              column_list = []
              for row in schema:
                  formatted_schema.append(bigquery.SchemaField(row['name'], row['type'], row['mode']))
                  column_list.append(row['name'])
              return formatted_schema, set(column_list)

            message_job_config = bigquery.LoadJobConfig()
            message_job_config.source_format = bigquery.SourceFormat.NEWLINE_DELIMITED_JSON
            message_job_config.schema, message_column_list = format_schema(message_table_schema)
            print(message_column_list)

            error_job_config = bigquery.LoadJobConfig()
            error_job_config.source_format = bigquery.SourceFormat.NEWLINE_DELIMITED_JSON
            error_job_config.schema, error_column_list = format_schema(error_table_schema)
            print(error_column_list)

            agg_job_config = bigquery.LoadJobConfig()
            agg_job_config.source_format = bigquery.SourceFormat.NEWLINE_DELIMITED_JSON
            agg_job_config.schema, agg_column_list = format_schema(table_schema_agri)
            print(agg_column_list)

            class ProcessMessage(beam.DoFn):

                def process(self, message):
                  json_object = []
                  try:
                    message_str = message.decode("utf-8")
                    json_message = eval(message_str)

                    if isinstance(json_message, dict) and set(message_column_list).issubset(json_message.keys()):
                      print('{}'.format(("Formatted message", json_message)))
                      yield tuple(["Formatted message", json_message])
                    else:
                      print('{}'.format(("Error_Message", {"Original_Message": json_message, "Error": 'Unable to map input to table schema'})))
                      yield tuple(["Error_Message", {"Original_Message": json_message, "Error": 'Unable to map input to table schema'}])
                  except Exception as e:
                      print('{}'.format(("Error_Message", {"Original_Message": message_str if message_str else message, "Error": str(e)})))
                      yield tuple(["Error_Message", {"Original_Message": message_str if message_str else message, "Error": str(e)}])

            class WriteDataframeToBQ(beam.DoFn):
                
                def process(self, message):
                  global project_id 
                  global bigquery_dataset
                  global bigquery_message_table
                  global bigquery_error_table
                  global bigquery_message_agg_table
                  global subscription_id
                  global client
                  global bq_dataset
                  global bq_msg_table
                  global bq_err_table
                  global bq_agg_table
                  if message[0] == 'Formatted message':
                    try:
                      print('Now inserting:', message[1])
                      job = client.load_table_from_json(message[1], BQ_MSG_TABLE, job_config = message_job_config)
                      job.result()  # Waits for the job to complete.
                      print('Total {} messages successfully written to BigQuery Messsage Table'.format(len(message[1])))
                      
                      # Creating aggregation dataframe 
                      df = pd.DataFrame(message[1]) 
                      count_df= df.groupby(['MessageType','Location','Building','ID','TID']).size().reset_index(name='no_of_rows')
                      job = client.load_table_from_json(count_df.to_dict('records'), BQ_AGG_TABLE, job_config = agg_job_config)
                      job.result()  # Waits for the job to complete.
                      print('Total {} messages successfully written to BigQuery Agree Messsage Table'.format(str(count_df.shape[0])))
                     
                    except Exception as e:
                     # print(job.errors)
                      raise e
                  else:
                    
                    try:          
                      error_rows = [{"Message": json.dumps(m["Original_Message"]) if isinstance(m["Original_Message"], dict) else str(m["Original_Message"]), "Error": m["Error"]} for m in message[1]]
                      print('Now inserting:', error_rows)
                      job = client.load_table_from_json(error_rows, BQ_ERR_TABLE, job_config = error_job_config)
                      job.result()  # Waits for the job to complete.
                      print('Total {} messages successfully written to BigQuery Error Table'.format(len(error_rows)))
                    except Exception as e:
                      #print(job.errors)
                      raise e       

            def run():  


                pipeline_options = {
                  'project': 'test-project',        
                  'runner': 'DataflowRunner',
                  'region': 'us-east1',
                  'staging_location': 'gs://bucket12345/tmp/',
                  'temp_location': 'gs://bucket12345/tmp/',
                  'template_location': 'gs://bucket12345/template/dataflow_test_template',
                  
                   'streaming': True
                  }
                pipeline_options = PipelineOptions.from_dictionary(pipeline_options)
                    table = 'test-project:bq_load.message_table'

                with beam.Pipeline(options=pipeline_options) as p:
                   (     p
                      | "ReadFromPubSub" >> beam.io.gcp.pubsub.ReadFromPubSub(
                          subscription=INPUT_SUBSCRIPTION, timestamp_attribute=None
                      )
                      | "Process Message" >> beam.ParDo(ProcessMessage())
                      | 'Fixed-size windows' >> beam.WindowInto(window.FixedWindows(WINDOW_SIZE))
                      | beam.GroupByKey()
                      | "Write Raw Data to Big Query" >> beam.ParDo(WriteDataframeToBQ())
                      | "Write Raw Data to BQ" >> beam.io.WriteToBigQuery(
                          table,custom_gcs_temp_location = 'gs://bucket12345/template/dataflow_test_template')
                  )

            run()