Spaces:
Build error
Build error
File size: 3,121 Bytes
1c75c98 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
# 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()
|