loader.py

"""Data Loader for TransitPulse Cloud Mode.
Uploads raw/aggregated parquets to Google Cloud Storage and loads them into BigQuery.
"""

from __future__ import annotations

import os
from pathlib import Path

from config import CFG


def upload_to_gcs(local_file_path: Path, gcs_blob_name: str) -> str:
    """Uploads a local file to the configured GCS bucket."""
    from google.cloud import storage
    
    if not CFG.gcs_bucket:
        raise ValueError("GCS_BUCKET environment variable is not configured.")
        
    client = storage.Client(project=CFG.gcp_project)
    bucket = client.bucket(CFG.gcs_bucket)
    
    # Create bucket if it doesn't exist
    if not bucket.exists():
        print(f"Creating GCS bucket: {CFG.gcs_bucket}...")
        bucket.create(location="US")
        
    blob = bucket.blob(gcs_blob_name)
    blob.upload_from_filename(str(local_file_path))
    
    gcs_uri = f"gs://{CFG.gcs_bucket}/{gcs_blob_name}"
    print(f"Successfully uploaded {local_file_path.name} to {gcs_uri}")
    return gcs_uri


def load_parquet_to_bigquery(gcs_uri: str, table_id: str) -> None:
    """Loads a parquet file from GCS into a BigQuery table, creating/overwriting it."""
    from google.cloud import bigquery
    
    if not CFG.gcp_project or not CFG.bq_dataset:
        raise ValueError("GCP_PROJECT or BQ_DATASET environment variables are not configured.")
        
    client = bigquery.Client(project=CFG.gcp_project)
    
    # Ensure dataset exists
    dataset_ref = client.dataset(CFG.bq_dataset)
    try:
        client.get_dataset(dataset_ref)
    except Exception:
        print(f"Creating BigQuery dataset: {CFG.bq_dataset}...")
        dataset = bigquery.Dataset(dataset_ref)
        dataset.location = "US"
        client.create_dataset(dataset)
        
    full_table_ref = f"{CFG.gcp_project}.{CFG.bq_dataset}.{table_id}"
    
    job_config = bigquery.LoadJobConfig(
        source_format=bigquery.SourceFormat.PARQUET,
        write_disposition=bigquery.WriteDisposition.WRITE_TRUNCATE, # Overwrite table
    )
    
    print(f"Loading {gcs_uri} into BigQuery table {full_table_ref}...")
    load_job = client.load_table_from_uri(gcs_uri, full_table_ref, job_config=job_config)
    
    # Wait for completion
    load_job.result()
    print(f"Successfully loaded BigQuery table {full_table_ref}.")


def push_to_cloud() -> None:
    """Pushes all aggregates generated by the pipeline to GCP."""
    print("=== Commencing Cloud Sync to Google Cloud Storage & BigQuery ===")
    
    aggregates = {
        "route_scores": CFG.output_dir / "route_scores.parquet",
        "daily_segment_metrics": CFG.output_dir / "daily_segment_metrics.parquet",
        "anomaly_events": CFG.output_dir / "anomaly_events.parquet"
    }
    
    for table_name, file_path in aggregates.items():
        if not file_path.exists():
            print(f"Error: Aggregate file {file_path} not found. Run pipeline first.")
            continue
            
        # 1. Upload to GCS
        blob_name = f"aggregates/{table_name}/{file_path.name}"
        gcs_uri = upload_to_gcs(file_path, blob_name)
        
        # 2. Ingest into BigQuery table
        load_parquet_to_bigquery(gcs_uri, table_name)
        
    print("=== Cloud Sync Complete! ===")


if __name__ == "__main__":
    push_to_cloud()