fn_registry/apps/dag_engine/store/store.go

package store

import (
	"database/sql"
	_ "embed"
	"fmt"
	"time"

	_ "github.com/mattn/go-sqlite3"
)

//go:embed migrations/001_init.sql
var migrationSQL string

// DB wraps a SQLite connection for DAG run persistence.
type DB struct {
	conn *sql.DB
	path string
}

// Open opens or creates a DAG engine database at the given path.
func Open(path string) (*DB, error) {
	conn, err := sql.Open("sqlite3", path+"?_journal_mode=WAL&_busy_timeout=5000&_foreign_keys=on")
	if err != nil {
		return nil, fmt.Errorf("store: open %s: %w", path, err)
	}
	if _, err := conn.Exec(migrationSQL); err != nil {
		conn.Close()
		return nil, fmt.Errorf("store: migrate: %w", err)
	}
	return &DB{conn: conn, path: path}, nil
}

// Close closes the database connection.
func (db *DB) Close() error {
	return db.conn.Close()
}

// --- DagRun CRUD ---

// DagRun mirrors infra.DagRun for the store layer.
type DagRun struct {
	ID         string
	DagName    string
	DagPath    string
	Status     string
	Trigger    string
	StartedAt  time.Time
	FinishedAt *time.Time
	Error      string
}

// CreateRun inserts a new run record.
func (db *DB) CreateRun(run *DagRun) error {
	_, err := db.conn.Exec(
		`INSERT INTO dag_runs (id, dag_name, dag_path, status, trigger, started_at, error)
		 VALUES (?, ?, ?, ?, ?, ?, ?)`,
		run.ID, run.DagName, run.DagPath, run.Status, run.Trigger,
		run.StartedAt.Format(time.RFC3339), run.Error,
	)
	return err
}

// UpdateRunStatus updates a run's status and optionally its finished_at and error.
func (db *DB) UpdateRunStatus(id, status string, finishedAt *time.Time, errMsg string) error {
	var fin *string
	if finishedAt != nil {
		s := finishedAt.Format(time.RFC3339)
		fin = &s
	}
	_, err := db.conn.Exec(
		`UPDATE dag_runs SET status=?, finished_at=?, error=? WHERE id=?`,
		status, fin, errMsg, id,
	)
	return err
}

// GetRun retrieves a single run by ID.
func (db *DB) GetRun(id string) (*DagRun, error) {
	row := db.conn.QueryRow(
		`SELECT id, dag_name, dag_path, status, trigger, started_at, finished_at, error
		 FROM dag_runs WHERE id=?`, id,
	)
	return scanRun(row)
}

// ListRuns returns runs, newest first, with optional dag name filter.
func (db *DB) ListRuns(dagName string, limit, offset int) ([]DagRun, int, error) {
	var total int
	var args []interface{}
	where := ""
	if dagName != "" {
		where = " WHERE dag_name=?"
		args = append(args, dagName)
	}
	err := db.conn.QueryRow("SELECT COUNT(*) FROM dag_runs"+where, args...).Scan(&total)
	if err != nil {
		return nil, 0, err
	}

	query := "SELECT id, dag_name, dag_path, status, trigger, started_at, finished_at, error FROM dag_runs" +
		where + " ORDER BY started_at DESC LIMIT ? OFFSET ?"
	args = append(args, limit, offset)

	rows, err := db.conn.Query(query, args...)
	if err != nil {
		return nil, 0, err
	}
	defer rows.Close()

	var runs []DagRun
	for rows.Next() {
		r, err := scanRunRows(rows)
		if err != nil {
			return nil, 0, err
		}
		runs = append(runs, *r)
	}
	return runs, total, rows.Err()
}

// --- DagStepResult CRUD ---

// DagStepResult mirrors infra.DagStepResult for the store layer.
type DagStepResult struct {
	ID         string
	RunID      string
	StepName   string
	Status     string
	ExitCode   int
	Stdout     string
	Stderr     string
	StartedAt  *time.Time
	FinishedAt *time.Time
	DurationMs int64
	Error      string
}

// InsertStepResult inserts a new step result.
func (db *DB) InsertStepResult(r *DagStepResult) error {
	var startedAt, finishedAt *string
	if r.StartedAt != nil {
		s := r.StartedAt.Format(time.RFC3339)
		startedAt = &s
	}
	if r.FinishedAt != nil {
		s := r.FinishedAt.Format(time.RFC3339)
		finishedAt = &s
	}
	_, err := db.conn.Exec(
		`INSERT INTO dag_step_results (id, run_id, step_name, status, exit_code, stdout, stderr, started_at, finished_at, duration_ms, error)
		 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
		r.ID, r.RunID, r.StepName, r.Status, r.ExitCode, r.Stdout, r.Stderr,
		startedAt, finishedAt, r.DurationMs, r.Error,
	)
	return err
}

// UpdateStepResult updates a step result by ID.
func (db *DB) UpdateStepResult(id, status string, exitCode int, stdout, stderr string, finishedAt *time.Time, durationMs int64, errMsg string) error {
	var fin *string
	if finishedAt != nil {
		s := finishedAt.Format(time.RFC3339)
		fin = &s
	}
	_, err := db.conn.Exec(
		`UPDATE dag_step_results SET status=?, exit_code=?, stdout=?, stderr=?, finished_at=?, duration_ms=?, error=? WHERE id=?`,
		status, exitCode, stdout, stderr, fin, durationMs, errMsg, id,
	)
	return err
}

// ListStepResults returns all step results for a given run.
func (db *DB) ListStepResults(runID string) ([]DagStepResult, error) {
	rows, err := db.conn.Query(
		`SELECT id, run_id, step_name, status, exit_code, stdout, stderr, started_at, finished_at, duration_ms, error
		 FROM dag_step_results WHERE run_id=? ORDER BY started_at ASC`, runID,
	)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var results []DagStepResult
	for rows.Next() {
		var r DagStepResult
		var startedAt, finishedAt sql.NullString
		if err := rows.Scan(&r.ID, &r.RunID, &r.StepName, &r.Status, &r.ExitCode,
			&r.Stdout, &r.Stderr, &startedAt, &finishedAt, &r.DurationMs, &r.Error); err != nil {
			return nil, err
		}
		if startedAt.Valid {
			t, _ := time.Parse(time.RFC3339, startedAt.String)
			r.StartedAt = &t
		}
		if finishedAt.Valid {
			t, _ := time.Parse(time.RFC3339, finishedAt.String)
			r.FinishedAt = &t
		}
		results = append(results, r)
	}
	return results, rows.Err()
}

// --- scan helpers ---

type scanner interface {
	Scan(dest ...interface{}) error
}

func scanRun(s scanner) (*DagRun, error) {
	var r DagRun
	var startedAt string
	var finishedAt sql.NullString
	if err := s.Scan(&r.ID, &r.DagName, &r.DagPath, &r.Status, &r.Trigger, &startedAt, &finishedAt, &r.Error); err != nil {
		if err == sql.ErrNoRows {
			return nil, nil
		}
		return nil, err
	}
	r.StartedAt, _ = time.Parse(time.RFC3339, startedAt)
	if finishedAt.Valid {
		t, _ := time.Parse(time.RFC3339, finishedAt.String)
		r.FinishedAt = &t
	}
	return &r, nil
}

func scanRunRows(rows *sql.Rows) (*DagRun, error) {
	return scanRun(rows)
}