insert_to_clickhouse.py

Go to the documentation of this file.

# ===========================================
# insert_to_clickhouse.py
# ===========================================
 

 
 
import argparse
import polars as pl
from clickhouse_driver import Client
from pipeline.schema import SCHEMA
from pipeline.utils import safe_vector_cast
from scripts.config import *
 
 
def insert_batch(batch_id: str) -> None:
    """!Filters and inserts a batch of records into ClickHouse.
 
    Loads data from the Parquet file at DB_PATH, filters by BatchID,
    and inserts the resulting records into the `benchmark.performance` table.
 
    @param batch_id The BatchID to filter the dataset on.
 
    @throws Exception If ClickHouse insert fails.
    """
    df = pl.read_parquet(DB_PATH)
    df = df.filter(pl.col("BatchID") == batch_id)
 
    # Optional: enforce schema casting
    df = safe_vector_cast(df, SCHEMA)
 
    client = Client(
        host=CLICKHOUSE_HOST,
        port=CLICKHOUSE_TCP_PORT,
        user=CLICKHOUSE_USER,
        password=CLICKHOUSE_PASSWORD
    )
 
    records = df.to_dicts()
 
    try:
        client.execute("INSERT INTO benchmark.performance VALUES", records)
    except Exception as e:
        print(f"[ERROR] Error inserting records into ClickHouse: {e}")
        raise
 
    print(f"[INFO] Inserted {len(records)} records into ClickHouse for batch '{batch_id}'.")
 
 
def main():
    """!CLI entrypoint for inserting a batch into ClickHouse.
 
    Parses --batchid from command-line arguments and performs the insert.
    """
    parser = argparse.ArgumentParser(description="Insert benchmarking logs into ClickHouse")
    parser.add_argument("--batchid", type=str, required=True, help="Batch ID to ingest")
    args = parser.parse_args()
 
    insert_batch(args.batchid)
 
 
if __name__ == "__main__":
    main()