adding transformations for bq table

Author: balit-raibot

scripts/noaa_gfs/dc_bq_ingest.py

@@ -23,39 +23,72 @@
 
 # --- FLAG DEFINITIONS ---
 FLAGS = flags.FLAGS
-flags.DEFINE_string('project_id', 'datcom-store', 'GCP Project ID.')
+flags.DEFINE_string('project_id', 'datcom-external', 'GCP Project ID.')
 flags.DEFINE_string('bucket_name', 'datcom-prod-imports', 'GCS Bucket containing the CSVs.')
 flags.DEFINE_string('gcs_prefix', 'scripts/noaa_gfs/NOAA_GlobalForecastSystem/output/', 'GCS prefix (folder path).')
-flags.DEFINE_string('dataset_id', 'dc_kg_latest', 'BigQuery Dataset ID.')
-flags.DEFINE_string('table_id', 'noaa_gfs_pdp', 'BigQuery Table ID.')
+flags.DEFINE_string('dataset_id', 'data_commons_noaa_gfs', 'BigQuery Dataset ID.')
+flags.DEFINE_string('table_id', 'Observation', 'BigQuery Table ID.')
+flags.DEFINE_string('staging_table_id', 'Observation_Staging', 'Temporary Staging Table ID.')
 
-def upload_gcs_to_bq(bq_client, gcs_uri):
+def run_mapping_query(bq_client):
     """
-    Triggers a BigQuery load job directly from a GCS URI.
+    Executes the SQL transformation to map data from Staging to Final table.
     """
-    table_ref = f"{FLAGS.project_id}.{FLAGS.dataset_id}.{FLAGS.table_id}"
+    final_table = f"{FLAGS.project_id}.{FLAGS.dataset_id}.{FLAGS.table_id}"
+    staging_table = f"{FLAGS.project_id}.{FLAGS.dataset_id}.{FLAGS.staging_table_id}"
+
+    query = f"""
+    INSERT INTO `{final_table}` (
+        observation_about,
+        variable_measured,
+        value,
+        observation_date,
+        measurement_method,
+        unit,
+        prov_id
+    )
+    SELECT 
+        placeName,
+        variableMeasured,
+        CAST(value AS STRING),
+        CAST(observationDate AS STRING),
+        measurementMethod,
+        unit,
+        'dc/base/NOAA_GlobalForecastSystem'
+    FROM `{staging_table}`;
+    """
+    
+    try:
+        logging.info("Starting transformation query...")
+        query_job = bq_client.query(query)
+        query_job.result() # Wait for completion
+        
+        # Optional: Truncate staging table after successful migration
+        bq_client.query(f"TRUNCATE TABLE `{staging_table}`").result()
+        logging.info("Transformation complete and staging table cleared.")
+        return True
+    except Exception as e:
+        logging.error(f"Mapping query failed: {e}")
+        return False
+
+def upload_gcs_to_staging(bq_client, gcs_uri):
+    """
+    Loads raw CSV data into the Staging table.
+    """
+    table_ref = f"{FLAGS.project_id}.{FLAGS.dataset_id}.{FLAGS.staging_table_id}"
 
-    # Configure the load job
     job_config = bigquery.LoadJobConfig(
         source_format=bigquery.SourceFormat.CSV,
-        skip_leading_rows=1,      # Skip the header row
-        autodetect=True,          # Automatically infer types
+        skip_leading_rows=1,
+        autodetect=True,
+        # WRITE_APPEND used here to collect all CSVs before the final SQL transformation
         write_disposition=bigquery.WriteDisposition.WRITE_APPEND,
     )
 
     try:
-        logging.info(f"Starting load job for: {gcs_uri}")
-        # Direct URI load
-        load_job = bq_client.load_table_from_uri(
-            gcs_uri, 
-            table_ref, 
-            job_config=job_config
-        )
-        
-        # Wait for the job to complete
+        logging.info(f"Loading to staging: {gcs_uri}")
+        load_job = bq_client.load_table_from_uri(gcs_uri, table_ref, job_config=job_config)
         load_job.result() 
-            
-        logging.info(f"Successfully loaded rows from {gcs_uri}. Job ID: {load_job.job_id}")
         return True
     except Exception as e:
         logging.error(f"Failed to load {gcs_uri}: {e}")
@@ -80,13 +113,17 @@ def main(argv):
 
     logging.info(f"Found {len(csv_uris)} files in GCS for ingestion.")
 
-    # Loop and Trigger Load Jobs
+    # Step 1: Bulk Load everything into Staging
     success_count = 0
     for uri in csv_uris:
-        if upload_gcs_to_bq(bq_client, uri):
+        if upload_gcs_to_staging(bq_client, uri):
             success_count += 1
 
     logging.info(f"Ingestion batch complete. {success_count}/{len(csv_uris)} URIs processed.")
 
+    # Step 2: Run Mapping SQL if at least some files loaded
+    if success_count > 0:
+        run_mapping_query(bq_client)
+
 if __name__ == "__main__":
     app.run(main)