scripts.setup Namespace Reference

Functions
	log (str msg)
	Prints an info message to stdout.
	err (str msg)
	Prints an error message to stdout.
	run_command (str cmd)
	Executes a shell command with logging.
Client	wait_for_clickhouse ()
	Waits for ClickHouse server to become ready, retries for up to 30 attempts.
	setup_clickhouse (Client client)
	Creates the ClickHouse database and performance table if they don't exist.
	load_db_to_clickhouse (Client client, Path db_path)
	Wipes previous data and loads data from a Parquet file into ClickHouse.
	main ()
	CLI entrypoint.

Function Documentation

err()

scripts.setup.err

(

str

msg

)

Prints an error message to stdout.

Parameters

msg	The message string.

Definition at line 58 of file setup.py.

def err(msg: str):
    """!Prints an error message to stdout.
    @param msg The message string.
    """
    print(f"[ERROR] {msg}")
 

load_db_to_clickhouse()

scripts.setup.load_db_to_clickhouse	(	Client	client,
		Path	db_path )

Wipes previous data and loads data from a Parquet file into ClickHouse.

Casts data using the shared schema and inserts it into the benchmark.performance table. Existing table data will be truncated.

Parameters

client	The connected ClickHouse client.
db_path	Path to the Parquet file to load.

Exceptions

FileNotFoundError	if the file does not exist.
Exception	if insertion fails.

Definition at line 102 of file setup.py.

def load_db_to_clickhouse(client: Client, db_path: Path):
    """!Wipes previous data and loads data from a Parquet file into ClickHouse.
 
    Casts data using the shared schema and inserts it into the benchmark.performance table.
    Existing table data will be truncated.
 
    @param client The connected ClickHouse client.
    @param db_path Path to the Parquet file to load.
    @throws FileNotFoundError if the file does not exist.
    @throws Exception if insertion fails.
    """
    if not db_path.exists():
        raise FileNotFoundError(f"{db_path} not found")
 
    log(f"Loading data from: {db_path}")
    df = pl.read_parquet(db_path)
    df = safe_vector_cast(df, SCHEMA)
 
    records = df.to_dicts()
    
    if not records:
        log("No records to insert.")
        return
    
    try:
        client.execute("TRUNCATE TABLE benchmark.performance")
        client.execute("INSERT INTO benchmark.performance VALUES", records)
    except Exception as e:
        err(f"Error inserting records into ClickHouse: {e}")
        raise
    
 

log()

scripts.setup.log

(

str

msg

)

Prints an info message to stdout.

Parameters

msg	The message string.

Returns

None

Definition at line 50 of file setup.py.

def log(msg: str):
    """!Prints an info message to stdout.
 
    \param msg The message string.
    \return None
    """
    print(f"[INFO] {msg}")
 

main()

scripts.setup.main

(

)

CLI entrypoint.

Parses arguments and coordinates Docker, schema setup, and data loading.

Definition at line 134 of file setup.py.

def main():
    """!CLI entrypoint. Parses arguments and coordinates Docker, schema setup, and data loading.
    """
    parser = argparse.ArgumentParser(description="Setup and load benchmark data into ClickHouse")
    parser.add_argument("--load-from-sample", action="store_true", help="Setup and restore from db_sample.parquet")
    parser.add_argument("--load-from-db", action="store_true", help="Setup and use existing db.parquet")
    parser.add_argument("--docker-compose", action="store_true", help="Use Docker Compose to start ClickHouse & Grafana")
    parser.add_argument("--setup-clickhouse", action="store_true", help="Setup ClickHouse database and table")
 
    args = parser.parse_args()
 
    os.makedirs("db", exist_ok=True)
    os.makedirs("samples", exist_ok=True)
    os.makedirs("db/logs", exist_ok=True)
 
    if args.docker_compose:
        log("Setting up docker instance to be ready for ClickHouse database and Grafana.")
        run_command("docker-compose up -d")
        time.sleep(5)  # Wait for ClickHouse to start
 
    client = wait_for_clickhouse()
 
    if args.setup_clickhouse or args.docker_compose:
        setup_clickhouse(client)
        
    if args.load_from_sample:
        log(f"Loading from sample data: {SAMPLE_PATH}")
        shutil.copy(SAMPLE_PATH, DB_PATH)
        load_db_to_clickhouse(client, DB_PATH)
 
    elif args.load_from_db:
        log(f"Loading from existing data: {DB_PATH}")
        load_db_to_clickhouse(client, DB_PATH)
 
    if args.load_from_sample or args.load_from_db:
        run_command("docker restart grafana")
 
 

run_command()

scripts.setup.run_command

(

str

cmd

)

Executes a shell command with logging.

Parameters

cmd	The command string to run.

Exceptions

subprocess.CalledProcessError

if the command fails.

Definition at line 64 of file setup.py.

def run_command(cmd: str):
    """!Executes a shell command with logging.
    @param cmd The command string to run.
    @throws subprocess.CalledProcessError if the command fails.
    """
    log(f"Running command: {cmd}")
    subprocess.run(cmd, shell=True, check=True)
 

setup_clickhouse()

scripts.setup.setup_clickhouse

(

Client

client

)

Creates the ClickHouse database and performance table if they don't exist.

Parameters

client

The connected ClickHouse client.

Definition at line 93 of file setup.py.

def setup_clickhouse(client: Client):
    """!Creates the ClickHouse database and performance table if they don't exist.
    @param client The connected ClickHouse client.
    """
    log("Setting up ClickHouse database and table.")
    client.execute("CREATE DATABASE IF NOT EXISTS benchmark")
    client.execute(generate_clickhouse_table())
    log("Schema loaded into ClickHouse.")
 

wait_for_clickhouse()

Client scripts.setup.wait_for_clickhouse

(

)

Waits for ClickHouse server to become ready, retries for up to 30 attempts.

Returns

A connected ClickHouse Client instance.

Exceptions

RuntimeError

if ClickHouse doesn't respond after all attempts.

Definition at line 72 of file setup.py.

def wait_for_clickhouse() -> Client:
    """!Waits for ClickHouse server to become ready, retries for up to 30 attempts.
    @return A connected ClickHouse Client instance.
    @throws RuntimeError if ClickHouse doesn't respond after all attempts.
    """
    log("Waiting for ClickHouse...")
 
    for attempt in range(30):
        try:
            log(f"Connecting to ClickHouse at {CLICKHOUSE_HOST}:{CLICKHOUSE_TCP_PORT} as user '{CLICKHOUSE_USER}'")
            client = Client(host=CLICKHOUSE_HOST, port=CLICKHOUSE_TCP_PORT, user=CLICKHOUSE_USER, password=CLICKHOUSE_PASSWORD)
            client.execute("SELECT 1")
            log("ClickHouse is ready.")
            return client
        
        except Exception as e:
            log(f"Attempt {attempt+1}/30 failed: {e}")
            time.sleep(2)
 
    raise RuntimeError("ClickHouse did not start after 30 attempts.")