storage: add spark_submissions table and SubmissionRepo
Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
tao.chen
@@ -47,6 +47,19 @@ CREATE TABLE IF NOT EXISTS upload_files (
|
||||
uploaded_at INTEGER NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_uploads_ts ON upload_files(uploaded_at DESC);
|
||||
CREATE TABLE IF NOT EXISTS spark_submissions (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
file_id TEXT NOT NULL,
|
||||
app_id TEXT NOT NULL,
|
||||
tracking_url TEXT,
|
||||
cluster_id TEXT NOT NULL,
|
||||
exit_code INTEGER NOT NULL,
|
||||
duration_ms INTEGER NOT NULL,
|
||||
submitted_at INTEGER NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_submissions_app_id ON spark_submissions(app_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_submissions_submitted_at ON spark_submissions(submitted_at DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_submissions_file_id ON spark_submissions(file_id);
|
||||
`
|
||||
|
||||
// DB is the storage handle.
|
||||
|
||||
@@ -0,0 +1,194 @@
|
||||
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
|
||||
}
|
||||
@@ -0,0 +1,151 @@
|
||||
package storage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func setupSubmissionTest(t *testing.T) (*SubmissionRepo, *UploadRepo, context.Context) {
|
||||
db, err := Open(":memory:")
|
||||
if err != nil {
|
||||
t.Fatalf("open: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
return db.Submissions(), db.Uploads(), context.Background()
|
||||
}
|
||||
|
||||
func TestSubmissionRepo_RoundTrip(t *testing.T) {
|
||||
repo, uploadRepo, ctx := setupSubmissionTest(t)
|
||||
|
||||
if err := uploadRepo.Create(ctx, "00000000000000000000000000000001", "alpha.py", 10, "deadbeef", time.Now()); err != nil {
|
||||
t.Fatalf("create upload: %v", err)
|
||||
}
|
||||
|
||||
s := Submission{
|
||||
FileID: "00000000000000000000000000000001",
|
||||
AppID: "application_1_0001",
|
||||
TrackingURL: "http://rm.example.com:8088/proxy/application_1_0001/",
|
||||
ClusterID: "cluster-a",
|
||||
ExitCode: 0,
|
||||
DurationMS: 1234,
|
||||
SubmittedAt: time.Unix(0, time.Now().UnixNano()),
|
||||
}
|
||||
if err := repo.Create(ctx, s); err != nil {
|
||||
t.Fatalf("create: %v", err)
|
||||
}
|
||||
|
||||
list, err := repo.List(ctx, SubmissionFilter{Limit: 10})
|
||||
if err != nil {
|
||||
t.Fatalf("list: %v", err)
|
||||
}
|
||||
if len(list) != 1 {
|
||||
t.Fatalf("len(list)=%d, want 1", len(list))
|
||||
}
|
||||
swf := list[0]
|
||||
if swf.FileName != "alpha.py" {
|
||||
t.Errorf("file_name=%q, want alpha.py", swf.FileName)
|
||||
}
|
||||
if swf.AppID != s.AppID {
|
||||
t.Errorf("app_id=%q, want %q", swf.AppID, s.AppID)
|
||||
}
|
||||
|
||||
id := swf.ID
|
||||
got, err := repo.Get(ctx, id)
|
||||
if err != nil {
|
||||
t.Fatalf("get: %v", err)
|
||||
}
|
||||
if got.ID != id || got.AppID != s.AppID {
|
||||
t.Errorf("got=%+v, want id=%d app_id=%s", got, id, s.AppID)
|
||||
}
|
||||
|
||||
if err := repo.Delete(ctx, id); err != nil {
|
||||
t.Fatalf("delete: %v", err)
|
||||
}
|
||||
if _, err := repo.Get(ctx, id); !isErrNotFound(err) {
|
||||
t.Errorf("expected ErrNotFound after delete, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSubmissionRepo_SearchByAppID(t *testing.T) {
|
||||
repo, uploadRepo, ctx := setupSubmissionTest(t)
|
||||
|
||||
if err := uploadRepo.Create(ctx, "00000000000000000000000000000001", "alpha.py", 10, "deadbeef", time.Now()); err != nil {
|
||||
t.Fatalf("create upload alpha: %v", err)
|
||||
}
|
||||
if err := uploadRepo.Create(ctx, "00000000000000000000000000000002", "beta.py", 10, "cafebabe", time.Now()); err != nil {
|
||||
t.Fatalf("create upload beta: %v", err)
|
||||
}
|
||||
if err := repo.Create(ctx, Submission{FileID: "00000000000000000000000000000001", AppID: "application_alpha_1", ClusterID: "c", ExitCode: 0, DurationMS: 1, SubmittedAt: time.Now()}); err != nil {
|
||||
t.Fatalf("create alpha: %v", err)
|
||||
}
|
||||
if err := repo.Create(ctx, Submission{FileID: "00000000000000000000000000000002", AppID: "application_beta_1", ClusterID: "c", ExitCode: 0, DurationMS: 1, SubmittedAt: time.Now()}); err != nil {
|
||||
t.Fatalf("create beta: %v", err)
|
||||
}
|
||||
|
||||
list, err := repo.List(ctx, SubmissionFilter{Search: "alpha", Limit: 10})
|
||||
if err != nil {
|
||||
t.Fatalf("search: %v", err)
|
||||
}
|
||||
if len(list) != 1 || list[0].AppID != "application_alpha_1" || list[0].FileName != "alpha.py" {
|
||||
t.Errorf("got %+v, want 1 alpha submission", list)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSubmissionRepo_SearchByFileName(t *testing.T) {
|
||||
db, err := Open(":memory:")
|
||||
if err != nil {
|
||||
t.Fatalf("open: %v", err)
|
||||
}
|
||||
defer db.Close()
|
||||
repo := db.Submissions()
|
||||
uploadRepo := db.Uploads()
|
||||
ctx := context.Background()
|
||||
|
||||
if err := uploadRepo.Create(ctx, "00000000000000000000000000000001", "report.py", 10, "deadbeef", time.Now()); err != nil {
|
||||
t.Fatalf("create upload: %v", err)
|
||||
}
|
||||
if err := repo.Create(ctx, Submission{FileID: "00000000000000000000000000000001", AppID: "application_x_1", ClusterID: "c", ExitCode: 0, DurationMS: 1, SubmittedAt: time.Now()}); err != nil {
|
||||
t.Fatalf("create: %v", err)
|
||||
}
|
||||
|
||||
list, err := repo.List(ctx, SubmissionFilter{Search: "report", Limit: 10})
|
||||
if err != nil {
|
||||
t.Fatalf("search: %v", err)
|
||||
}
|
||||
if len(list) != 1 || list[0].FileName != "report.py" {
|
||||
t.Errorf("got %+v, want 1 report submission", list)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSubmissionRepo_FilterByFileID(t *testing.T) {
|
||||
db, err := Open(":memory:")
|
||||
if err != nil {
|
||||
t.Fatalf("open: %v", err)
|
||||
}
|
||||
defer db.Close()
|
||||
repo := db.Submissions()
|
||||
ctx := context.Background()
|
||||
|
||||
if err := repo.Create(ctx, Submission{FileID: "f1", AppID: "a1", ClusterID: "c", ExitCode: 0, DurationMS: 1, SubmittedAt: time.Now()}); err != nil {
|
||||
t.Fatalf("create f1: %v", err)
|
||||
}
|
||||
if err := repo.Create(ctx, Submission{FileID: "f2", AppID: "a2", ClusterID: "c", ExitCode: 0, DurationMS: 1, SubmittedAt: time.Now()}); err != nil {
|
||||
t.Fatalf("create f2: %v", err)
|
||||
}
|
||||
|
||||
list, err := repo.List(ctx, SubmissionFilter{FileID: "f1", Limit: 10})
|
||||
if err != nil {
|
||||
t.Fatalf("list: %v", err)
|
||||
}
|
||||
if len(list) != 1 || list[0].FileID != "f1" {
|
||||
t.Errorf("got %+v, want f1 submission", list)
|
||||
}
|
||||
}
|
||||
|
||||
func isErrNotFound(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
return err.Error() == ErrNotFound.Error()
|
||||
}
|
||||
Reference in New Issue
Block a user