spark-mcp/internal/storage/submission_repo.go

package storage

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"strings"
	"time"
)

// Submission is one persisted spark_submit invocation.
type Submission struct {
	ID          int64     `json:"id"`
	FileID      string    `json:"file_id"`
	AppID       string    `json:"app_id"`
	TrackingURL string    `json:"tracking_url,omitempty"`
	ClusterID   string    `json:"cluster_id"`
	ExitCode    int       `json:"exit_code"`
	DurationMS  int64     `json:"duration_ms"`
	SubmittedAt time.Time `json:"submitted_at"`
}

// SubmissionWithFile joins a submission with the original upload file name.
// FileName is empty when the referenced upload has been deleted.
type SubmissionWithFile struct {
	Submission
	FileName string `json:"file_name"`
}

// SubmissionFilter controls the rows returned by SubmissionRepo.List.
type SubmissionFilter struct {
	Search    string
	FileID    string
	ClusterID string
	AppID     string
	Limit     int
	Offset    int
}

// SubmissionRepo provides CRUD operations for spark_submissions.
type SubmissionRepo struct {
	db *DB
}

// Submissions returns a repository bound to this DB handle.
func (d *DB) Submissions() *SubmissionRepo {
	return &SubmissionRepo{db: d}
}

// Create inserts a new submission record.
func (r *SubmissionRepo) Create(ctx context.Context, s Submission) error {
	if s.SubmittedAt.IsZero() {
		s.SubmittedAt = time.Now()
	}
	_, err := r.db.sqlDB.ExecContext(ctx, `
		INSERT INTO spark_submissions (file_id, app_id, tracking_url, cluster_id, exit_code, duration_ms, submitted_at)
		VALUES (?, ?, ?, ?, ?, ?, ?)`,
		s.FileID, s.AppID, s.TrackingURL, s.ClusterID, s.ExitCode, s.DurationMS, s.SubmittedAt.UnixNano(),
	)
	if err != nil {
		return fmt.Errorf("storage: create submission: %w", err)
	}
	return nil
}

// Get retrieves a submission by id.
func (r *SubmissionRepo) Get(ctx context.Context, id int64) (Submission, error) {
	row := r.db.sqlDB.QueryRowContext(ctx, `
		SELECT id, file_id, app_id, tracking_url, cluster_id, exit_code, duration_ms, submitted_at
		FROM spark_submissions
		WHERE id = ?`, id)
	s, err := scanSubmission(row.Scan)
	if errors.Is(err, sql.ErrNoRows) {
		return Submission{}, ErrNotFound
	}
	if err != nil {
		return Submission{}, fmt.Errorf("storage: get submission %d: %w", id, err)
	}
	return s, nil
}

// List returns submissions ordered by submitted_at descending.
//
// Search matches app_id or the joined upload file name. Empty filter strings
// mean no restriction. Limit defaults to 100 and is capped at 500.
func (r *SubmissionRepo) List(ctx context.Context, filter SubmissionFilter) ([]SubmissionWithFile, error) {
	if filter.Limit <= 0 {
		filter.Limit = 100
	}
	if filter.Limit > 500 {
		filter.Limit = 500
	}
	if filter.Offset < 0 {
		filter.Offset = 0
	}

	args := []any{}
	var where []string
	if filter.Search != "" {
		like := "%" + filter.Search + "%"
		args = append(args, like, like)
		where = append(where, "(s.app_id LIKE ? OR COALESCE(u.name, '') LIKE ?)")
	}
	if filter.FileID != "" {
		args = append(args, filter.FileID)
		where = append(where, "s.file_id = ?")
	}
	if filter.ClusterID != "" {
		args = append(args, filter.ClusterID)
		where = append(where, "s.cluster_id = ?")
	}
	if filter.AppID != "" {
		args = append(args, filter.AppID)
		where = append(where, "s.app_id = ?")
	}

	// WHERE args are in the order the conditions were appended; append LIMIT/OFFSET last.
	args = append(args, filter.Limit, filter.Offset)

	q := `SELECT s.id, s.file_id, s.app_id, s.tracking_url, s.cluster_id, s.exit_code, s.duration_ms, s.submitted_at, COALESCE(u.name, '') AS file_name
		FROM spark_submissions s
		LEFT JOIN upload_files u ON s.file_id = u.file_id`
	if len(where) > 0 {
		q += "\n\t\tWHERE " + strings.Join(where, " AND ")
	}
	q += "\n\t\tORDER BY s.submitted_at DESC\n\t\tLIMIT ? OFFSET ?"

	rows, err := r.db.sqlDB.QueryContext(ctx, q, args...)
	if err != nil {
		return nil, fmt.Errorf("storage: list submissions: %w", err)
	}
	defer rows.Close()

	var out []SubmissionWithFile
	for rows.Next() {
		swf, err := scanSubmissionWithFile(rows.Scan)
		if err != nil {
			return nil, fmt.Errorf("storage: list submissions: %w", err)
		}
		out = append(out, swf)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("storage: list submissions: %w", err)
	}
	if out == nil {
		out = []SubmissionWithFile{}
	}
	return out, nil
}

// Delete removes a submission by id.
func (r *SubmissionRepo) Delete(ctx context.Context, id int64) error {
	res, err := r.db.sqlDB.ExecContext(ctx, `DELETE FROM spark_submissions WHERE id = ?`, id)
	if err != nil {
		return fmt.Errorf("storage: delete submission %d: %w", id, err)
	}
	n, err := res.RowsAffected()
	if err != nil {
		return fmt.Errorf("storage: delete submission %d: rows affected: %w", id, err)
	}
	if n == 0 {
		return ErrNotFound
	}
	return nil
}

func scanSubmission(scan scanFunc) (Submission, error) {
	var s Submission
	var trackingURL sql.NullString
	var submittedAt int64
	if err := scan(&s.ID, &s.FileID, &s.AppID, &trackingURL, &s.ClusterID, &s.ExitCode, &s.DurationMS, &submittedAt); err != nil {
		return Submission{}, err
	}
	if trackingURL.Valid {
		s.TrackingURL = trackingURL.String
	}
	s.SubmittedAt = time.Unix(0, submittedAt)
	return s, nil
}

func scanSubmissionWithFile(scan scanFunc) (SubmissionWithFile, error) {
	var swf SubmissionWithFile
	var trackingURL sql.NullString
	var submittedAt int64
	if err := scan(&swf.ID, &swf.FileID, &swf.AppID, &trackingURL, &swf.ClusterID, &swf.ExitCode, &swf.DurationMS, &submittedAt, &swf.FileName); err != nil {
		return SubmissionWithFile{}, err
	}
	if trackingURL.Valid {
		swf.TrackingURL = trackingURL.String
	}
	swf.SubmittedAt = time.Unix(0, submittedAt)
	return swf, nil
}