uploads: dual-write uploads to DB, startup backfill, and audit log

Store now indexes every Save into the upload_files table via SetRepo.

DB failures are logged as warnings and do not fail the upload because

.meta.json remains the source of truth and startup backfill recovers.

Add Store.Backfill to walk Root at startup and insert index rows for

any pre-existing .meta.json sidecars, swallowing duplicate-key races.

The upload_file MCP Tool now writes an audit_log entry on success.

Tests cover dual-write args, repo-error non-failure, and backfill

skipping existing rows.

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
tao.chen
2026-07-14 11:30:39 +08:00
co-authored by Claude
parent 1796b5b872
commit 876f1c9edb
7 changed files with 316 additions and 17 deletions
+2
View File
@@ -18,6 +18,8 @@ const (
ActionClusterCreate Action = "cluster.create" ActionClusterCreate Action = "cluster.create"
ActionClusterUpdate Action = "cluster.update" ActionClusterUpdate Action = "cluster.update"
ActionClusterDelete Action = "cluster.delete" ActionClusterDelete Action = "cluster.delete"
ActionUploadCreate Action = "upload.create"
ActionUploadDelete Action = "upload.delete"
) )
// Entry is one admin write operation recorded for accountability. // Entry is one admin write operation recorded for accountability.
+2
View File
@@ -5,6 +5,7 @@ import (
"time" "time"
"spark-mcp-go/internal/analyzer" "spark-mcp-go/internal/analyzer"
"spark-mcp-go/internal/audit"
"spark-mcp-go/internal/httpclient" "spark-mcp-go/internal/httpclient"
"spark-mcp-go/internal/storage" "spark-mcp-go/internal/storage"
"spark-mcp-go/internal/uploads" "spark-mcp-go/internal/uploads"
@@ -13,6 +14,7 @@ import (
// Deps bundles the dependencies shared by all MCP Tool handlers. // Deps bundles the dependencies shared by all MCP Tool handlers.
type Deps struct { type Deps struct {
Logger *slog.Logger Logger *slog.Logger
AuditRepo *audit.Repo
ClusterRepo *storage.ClusterRepo ClusterRepo *storage.ClusterRepo
SparkSubmitTimeout time.Duration SparkSubmitTimeout time.Duration
HTTPClient *httpclient.Client HTTPClient *httpclient.Client
+24 -7
View File
@@ -16,6 +16,7 @@ import (
"github.com/mark3labs/mcp-go/mcp" "github.com/mark3labs/mcp-go/mcp"
"spark-mcp-go/internal/audit"
"spark-mcp-go/internal/cluster" "spark-mcp-go/internal/cluster"
"spark-mcp-go/internal/httpclient" "spark-mcp-go/internal/httpclient"
"spark-mcp-go/internal/storage" "spark-mcp-go/internal/storage"
@@ -24,7 +25,7 @@ import (
// testDepsWithDataDir returns dependencies backed by an in-memory DB and a // testDepsWithDataDir returns dependencies backed by an in-memory DB and a
// temporary data directory. // temporary data directory.
func testDepsWithDataDir(t *testing.T) (*Deps, *storage.ClusterRepo) { func testDepsWithDataDir(t *testing.T) (*Deps, *storage.ClusterRepo, *audit.Repo) {
t.Helper() t.Helper()
db, err := storage.Open(":memory:") db, err := storage.Open(":memory:")
if err != nil { if err != nil {
@@ -43,11 +44,12 @@ func testDepsWithDataDir(t *testing.T) (*Deps, *storage.ClusterRepo) {
Timeout: 5 * time.Second, Timeout: 5 * time.Second,
MaxResponseBytes: 1 << 20, MaxResponseBytes: 1 << 20,
}), }),
AuditRepo: audit.NewRepo(db),
ClusterRepo: db.Clusters(), ClusterRepo: db.Clusters(),
MaxResponseBytes: 1 << 20, MaxResponseBytes: 1 << 20,
DataDir: dataDir, DataDir: dataDir,
UploadStore: uploadStore, UploadStore: uploadStore,
}, db.Clusters() }, db.Clusters(), audit.NewRepo(db)
} }
// createCluster creates a cluster in the repository with the given fields. // createCluster creates a cluster in the repository with the given fields.
@@ -98,7 +100,7 @@ func TestFetchURL(t *testing.T) {
})) }))
defer srv.Close() defer srv.Close()
deps, repo := testDepsWithDataDir(t) deps, repo, _ := testDepsWithDataDir(t)
createCluster(t, repo, &cluster.Cluster{ createCluster(t, repo, &cluster.Cluster{
ID: "cluster-a", ID: "cluster-a",
Name: "Cluster A", Name: "Cluster A",
@@ -292,7 +294,7 @@ func TestFetchURL_RedirectPreservesAuth(t *testing.T) {
})) }))
defer server1.Close() defer server1.Close()
deps, repo := testDepsWithDataDir(t) deps, repo, _ := testDepsWithDataDir(t)
createCluster(t, repo, &cluster.Cluster{ createCluster(t, repo, &cluster.Cluster{
ID: "cluster-redirect", ID: "cluster-redirect",
Name: "Redirect", Name: "Redirect",
@@ -321,7 +323,7 @@ func TestFetchURL_RedirectPreservesAuth(t *testing.T) {
} }
func TestUploadFile(t *testing.T) { func TestUploadFile(t *testing.T) {
deps, _ := testDepsWithDataDir(t) deps, _, auditRepo := testDepsWithDataDir(t)
req := newToolRequest(UploadFileName, map[string]any{ req := newToolRequest(UploadFileName, map[string]any{
"filename": "hello.txt", "filename": "hello.txt",
@@ -367,10 +369,25 @@ func TestUploadFile(t *testing.T) {
if info.Mode().Perm() != 0o640 { if info.Mode().Perm() != 0o640 {
t.Errorf("mode=%o, want %o", info.Mode().Perm(), 0o640) t.Errorf("mode=%o, want %o", info.Mode().Perm(), 0o640)
} }
entries, err := auditRepo.List(context.Background(), 100)
if err != nil {
t.Fatalf("list audit: %v", err)
}
found := false
for _, e := range entries {
if e.Action == audit.ActionUploadCreate && e.Actor == "tool:upload_file" && e.ClusterID == fileID {
found = true
break
}
}
if !found {
t.Errorf("missing upload.create audit entry for file_id %s", fileID)
}
} }
func TestUploadFile_PathTraversal(t *testing.T) { func TestUploadFile_PathTraversal(t *testing.T) {
deps, _ := testDepsWithDataDir(t) deps, _, _ := testDepsWithDataDir(t)
cases := []string{ cases := []string{
"../../../etc/passwd", "../../../etc/passwd",
@@ -401,7 +418,7 @@ func TestUploadFile_PathTraversal(t *testing.T) {
} }
func TestUploadFile_Base64(t *testing.T) { func TestUploadFile_Base64(t *testing.T) {
deps, _ := testDepsWithDataDir(t) deps, _, _ := testDepsWithDataDir(t)
req := newToolRequest(UploadFileName, map[string]any{ req := newToolRequest(UploadFileName, map[string]any{
"filename": "hello.bin", "filename": "hello.bin",
+5 -5
View File
@@ -14,7 +14,7 @@ import (
) )
func TestSparkSubmit_StructuredCommand(t *testing.T) { func TestSparkSubmit_StructuredCommand(t *testing.T) {
deps, repo := testDepsWithDataDir(t) deps, repo, _ := testDepsWithDataDir(t)
deps.SparkSubmitTimeout = 5 * time.Second deps.SparkSubmitTimeout = 5 * time.Second
store := deps.UploadStore store := deps.UploadStore
@@ -98,7 +98,7 @@ func TestSparkSubmit_StructuredCommand(t *testing.T) {
} }
func TestSparkSubmit_MissingRequiredField(t *testing.T) { func TestSparkSubmit_MissingRequiredField(t *testing.T) {
deps, _ := testDepsWithDataDir(t) deps, _, _ := testDepsWithDataDir(t)
store := deps.UploadStore store := deps.UploadStore
_, _, _, _, mintedPath, err := store.Save([]byte("# dummy\n"), "dummy.py") _, _, _, _, mintedPath, err := store.Save([]byte("# dummy\n"), "dummy.py")
if err != nil { if err != nil {
@@ -131,7 +131,7 @@ func TestSparkSubmit_MissingRequiredField(t *testing.T) {
} }
func TestSparkSubmit_BadSparkConfValue(t *testing.T) { func TestSparkSubmit_BadSparkConfValue(t *testing.T) {
deps, _ := testDepsWithDataDir(t) deps, _, _ := testDepsWithDataDir(t)
store := deps.UploadStore store := deps.UploadStore
_, _, _, _, mintedPath, err := store.Save([]byte("# dummy\n"), "dummy.py") _, _, _, _, mintedPath, err := store.Save([]byte("# dummy\n"), "dummy.py")
if err != nil { if err != nil {
@@ -166,7 +166,7 @@ func TestSparkSubmit_BadSparkConfValue(t *testing.T) {
} }
func TestSparkSubmit_RejectsNonMintedPath(t *testing.T) { func TestSparkSubmit_RejectsNonMintedPath(t *testing.T) {
deps, _ := testDepsWithDataDir(t) deps, _, _ := testDepsWithDataDir(t)
unminted := filepath.Join(t.TempDir(), "unminted.py") unminted := filepath.Join(t.TempDir(), "unminted.py")
if err := os.WriteFile(unminted, []byte("print('not from upload_file')\n"), 0o644); err != nil { if err := os.WriteFile(unminted, []byte("print('not from upload_file')\n"), 0o644); err != nil {
@@ -200,7 +200,7 @@ func TestSparkSubmit_RejectsNonMintedPath(t *testing.T) {
} }
func TestSparkSubmit_EmptyMaster(t *testing.T) { func TestSparkSubmit_EmptyMaster(t *testing.T) {
deps, _ := testDepsWithDataDir(t) deps, _, _ := testDepsWithDataDir(t)
store := deps.UploadStore store := deps.UploadStore
_, _, _, _, mintedPath, err := store.Save([]byte("# dummy\n"), "dummy.py") _, _, _, _, mintedPath, err := store.Save([]byte("# dummy\n"), "dummy.py")
if err != nil { if err != nil {
+17
View File
@@ -8,6 +8,8 @@ import (
"regexp" "regexp"
"github.com/mark3labs/mcp-go/mcp" "github.com/mark3labs/mcp-go/mcp"
"spark-mcp-go/internal/audit"
) )
const UploadFileName = "upload_file" const UploadFileName = "upload_file"
@@ -76,6 +78,21 @@ func (d *Deps) UploadFileHandler(ctx context.Context, req mcp.CallToolRequest) (
return errResult("upload_file: save upload: " + err.Error()), nil return errResult("upload_file: save upload: " + err.Error()), nil
} }
if d.AuditRepo != nil {
details, _ := audit.MarshalDetails(map[string]any{
"file_id": fileID,
"name": filename,
"size": size,
"sha256": sha256Hex,
})
_ = d.AuditRepo.Insert(ctx, &audit.Entry{
Actor: "tool:upload_file",
Action: audit.ActionUploadCreate,
ClusterID: fileID,
Details: details,
})
}
result := map[string]any{ result := map[string]any{
"file_id": fileID, "file_id": fileID,
"path": absPath, "path": absPath,
+86 -5
View File
@@ -9,12 +9,16 @@ import (
"crypto/sha256" "crypto/sha256"
"encoding/hex" "encoding/hex"
"encoding/json" "encoding/json"
"errors"
"fmt" "fmt"
"log/slog"
"os" "os"
"path/filepath" "path/filepath"
"regexp" "regexp"
"strings" "strings"
"time" "time"
"spark-mcp-go/internal/storage"
) )
var fileIDRegex = regexp.MustCompile(`^[0-9a-f]{32}$`) var fileIDRegex = regexp.MustCompile(`^[0-9a-f]{32}$`)
@@ -22,6 +26,20 @@ var fileIDRegex = regexp.MustCompile(`^[0-9a-f]{32}$`)
// Store is a local file store rooted at Root. // Store is a local file store rooted at Root.
type Store struct { type Store struct {
Root string Root string
repo UploadsDB
}
// UploadsDB is the minimal surface the Store needs from an upload index.
type UploadsDB interface {
Create(ctx context.Context, fileID, name string, size int64, sha256Hex string, uploadedAt time.Time) error
Get(ctx context.Context, fileID string) (storage.UploadMeta, error)
}
// SetRepo binds an upload index repository for dual-write. The DB is an
// index only; failures are logged but never fail the file write because
// .meta.json is the source of truth and startup backfill can recover.
func (s *Store) SetRepo(repo UploadsDB) {
s.repo = repo
} }
// New creates a Store rooted at root. It creates root with mode 0o750 if it // New creates a Store rooted at root. It creates root with mode 0o750 if it
@@ -111,6 +129,13 @@ func (s *Store) Save(data []byte, originalName string) (fileID, name string, siz
return "", "", 0, "", "", fmt.Errorf("uploads: write sidecar: %w", err) return "", "", 0, "", "", fmt.Errorf("uploads: write sidecar: %w", err)
} }
// Index the upload in the DB as a best-effort mirror of the sidecar.
if s.repo != nil {
if dbErr := s.repo.Create(context.Background(), fileID, originalName, meta.Size, sha256Hex, meta.UploadedAt); dbErr != nil {
slog.Default().Warn("uploads: failed to index upload in DB", "file_id", fileID, "err", dbErr)
}
}
return fileID, originalName, meta.Size, sha256Hex, absPath, nil return fileID, originalName, meta.Size, sha256Hex, absPath, nil
} }
@@ -240,6 +265,54 @@ func (s *Store) Sweep(ctx context.Context, ttl time.Duration) (deleted int, err
return deleted, nil return deleted, nil
} }
// Backfill walks Root and inserts a DB index row for every existing
// .meta.json sidecar that is not already indexed. It is the recovery path
// for sidecars created before the DB table existed.
func (s *Store) Backfill(ctx context.Context, repo UploadsDB) (inserted int, err error) {
entries, err := os.ReadDir(s.Root)
if err != nil {
return 0, fmt.Errorf("uploads: backfill read root: %w", err)
}
for _, e := range entries {
if err := ctx.Err(); err != nil {
return inserted, err
}
name := e.Name()
base := strings.TrimSuffix(name, ".meta.json")
if base == name || !fileIDRegex.MatchString(base) {
continue
}
metaPath := filepath.Join(s.Root, name)
sc, err := readSidecar(metaPath)
if err != nil {
// Skip corrupt sidecars; Sweep will reap them later.
continue
}
_, getErr := repo.Get(ctx, base)
if getErr == nil {
continue
}
if !errors.Is(getErr, storage.ErrNotFound) {
return inserted, fmt.Errorf("uploads: backfill lookup %s: %w", base, getErr)
}
createErr := repo.Create(ctx, base, sc.Name, sc.Size, sc.Sha256, sc.UploadedAt)
if createErr != nil {
if strings.Contains(createErr.Error(), "UNIQUE constraint failed") {
continue
}
return inserted, fmt.Errorf("uploads: backfill index %s: %w", base, createErr)
}
inserted++
}
return inserted, nil
}
func newFileID() (string, error) { func newFileID() (string, error) {
var b [16]byte var b [16]byte
if _, err := rand.Read(b[:]); err != nil { if _, err := rand.Read(b[:]); err != nil {
@@ -249,16 +322,24 @@ func newFileID() (string, error) {
} }
func readUploadedAt(path string) (time.Time, error) { func readUploadedAt(path string) (time.Time, error) {
data, err := os.ReadFile(path) sc, err := readSidecar(path)
if err != nil { if err != nil {
return time.Time{}, err return time.Time{}, err
} }
var sc sidecar
if err := json.Unmarshal(data, &sc); err != nil {
return time.Time{}, err
}
if sc.UploadedAt.IsZero() { if sc.UploadedAt.IsZero() {
return time.Time{}, fmt.Errorf("uploads: missing uploaded_at") return time.Time{}, fmt.Errorf("uploads: missing uploaded_at")
} }
return sc.UploadedAt, nil return sc.UploadedAt, nil
} }
func readSidecar(path string) (sidecar, error) {
var sc sidecar
data, err := os.ReadFile(path)
if err != nil {
return sc, err
}
if err := json.Unmarshal(data, &sc); err != nil {
return sc, err
}
return sc, nil
}
+180
View File
@@ -5,12 +5,15 @@ import (
"crypto/sha256" "crypto/sha256"
"encoding/hex" "encoding/hex"
"encoding/json" "encoding/json"
"errors"
"fmt" "fmt"
"os" "os"
"path/filepath" "path/filepath"
"regexp" "regexp"
"testing" "testing"
"time" "time"
"spark-mcp-go/internal/storage"
) )
func TestStore_SaveAndRetrieve(t *testing.T) { func TestStore_SaveAndRetrieve(t *testing.T) {
@@ -351,3 +354,180 @@ func TestStore_Sweep_DeletesDataWithCorruptSidecar(t *testing.T) {
t.Errorf("corrupt sidecar still exists") t.Errorf("corrupt sidecar still exists")
} }
} }
type fakeUploadRepo struct {
calls []storage.UploadMeta
err error
}
func (f *fakeUploadRepo) Create(ctx context.Context, fileID, name string, size int64, sha256Hex string, uploadedAt time.Time) error {
f.calls = append(f.calls, storage.UploadMeta{
FileID: fileID, Name: name, Size: size, SHA256: sha256Hex, UploadedAt: uploadedAt,
})
return f.err
}
func (f *fakeUploadRepo) Get(ctx context.Context, fileID string) (storage.UploadMeta, error) {
for _, m := range f.calls {
if m.FileID == fileID {
return m, nil
}
}
return storage.UploadMeta{}, storage.ErrNotFound
}
func (f *fakeUploadRepo) List(ctx context.Context, search string, limit int) ([]storage.UploadMeta, error) {
return f.calls, nil
}
func (f *fakeUploadRepo) Delete(ctx context.Context, fileID string) error {
return nil
}
func TestStore_Save_DualWrite(t *testing.T) {
store, err := New(t.TempDir())
if err != nil {
t.Fatalf("new store: %v", err)
}
repo := &fakeUploadRepo{}
store.SetRepo(repo)
data := []byte("dual-write test")
fileID, _, size, sha256Hex, absPath, err := store.Save(data, "dual.txt")
if err != nil {
t.Fatalf("save: %v", err)
}
if absPath == "" {
t.Fatal("absPath empty")
}
if len(repo.calls) != 1 {
t.Fatalf("repo.Create calls: got %d, want 1", len(repo.calls))
}
call := repo.calls[0]
if call.FileID != fileID {
t.Errorf("FileID: got %q, want %q", call.FileID, fileID)
}
if call.Name != "dual.txt" {
t.Errorf("Name: got %q, want dual.txt", call.Name)
}
if call.Size != size {
t.Errorf("Size: got %d, want %d", call.Size, size)
}
if call.SHA256 != sha256Hex {
t.Errorf("SHA256: got %q, want %q", call.SHA256, sha256Hex)
}
if call.UploadedAt.IsZero() {
t.Errorf("UploadedAt is zero")
}
}
func TestStore_Save_RepoErrorDoesNotFailUpload(t *testing.T) {
store, err := New(t.TempDir())
if err != nil {
t.Fatalf("new store: %v", err)
}
repo := &fakeUploadRepo{err: errors.New("db down")}
store.SetRepo(repo)
data := []byte("db error test")
fileID, _, _, _, absPath, err := store.Save(data, "error.txt")
if err != nil {
t.Fatalf("save should not fail when repo errors: %v", err)
}
if _, err := os.Stat(absPath); err != nil {
t.Errorf("data file missing: %v", err)
}
if _, err := os.Stat(absPath + ".meta.json"); err != nil {
t.Errorf("sidecar missing: %v", err)
}
if len(repo.calls) != 1 {
t.Errorf("repo.Create calls: got %d, want 1", len(repo.calls))
}
if repo.calls[0].FileID != fileID {
t.Errorf("repo call file_id: got %q, want %q", repo.calls[0].FileID, fileID)
}
}
func TestStore_Backfill(t *testing.T) {
store, err := New(t.TempDir())
if err != nil {
t.Fatalf("new store: %v", err)
}
repo := &fakeUploadRepo{}
now := time.Now()
ids := []string{
"00000000000000000000000000000001",
"00000000000000000000000000000002",
"00000000000000000000000000000003",
}
for i, id := range ids {
sc := map[string]any{
"name": fmt.Sprintf("file%d.txt", i),
"size": i + 1,
"sha256": fmt.Sprintf("sha%d", i),
"uploaded_at": now.Add(time.Duration(i) * time.Second).Format(time.RFC3339Nano),
}
b, _ := json.Marshal(sc)
if err := os.WriteFile(filepath.Join(store.Root, id+".meta.json"), b, 0o600); err != nil {
t.Fatalf("create sidecar %s: %v", id, err)
}
}
inserted, err := store.Backfill(context.Background(), repo)
if err != nil {
t.Fatalf("backfill: %v", err)
}
if inserted != 3 {
t.Errorf("inserted=%d, want 3", inserted)
}
if len(repo.calls) != 3 {
t.Errorf("repo.Create calls: got %d, want 3", len(repo.calls))
}
// Second backfill must skip existing rows.
inserted, err = store.Backfill(context.Background(), repo)
if err != nil {
t.Fatalf("second backfill: %v", err)
}
if inserted != 0 {
t.Errorf("second inserted=%d, want 0", inserted)
}
if len(repo.calls) != 3 {
t.Errorf("repo.Create calls after second backfill: got %d, want 3", len(repo.calls))
}
}
func TestStore_Backfill_SkipsExisting(t *testing.T) {
store, err := New(t.TempDir())
if err != nil {
t.Fatalf("new store: %v", err)
}
repo := &fakeUploadRepo{}
id := "00000000000000000000000000000004"
now := time.Now()
sc := map[string]any{
"name": "existing.txt",
"size": 5,
"sha256": "sha",
"uploaded_at": now.Format(time.RFC3339Nano),
}
b, _ := json.Marshal(sc)
if err := os.WriteFile(filepath.Join(store.Root, id+".meta.json"), b, 0o600); err != nil {
t.Fatalf("create sidecar: %v", err)
}
// Pre-seed the repo so the sidecar is already indexed.
if err := repo.Create(context.Background(), id, "existing.txt", 5, "sha", now); err != nil {
t.Fatalf("seed repo: %v", err)
}
inserted, err := store.Backfill(context.Background(), repo)
if err != nil {
t.Fatalf("backfill: %v", err)
}
if inserted != 0 {
t.Errorf("inserted=%d, want 0", inserted)
}
}