tracklight_server/db/duckdb.py

# tracklight_server/db/duckdb.py

import duckdb
import os

DB_FILE = "tracklight.db"
TABLE_NAME = "metrics"

# Define the database path in the user's home directory
HOME_DIR = os.path.expanduser("~")
TRACKLIGHT_DIR = os.path.join(HOME_DIR, ".tracklight")
DB_PATH = os.path.join(TRACKLIGHT_DIR, DB_FILE)

# Ensure the .tracklight directory exists
os.makedirs(TRACKLIGHT_DIR, exist_ok=True)

def get_connection():
    """Returns a connection to the DuckDB database."""
    return duckdb.connect(DB_PATH)

def create_tables():
    """Creates the metrics and config tables if they don't exist."""
    with get_connection() as con:
        con.execute(f"""
            CREATE TABLE IF NOT EXISTS {TABLE_NAME} (
                run_id VARCHAR,
                project VARCHAR,
                user_name VARCHAR,
                metric_name VARCHAR,
                value DOUBLE,
                timestamp TIMESTAMP
            )
        """)
        con.execute("""
            CREATE TABLE IF NOT EXISTS config (
                run_id VARCHAR,
                config_name VARCHAR,
                value VARCHAR
            )
        """)

def insert_metrics(metrics: list):
    """Inserts a list of metrics into the database."""
    with get_connection() as con:
        con.executemany(f"INSERT INTO {TABLE_NAME} VALUES (?, ?, ?, ?, ?, ?)", 
                        [(m['run_id'], m['project'], m['user'], m['metric_name'], m['value'], m['timestamp']) for m in metrics])

def insert_config(configs: list):
    """Inserts a list of config values into the database."""
    with get_connection() as con:
        con.executemany("INSERT INTO config VALUES (?, ?, ?)", 
                        [(c['run_id'], c['config_name'], c['value']) for c in configs])

def get_projects():
    """Retrieves all projects with their creation timestamp."""
    with get_connection() as con:
        return con.execute(f"""
            SELECT project, MIN(timestamp) as creation_time
            FROM {TABLE_NAME}
            GROUP BY project
            ORDER BY creation_time DESC
        """).fetchall()

def get_runs(project: str, user: str):
    """Retrieves all runs for a given project and user, sorted by the most recent timestamp."""
    with get_connection() as con:
        return con.execute(f"""
            SELECT run_id 
            FROM {TABLE_NAME} 
            WHERE project = ? AND user_name = ?
            GROUP BY run_id
            ORDER BY MAX(timestamp) DESC
        """, [project, user]).fetchall()

def get_metrics_for_run(run_id: str):
    """Retrieves all metrics for a given run, excluding config parameters."""
    with get_connection() as con:
        return con.execute(f"SELECT metric_name, value, timestamp FROM {TABLE_NAME} WHERE run_id = ? AND metric_name NOT LIKE 'config/%' ORDER BY timestamp", [run_id]).fetchall()

def get_config_for_run(run_id: str):
    """Retrieves all configuration parameters for a given run."""
    with get_connection() as con:
        return con.execute(f"SELECT config_name, value FROM config WHERE run_id = ?", [run_id]).fetchall()


# Initialize the database
create_tables()