fetch_ml/internal/store/store.go

// Package store provides SQLite storage for TUI local mode
package store

import (
	"database/sql"
	"fmt"
	"os"
	"path/filepath"

	_ "modernc.org/sqlite" // SQLite driver
)

// Store handles SQLite operations for the TUI
type Store struct {
	db     *sql.DB
	dbPath string
}

// RunInfo represents a local run from SQLite
type RunInfo struct {
	EndTime      *string
	PID          *int64
	RunID        string
	ExperimentID string
	Name         string
	Status       string
	StartTime    string
	Synced       bool
}

// Metric represents a logged metric
type Metric struct {
	Key   string
	Value float64
	Step  int64
}

// Param represents a logged parameter
type Param struct {
	Key   string
	Value string
}

// Open opens the SQLite database at the given path
func Open(dbPath string) (*Store, error) {
	dir := filepath.Dir(dbPath)
	if err := os.MkdirAll(dir, 0750); err != nil {
		return nil, fmt.Errorf("failed to create directory: %w", err)
	}

	db, err := sql.Open("sqlite", dbPath)
	if err != nil {
		return nil, fmt.Errorf("failed to open database: %w", err)
	}

	if _, err := db.Exec("PRAGMA journal_mode=WAL;"); err != nil {
		_ = db.Close()
		return nil, fmt.Errorf("failed to enable WAL: %w", err)
	}

	if _, err := db.Exec("PRAGMA synchronous=NORMAL;"); err != nil {
		_ = db.Close()
		return nil, fmt.Errorf("failed to set synchronous: %w", err)
	}

	store := &Store{
		db:     db,
		dbPath: dbPath,
	}

	if err := store.initSchema(); err != nil {
		_ = db.Close()
		return nil, fmt.Errorf("failed to init schema: %w", err)
	}

	return store, nil
}

// Close closes the database connection
func (s *Store) Close() error {
	if s.db != nil {
		_, _ = s.db.Exec("PRAGMA wal_checkpoint(TRUNCATE);")
		return s.db.Close()
	}
	return nil
}

// initSchema creates the required tables if they don't exist
func (s *Store) initSchema() error {
	schema := `
CREATE TABLE IF NOT EXISTS ml_experiments (
    experiment_id TEXT PRIMARY KEY,
    name TEXT NOT NULL,
    description TEXT,
    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
    updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
    status TEXT DEFAULT 'active'
);

CREATE TABLE IF NOT EXISTS ml_runs (
    run_id TEXT PRIMARY KEY,
    experiment_id TEXT NOT NULL,
    name TEXT,
    status TEXT NOT NULL,
    start_time DATETIME DEFAULT CURRENT_TIMESTAMP,
    end_time DATETIME,
    pid INTEGER,
    synced INTEGER DEFAULT 0,
    FOREIGN KEY (experiment_id) REFERENCES ml_experiments(experiment_id)
);

CREATE TABLE IF NOT EXISTS ml_metrics (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    run_id TEXT NOT NULL,
    key TEXT NOT NULL,
    value REAL NOT NULL,
    step INTEGER DEFAULT 0,
    timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
    FOREIGN KEY (run_id) REFERENCES ml_runs(run_id)
);

CREATE TABLE IF NOT EXISTS ml_params (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    run_id TEXT NOT NULL,
    key TEXT NOT NULL,
    value TEXT NOT NULL,
    FOREIGN KEY (run_id) REFERENCES ml_runs(run_id)
);

CREATE TABLE IF NOT EXISTS ml_tags (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    run_id TEXT NOT NULL,
    key TEXT NOT NULL,
    value TEXT NOT NULL,
    FOREIGN KEY (run_id) REFERENCES ml_runs(run_id)
);

CREATE INDEX IF NOT EXISTS idx_runs_experiment ON ml_runs(experiment_id);
CREATE INDEX IF NOT EXISTS idx_runs_status ON ml_runs(status);
CREATE INDEX IF NOT EXISTS idx_metrics_run ON ml_metrics(run_id);
CREATE INDEX IF NOT EXISTS idx_params_run ON ml_params(run_id);
CREATE INDEX IF NOT EXISTS idx_tags_run ON ml_tags(run_id);
`
	_, err := s.db.Exec(schema)
	return err
}

// GetUnsyncedRuns returns all runs that haven't been synced to the server
func (s *Store) GetUnsyncedRuns() ([]RunInfo, error) {
	rows, err := s.db.Query(`
		SELECT run_id, experiment_id, name, status, start_time, end_time, pid, synced
		FROM ml_runs
		WHERE synced = 0
		ORDER BY start_time DESC
	`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var runs []RunInfo
	for rows.Next() {
		var r RunInfo
		var endTime sql.NullString
		var pid sql.NullInt64
		err := rows.Scan(&r.RunID, &r.ExperimentID, &r.Name, &r.Status, &r.StartTime, &endTime, &pid, &r.Synced)
		if err != nil {
			return nil, err
		}
		if endTime.Valid {
			r.EndTime = &endTime.String
		}
		if pid.Valid {
			r.PID = &pid.Int64
		}
		runs = append(runs, r)
	}
	return runs, rows.Err()
}

// GetRunsByExperiment returns all runs for a given experiment
func (s *Store) GetRunsByExperiment(expID string) ([]RunInfo, error) {
	rows, err := s.db.Query(`
		SELECT run_id, experiment_id, name, status, start_time, end_time, pid, synced
		FROM ml_runs
		WHERE experiment_id = ?
		ORDER BY start_time DESC
	`, expID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var runs []RunInfo
	for rows.Next() {
		var r RunInfo
		var endTime sql.NullString
		var pid sql.NullInt64
		err := rows.Scan(&r.RunID, &r.ExperimentID, &r.Name, &r.Status, &r.StartTime, &endTime, &pid, &r.Synced)
		if err != nil {
			return nil, err
		}
		if endTime.Valid {
			r.EndTime = &endTime.String
		}
		if pid.Valid {
			r.PID = &pid.Int64
		}
		runs = append(runs, r)
	}
	return runs, rows.Err()
}

// MarkRunSynced marks a run as synced to the server
func (s *Store) MarkRunSynced(runID string) error {
	_, err := s.db.Exec("UPDATE ml_runs SET synced = 1 WHERE run_id = ?", runID)
	return err
}

// GetRunMetrics returns all metrics for a run
func (s *Store) GetRunMetrics(runID string) ([]Metric, error) {
	rows, err := s.db.Query(`
		SELECT key, value, step
		FROM ml_metrics
		WHERE run_id = ?
		ORDER BY step ASC, key ASC
	`, runID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var metrics []Metric
	for rows.Next() {
		var m Metric
		if err := rows.Scan(&m.Key, &m.Value, &m.Step); err != nil {
			return nil, err
		}
		metrics = append(metrics, m)
	}
	return metrics, rows.Err()
}

// GetRunParams returns all params for a run
func (s *Store) GetRunParams(runID string) ([]Param, error) {
	rows, err := s.db.Query(`
		SELECT key, value
		FROM ml_params
		WHERE run_id = ?
		ORDER BY key ASC
	`, runID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var params []Param
	for rows.Next() {
		var p Param
		if err := rows.Scan(&p.Key, &p.Value); err != nil {
			return nil, err
		}
		params = append(params, p)
	}
	return params, rows.Err()
}

// DB returns the underlying database connection
func (s *Store) DB() *sql.DB {
	return s.db
}