Skip to main content

entrenar/storage/sqlite/backend/
schema.rs

1//! SQLite schema definition and migration.
2//!
3//! Defines all tables for the experiment store and handles schema initialization.
4
5use rusqlite::Connection;
6
7/// Current schema version
8pub const CURRENT_VERSION: &str = "1.0.0";
9
10/// Initialize the database schema, creating tables if they don't exist.
11///
12/// Also configures WAL mode and performance pragmas matching the
13/// paiml-mcp-agent-toolkit `.pmat/context.db` pattern.
14pub fn init_schema(conn: &Connection) -> Result<(), rusqlite::Error> {
15    // Performance pragmas (matching paiml-mcp-agent-toolkit)
16    conn.execute_batch(
17        "PRAGMA journal_mode = WAL;
18         PRAGMA synchronous = NORMAL;
19         PRAGMA busy_timeout = 5000;
20         PRAGMA cache_size = -64000;
21         PRAGMA temp_store = MEMORY;",
22    )?;
23
24    conn.execute_batch(SCHEMA_SQL)?;
25
26    // Insert schema version if not present
27    let count: i64 = conn.query_row("SELECT COUNT(*) FROM schema_version", [], |row| row.get(0))?;
28    if count == 0 {
29        conn.execute("INSERT INTO schema_version (version) VALUES (?1)", [CURRENT_VERSION])?;
30    }
31
32    Ok(())
33}
34
35const SCHEMA_SQL: &str = "
36CREATE TABLE IF NOT EXISTS schema_version (
37    version TEXT NOT NULL,
38    applied_at TEXT NOT NULL DEFAULT (datetime('now'))
39);
40
41CREATE TABLE IF NOT EXISTS experiments (
42    id TEXT PRIMARY KEY,
43    name TEXT NOT NULL,
44    description TEXT,
45    config TEXT,
46    tags TEXT,
47    created_at TEXT NOT NULL DEFAULT (datetime('now')),
48    updated_at TEXT NOT NULL DEFAULT (datetime('now'))
49);
50
51CREATE TABLE IF NOT EXISTS runs (
52    id TEXT PRIMARY KEY,
53    experiment_id TEXT NOT NULL,
54    status TEXT NOT NULL DEFAULT 'pending',
55    start_time TEXT,
56    end_time TEXT,
57    tags TEXT,
58    FOREIGN KEY (experiment_id) REFERENCES experiments(id)
59);
60CREATE INDEX IF NOT EXISTS idx_runs_experiment ON runs(experiment_id);
61CREATE INDEX IF NOT EXISTS idx_runs_status ON runs(status);
62
63CREATE TABLE IF NOT EXISTS params (
64    run_id TEXT NOT NULL,
65    key TEXT NOT NULL,
66    value TEXT NOT NULL,
67    type TEXT NOT NULL,
68    PRIMARY KEY (run_id, key),
69    FOREIGN KEY (run_id) REFERENCES runs(id)
70);
71
72CREATE TABLE IF NOT EXISTS metrics (
73    run_id TEXT NOT NULL,
74    key TEXT NOT NULL,
75    step INTEGER NOT NULL,
76    value REAL NOT NULL,
77    timestamp TEXT NOT NULL DEFAULT (datetime('now')),
78    FOREIGN KEY (run_id) REFERENCES runs(id)
79);
80CREATE INDEX IF NOT EXISTS idx_metrics_run_key ON metrics(run_id, key);
81
82CREATE TABLE IF NOT EXISTS artifacts (
83    id TEXT PRIMARY KEY,
84    run_id TEXT NOT NULL,
85    path TEXT NOT NULL,
86    size_bytes INTEGER,
87    sha256 TEXT NOT NULL,
88    data BLOB,
89    created_at TEXT NOT NULL DEFAULT (datetime('now')),
90    FOREIGN KEY (run_id) REFERENCES runs(id)
91);
92CREATE INDEX IF NOT EXISTS idx_artifacts_run ON artifacts(run_id);
93CREATE INDEX IF NOT EXISTS idx_artifacts_sha ON artifacts(sha256);
94
95CREATE TABLE IF NOT EXISTS span_ids (
96    run_id TEXT PRIMARY KEY,
97    span_id TEXT NOT NULL,
98    FOREIGN KEY (run_id) REFERENCES runs(id)
99);
100";