From a2232cecc407afd8f534531cad51c3c2bacde841 Mon Sep 17 00:00:00 2001 From: "tao.chen" <93983997+taochen-ct@users.noreply.github.com> Date: Tue, 14 Jul 2026 13:45:03 +0800 Subject: [PATCH] storage: add spark_submissions table and SubmissionRepo Co-Authored-By: Claude --- internal/storage/sqlite.go | 13 ++ internal/storage/submission_repo.go | 194 +++++++++++++++++++++++ internal/storage/submission_repo_test.go | 151 ++++++++++++++++++ 3 files changed, 358 insertions(+) create mode 100644 internal/storage/submission_repo.go create mode 100644 internal/storage/submission_repo_test.go diff --git a/internal/storage/sqlite.go b/internal/storage/sqlite.go index 1808ac4..8725b66 100644 --- a/internal/storage/sqlite.go +++ b/internal/storage/sqlite.go @@ -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. diff --git a/internal/storage/submission_repo.go b/internal/storage/submission_repo.go new file mode 100644 index 0000000..7ed94bd --- /dev/null +++ b/internal/storage/submission_repo.go @@ -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 +} diff --git a/internal/storage/submission_repo_test.go b/internal/storage/submission_repo_test.go new file mode 100644 index 0000000..a19af41 --- /dev/null +++ b/internal/storage/submission_repo_test.go @@ -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() +}