feat(benchmark): add SQLite storage backend
Persistent storage for benchmark history with WAL mode for concurrent access.
This commit is contained in:
179
src/veritext/benchmark/storage.py
Normal file
179
src/veritext/benchmark/storage.py
Normal file
@@ -0,0 +1,179 @@
|
||||
"""SQLite storage for benchmark history."""
|
||||
|
||||
import json
|
||||
import sqlite3
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
from veritext.benchmark.models import BenchmarkRun
|
||||
from veritext.core.exceptions import StorageError
|
||||
|
||||
|
||||
class BenchmarkStorage:
|
||||
"""SQLite-backed storage for benchmark runs."""
|
||||
|
||||
def __init__(self, db_path: Path) -> None:
|
||||
"""
|
||||
Initialise storage, creating tables if needed.
|
||||
|
||||
Args:
|
||||
db_path: Path to the SQLite database file.
|
||||
"""
|
||||
self._db_path = db_path
|
||||
self._ensure_parent_exists()
|
||||
self._init_database()
|
||||
|
||||
def _ensure_parent_exists(self) -> None:
|
||||
"""Ensure the parent directory exists."""
|
||||
self._db_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def _get_connection(self) -> sqlite3.Connection:
|
||||
"""Get a database connection with WAL mode enabled."""
|
||||
conn = sqlite3.connect(str(self._db_path), timeout=30.0)
|
||||
conn.execute("PRAGMA journal_mode=WAL")
|
||||
conn.execute("PRAGMA foreign_keys=ON")
|
||||
conn.row_factory = sqlite3.Row
|
||||
return conn
|
||||
|
||||
def _init_database(self) -> None:
|
||||
"""Create tables if they don't exist."""
|
||||
try:
|
||||
with self._get_connection() as conn:
|
||||
conn.executescript("""
|
||||
CREATE TABLE IF NOT EXISTS benchmark_runs (
|
||||
id TEXT PRIMARY KEY,
|
||||
benchmark_name TEXT NOT NULL,
|
||||
timestamp TEXT NOT NULL,
|
||||
veritext_version TEXT NOT NULL,
|
||||
sample_count INTEGER NOT NULL,
|
||||
metadata TEXT
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS benchmark_metrics (
|
||||
run_id TEXT REFERENCES benchmark_runs(id) ON DELETE CASCADE,
|
||||
metric_name TEXT NOT NULL,
|
||||
value REAL NOT NULL,
|
||||
PRIMARY KEY (run_id, metric_name)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_benchmark_name
|
||||
ON benchmark_runs(benchmark_name, timestamp DESC);
|
||||
""")
|
||||
except sqlite3.Error as e:
|
||||
raise StorageError(f"Failed to initialise database: {e}") from e
|
||||
|
||||
def save_run(self, run: BenchmarkRun) -> None:
|
||||
"""
|
||||
Persist a benchmark run.
|
||||
|
||||
Args:
|
||||
run: The benchmark run to save.
|
||||
|
||||
Raises:
|
||||
StorageError: If the save operation fails.
|
||||
"""
|
||||
try:
|
||||
with self._get_connection() as conn:
|
||||
# Insert the run
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO benchmark_runs
|
||||
(id, benchmark_name, timestamp, veritext_version, sample_count, metadata)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
run.id,
|
||||
run.benchmark_name,
|
||||
run.timestamp.isoformat(),
|
||||
run.veritext_version,
|
||||
run.sample_count,
|
||||
json.dumps(run.metadata) if run.metadata else None,
|
||||
),
|
||||
)
|
||||
|
||||
# Insert metrics
|
||||
for metric_name, value in run.metrics.items():
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO benchmark_metrics (run_id, metric_name, value)
|
||||
VALUES (?, ?, ?)
|
||||
""",
|
||||
(run.id, metric_name, value),
|
||||
)
|
||||
except sqlite3.IntegrityError as e:
|
||||
raise StorageError(f"Run with id '{run.id}' already exists") from e
|
||||
except sqlite3.Error as e:
|
||||
raise StorageError(f"Failed to save benchmark run: {e}") from e
|
||||
|
||||
def get_runs(
|
||||
self,
|
||||
benchmark_name: str,
|
||||
limit: int | None = None,
|
||||
) -> list[BenchmarkRun]:
|
||||
"""
|
||||
Retrieve runs for a benchmark, most recent first.
|
||||
|
||||
Args:
|
||||
benchmark_name: Name of the benchmark to retrieve runs for.
|
||||
limit: Maximum number of runs to return.
|
||||
|
||||
Returns:
|
||||
List of BenchmarkRun objects, most recent first.
|
||||
|
||||
Raises:
|
||||
StorageError: If the retrieval fails.
|
||||
"""
|
||||
try:
|
||||
with self._get_connection() as conn:
|
||||
query = """
|
||||
SELECT id, benchmark_name, timestamp, veritext_version,
|
||||
sample_count, metadata
|
||||
FROM benchmark_runs
|
||||
WHERE benchmark_name = ?
|
||||
ORDER BY timestamp DESC
|
||||
"""
|
||||
if limit is not None:
|
||||
query += " LIMIT ?"
|
||||
rows = conn.execute(query, (benchmark_name, limit)).fetchall()
|
||||
else:
|
||||
rows = conn.execute(query, (benchmark_name,)).fetchall()
|
||||
|
||||
runs = []
|
||||
for row in rows:
|
||||
# Get metrics for this run
|
||||
metrics_rows = conn.execute(
|
||||
"SELECT metric_name, value FROM benchmark_metrics WHERE run_id = ?",
|
||||
(row["id"],),
|
||||
).fetchall()
|
||||
metrics = {m["metric_name"]: m["value"] for m in metrics_rows}
|
||||
|
||||
metadata = json.loads(row["metadata"]) if row["metadata"] else {}
|
||||
|
||||
runs.append(
|
||||
BenchmarkRun(
|
||||
id=row["id"],
|
||||
benchmark_name=row["benchmark_name"],
|
||||
timestamp=datetime.fromisoformat(row["timestamp"]),
|
||||
veritext_version=row["veritext_version"],
|
||||
sample_count=row["sample_count"],
|
||||
metrics=metrics,
|
||||
metadata=metadata,
|
||||
)
|
||||
)
|
||||
|
||||
return runs
|
||||
except sqlite3.Error as e:
|
||||
raise StorageError(f"Failed to retrieve benchmark runs: {e}") from e
|
||||
|
||||
def get_latest_run(self, benchmark_name: str) -> BenchmarkRun | None:
|
||||
"""
|
||||
Get the most recent run for a benchmark.
|
||||
|
||||
Args:
|
||||
benchmark_name: Name of the benchmark.
|
||||
|
||||
Returns:
|
||||
The most recent BenchmarkRun, or None if no runs exist.
|
||||
"""
|
||||
runs = self.get_runs(benchmark_name, limit=1)
|
||||
return runs[0] if runs else None
|
||||
Reference in New Issue
Block a user