Compare commits
38
Commits
1a7c692e79
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2f026946b6 | ||
|
|
c02479381f | ||
|
|
ec5292f9f7 | ||
|
|
fefd816a21 | ||
|
|
2571ce69c0 | ||
|
|
160c9f32c8 | ||
|
|
e5f7d661b4 | ||
|
|
6f4c3a3e21 | ||
|
|
44a30c4a8a | ||
|
|
287d098e06 | ||
|
|
e2bbbd5c45 | ||
|
|
1cf4f90452 | ||
|
|
2bd67a5afc | ||
|
|
1460303b3e | ||
|
|
02499922b4 | ||
|
|
25c2fc66be | ||
|
|
80e644483a | ||
|
|
25f93dfb59 | ||
|
|
573397dbb4 | ||
|
|
a2232cecc4 | ||
|
|
d6e3952a4b | ||
|
|
df9b01272f | ||
|
|
876f1c9edb | ||
|
|
1796b5b872 | ||
|
|
4af166ea48 | ||
|
|
157848d87a | ||
|
|
97f6184f65 | ||
|
|
dab4cd062e | ||
|
|
4a28886502 | ||
|
|
67d1e9a6c5 | ||
|
|
44c9f415b7 | ||
|
|
c9e54a3f27 | ||
|
|
7920dc9597 | ||
|
|
2570713fbf | ||
|
|
f374ebfdb4 | ||
|
|
210aeb1315 | ||
|
|
475e353850 | ||
|
|
3bc8bd49f6 |
@@ -22,6 +22,9 @@ MAX_RESPONSE_BYTES=1048576
|
||||
# Timeout for spark-submit process execution (default: 60s)
|
||||
SPARK_SUBMIT_TIMEOUT=60s
|
||||
|
||||
# How long uploaded files are kept. Swept on server startup. Use Go duration syntax (e.g. 24h, 168h, 720h).
|
||||
UPLOAD_TTL=168h
|
||||
|
||||
# Log directory (default: ./data/logs)
|
||||
LOG_DIR=./data/logs
|
||||
|
||||
|
||||
Executable
+37
@@ -0,0 +1,37 @@
|
||||
#!/usr/bin/env bash
|
||||
# bark-notify.sh — push a notification to the local Bark server.
|
||||
#
|
||||
# Usage: bark-notify.sh <title> <body>
|
||||
#
|
||||
# The Bark endpoint, device_key, and group are pinned in this script
|
||||
# (it is a deployment detail, not a secret). The group is "Gitea" —
|
||||
# change it here if you want to sort notifications by category.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
if [ "$#" -ne 2 ]; then
|
||||
echo "usage: $0 <title> <body>" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
TITLE="$1"
|
||||
BODY="$2"
|
||||
|
||||
PAYLOAD=$(cat <<EOF
|
||||
{
|
||||
"device_key": "6ZFVsxM95cuAjYoNLWSWaN",
|
||||
"title": "${TITLE}",
|
||||
"body": "${BODY}",
|
||||
"group": "Gitea",
|
||||
"isArchive": 1,
|
||||
"ttl": 3600
|
||||
}
|
||||
EOF
|
||||
)
|
||||
|
||||
curl -s \
|
||||
-X POST http://101.43.40.124:81/push \
|
||||
-H "Content-Type: application/json; charset=utf-8" \
|
||||
-d "${PAYLOAD}"
|
||||
|
||||
echo "bark notified: ${TITLE}"
|
||||
@@ -0,0 +1,132 @@
|
||||
# This workflow runs on Gitea Actions, which is API-compatible with GitHub Actions.
|
||||
# It performs CI on every push to main and every pull request, and automatically
|
||||
# builds and attaches a Linux amd64 release binary when a version tag (v*) is pushed.
|
||||
#
|
||||
# Go is installed manually from a Chinese mirror (NJU) instead of using
|
||||
# actions/setup-go, because the runner is in a region where the official Go
|
||||
# download endpoint is slow. GOPROXY is set to goproxy.cn for the same reason.
|
||||
name: CI / Release
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
tags:
|
||||
- 'v*'
|
||||
pull_request:
|
||||
|
||||
jobs:
|
||||
ci:
|
||||
name: CI
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v6
|
||||
|
||||
- name: Install Go (Mirror in China)
|
||||
run: |
|
||||
if [ ! -d "/usr/local/go" ]; then
|
||||
wget https://mirrors.nju.edu.cn/golang/go1.25.5.linux-amd64.tar.gz
|
||||
sudo tar -C /usr/local -xzf go1.25.5.linux-amd64.tar.gz
|
||||
rm go1.25.5.linux-amd64.tar.gz
|
||||
fi
|
||||
echo "/usr/local/go/bin" >> $GITHUB_PATH
|
||||
echo "GOPROXY=https://goproxy.cn/,https://goproxy.io,direct" >> $GITHUB_ENV
|
||||
|
||||
- name: Build
|
||||
run: go build ./...
|
||||
|
||||
- name: Vet
|
||||
run: go vet ./...
|
||||
|
||||
- name: Test
|
||||
run: go test ./...
|
||||
|
||||
- name: Check formatting
|
||||
run: |
|
||||
if [ -n "$(gofmt -l .)" ]; then
|
||||
gofmt -l .
|
||||
exit 1
|
||||
fi
|
||||
|
||||
- name: Notify CI via Bark
|
||||
if: always()
|
||||
run: |
|
||||
REF_NAME="${{ github.ref_name }}"
|
||||
STATUS="${{ job.status }}"
|
||||
if [ "${STATUS}" = "success" ]; then
|
||||
EMOJI="✅"
|
||||
else
|
||||
EMOJI="❌"
|
||||
fi
|
||||
BODY="${EMOJI} ${REF_NAME} ${STATUS}: ${GITHUB_SERVER_URL}/${GITHUB_REPOSITORY}/actions/runs/${{ github.run_id }}"
|
||||
bash .github/scripts/bark-notify.sh "CI ${REF_NAME} ${STATUS}" "${BODY}"
|
||||
|
||||
release:
|
||||
name: Release
|
||||
needs: ci
|
||||
if: startsWith(github.ref, 'refs/tags/v')
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v6
|
||||
|
||||
- name: Install Go (Mirror in China)
|
||||
run: |
|
||||
if [ ! -d "/usr/local/go" ]; then
|
||||
wget https://mirrors.nju.edu.cn/golang/go1.25.5.linux-amd64.tar.gz
|
||||
sudo tar -C /usr/local -xzf go1.25.5.linux-amd64.tar.gz
|
||||
rm go1.25.5.linux-amd64.tar.gz
|
||||
fi
|
||||
echo "/usr/local/go/bin" >> $GITHUB_PATH
|
||||
echo "GOPROXY=https://goproxy.cn/,https://goproxy.io,direct" >> $GITHUB_ENV
|
||||
|
||||
- name: Build release binary
|
||||
run: GOOS=linux GOARCH=amd64 go build -o spark-mcp-linux-amd64 ./main.go
|
||||
|
||||
- name: Upload release asset via Gitea API
|
||||
env:
|
||||
GITEA_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
run: |
|
||||
# 1. 获取当前 Tag 名字 (例如 v0.0.1-beta)
|
||||
TAG_NAME="${{ github.ref_name }}"
|
||||
|
||||
echo "正在为 Tag ${TAG_NAME} 创建 Gitea Release..."
|
||||
|
||||
# 2. 调用 Gitea API 创建一个 Release 并获取其 ID
|
||||
# 注意:${{ github.api_url }} 会自动解析为你本地 Gitea 的 API 地址 (如 http://your-gitea/api/v1)
|
||||
RELEASE_ID=$(curl -X POST "${{ github.api_url }}/repos/${{ github.repository }}/releases" \
|
||||
-H "accept: application/json" \
|
||||
-H "authorization: token ${GITEA_TOKEN}" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d "{
|
||||
\"tag_name\": \"${TAG_NAME}\",
|
||||
\"target_commitish\": \"${{ github.sha }}\",
|
||||
\"name\": \"Release ${TAG_NAME}\",
|
||||
\"body\": \"Auto generated release assets.\",
|
||||
\"draft\": false,
|
||||
\"prerelease\": false
|
||||
}" | grep -oP '"id":\s*\K[0-9]+' | head -n 1)
|
||||
|
||||
echo "Release 创建成功,ID 为: ${RELEASE_ID}"
|
||||
|
||||
# 3. 将你的二进制文件上传到该 Release 中
|
||||
echo "开始上传文件 spark-mcp-linux-amd64 ..."
|
||||
curl -X POST "${{ github.api_url }}/repos/${{ github.repository }}/releases/${RELEASE_ID}/assets" \
|
||||
-H "accept: application/json" \
|
||||
-H "authorization: token ${GITEA_TOKEN}" \
|
||||
-H "Content-Type: multipart/form-data" \
|
||||
-F "attachment=@spark-mcp-linux-amd64"
|
||||
|
||||
- name: Notify release via Bark
|
||||
if: always()
|
||||
run: |
|
||||
TITLE="${{ github.ref_name }}"
|
||||
STATUS="${{ job.status }}"
|
||||
if [ "${STATUS}" = "success" ]; then
|
||||
EMOJI="✅"
|
||||
else
|
||||
EMOJI="❌"
|
||||
fi
|
||||
BODY="${EMOJI} spark-mcp ${TITLE} ${STATUS}: ${GITHUB_SERVER_URL}/${GITHUB_REPOSITORY}/releases/tag/${TITLE}"
|
||||
bash .github/scripts/bark-notify.sh "${TITLE} ${STATUS}" "${BODY}"
|
||||
@@ -29,6 +29,8 @@ Standard Go toolchain (Go 1.25.5 per `go.mod`). All commands run from the repo r
|
||||
|
||||
There is no `Makefile`, no `golangci-lint` config, and no test files yet — do not assume they exist.
|
||||
|
||||
Cross-compile outputs are gitignored via the pattern `spark-mcp-*-*` (see `.gitignore`); do not commit stray build artifacts.
|
||||
|
||||
## Key Pinned Dependencies
|
||||
|
||||
`go.mod` declares these as `// indirect` (they will become direct once code is added that imports them):
|
||||
@@ -56,3 +58,71 @@ The codebase is empty enough that a few decisions are not yet pinned down and wi
|
||||
## Memory
|
||||
|
||||
This is a fresh repo with no `MEMORY.md` or project-level memory file. Decisions made during early development (transport choice, Spark backend interface, persistence schema) are worth recording once they are made — see the `claude-md-management` skill conventions.
|
||||
|
||||
---
|
||||
|
||||
# Working in this Repo
|
||||
|
||||
The rules in `~/.claude/CLAUDE.md` apply (Codex-first routing, mandatory Codex MCP invocation parameters, Linus three-question review, 4-phase workflow, Codex prompt template, auto-confirm vs pause boundaries). This section adds the **project-specific calibration** for applying them here.
|
||||
|
||||
## Routing: Codex vs CC
|
||||
|
||||
| Task in this repo | Executor | Why |
|
||||
|---|---|---|
|
||||
| Edit `main.go` to wire `cfg.GinMode` | Codex | One-line code change in production path |
|
||||
| Add `loadDotenv` to `internal/config/config.go` | Codex | New function, real logic |
|
||||
| Add a new test file | Codex | Even small, code generation is the strength |
|
||||
| Edit `internal/httpclient/ssrf.go` to silence S1017 | Codex | Behavior must stay identical — needs verification |
|
||||
| Update `.gitignore` line | CC | Single-line, no logic, mechanical |
|
||||
| Update this `CLAUDE.md` | CC | Pure docs |
|
||||
| Fix typo in a comment | CC | <20 lines, no logic |
|
||||
| Update `.env.example` header comment | CC | Single line of prose |
|
||||
|
||||
**Threshold for "trivial enough for CC" in this repo:** purely textual, no code generation, <20 lines, no risk of subtle behavior change. When in doubt, default to Codex — the cost difference is small and Codex is more consistent at preserving invariants.
|
||||
|
||||
## Codebase Patterns to Reuse
|
||||
|
||||
When working in a package, **read the package's existing test file first** to see the established style before writing new code:
|
||||
|
||||
- **HTTP client mocks** — `internal/httpclient/redirect_test.go` uses `httptest.NewServer` with path-switch handlers. For SSRF / redirect tests, pass `"127.0.0.1"` as the allowed host.
|
||||
- **slog capture in tests** — `internal/logging/tool_call_test.go` shows the pattern: `bytes.Buffer` + `slog.NewJSONHandler`, decode each line with `json.Unmarshal`, assert on `rec["msg"]`, `rec["level"]`, `rec["tool"]`. Reuse this for any new logger test.
|
||||
- **MCP tool registration** — every tool lives in `internal/mcp/tools/<name>.go`. Schema is the first arg, handler is a method on the package's client struct. The tool's Go test file goes in the same directory.
|
||||
|
||||
## Pre-Commit Checklist (project-specific)
|
||||
|
||||
In addition to the global acceptance list, this repo's convention is:
|
||||
|
||||
- [ ] `go build ./...` — zero errors
|
||||
- [ ] `go vet ./...` — zero warnings (catches the stringsseq / S1017 family)
|
||||
- [ ] `go test ./...` — all green; new tests added for new behavior
|
||||
- [ ] `gofmt -l .` — empty output (no diffs pending)
|
||||
- [ ] No stray `spark-mcp-*-*` binary in repo root
|
||||
- [ ] If behavior gap is intentionally deferred: add a deferred test that **locks the current behavior** (so a future fix must update the test deliberately) and a `TODO(<area>)` comment in the source with rationale
|
||||
- [ ] If a public function signature changed: explain backward compatibility in the commit body
|
||||
|
||||
## Commit Message Convention
|
||||
|
||||
Observed in this repo's recent history:
|
||||
|
||||
```
|
||||
<package>: <verb-led summary in Chinese>
|
||||
|
||||
<optional body: motivation, design tradeoffs, why-not-fix-X, what tests were added>
|
||||
<optional "未提交: <file>" trailer for intentionally uncommitted artifacts>
|
||||
|
||||
Co-Authored-By: Claude <noreply@anthropic.com>
|
||||
```
|
||||
|
||||
- Subject line: `<package>:` prefix (e.g. `httpclient:`, `rm:`, `admin:`, `config:`, `logging:`, `gitignore:`), then a short verb-led summary in Chinese.
|
||||
- Body in Chinese; technical English terms (RFC numbers, status codes, test names, package paths) are fine in code form.
|
||||
- Always end with the `Co-Authored-By` trailer.
|
||||
- Body should answer: *what was the gap, why this fix, what could it have broken, what tests lock it in.* This is what makes the commit history a useful design record.
|
||||
|
||||
## Project-Specific Gotchas
|
||||
|
||||
- **`.env` is uncommitted and gitignored** — never try to read it, and never claim to know what's in it. Treat its existence as opaque.
|
||||
- **Cross-compile artifacts** — `go build -o spark-mcp-linux-amd64 .` and similar produce files covered by the `spark-mcp-*-*` gitignore. Don't try to `git add` them.
|
||||
- **`internal/` is enforced** — Go's `internal/` rule means sub-packages of `spark-mcp-go/internal/...` are importable only by `spark-mcp-go/...`. Cross-repo imports will fail; design packages accordingly.
|
||||
- **Deferred-behavior tests are intentional** — `TestDoWithRedirect_DeferredRFC9110Gaps` is a *gap-locking* test, not a regression test. The assertions verify the gap, not the fix. Read the comment block above the test before "fixing" what looks like a bug.
|
||||
- **YARN / SHS API surface is the source of truth** — the upstream REST API behavior (e.g. amContainerLogs being a 307 redirect, not a 303) drives the design. Don't invent redirect semantics; check what the cluster actually returns.
|
||||
- **`log/slog` uses a `multiHandler` for dual terminal+file writing** — the `base *slog.Logger` passed to `StartToolCall` already routes through the main logger. Adding a new event is one `base.Info(...)` call; do not add a separate handler.
|
||||
|
||||
+151
-1
@@ -5,8 +5,12 @@ package admin
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
|
||||
@@ -14,11 +18,12 @@ import (
|
||||
"spark-mcp-go/internal/cluster"
|
||||
"spark-mcp-go/internal/middleware"
|
||||
"spark-mcp-go/internal/storage"
|
||||
"spark-mcp-go/internal/uploads"
|
||||
)
|
||||
|
||||
// Mount attaches the /admin sub-router to r, protecting every route with
|
||||
// bearer-token admin authentication.
|
||||
func Mount(r *gin.Engine, repo *storage.ClusterRepo, auditRepo *audit.Repo, adminTokens []string) {
|
||||
func Mount(r *gin.Engine, repo *storage.ClusterRepo, uploadRepo *storage.UploadRepo, submissionRepo *storage.SubmissionRepo, auditRepo *audit.Repo, uploadStore *uploads.Store, adminTokens []string) {
|
||||
// HTML page is public — its own modal prompts for the token.
|
||||
// All other /admin/* endpoints (API + OpenAPI docs) still require auth.
|
||||
gPublic := r.Group("/admin")
|
||||
@@ -32,6 +37,12 @@ func Mount(r *gin.Engine, repo *storage.ClusterRepo, auditRepo *audit.Repo, admi
|
||||
g.PUT("/clusters/:id", updateCluster(repo, auditRepo))
|
||||
g.DELETE("/clusters/:id", deleteCluster(repo, auditRepo))
|
||||
g.GET("/audit", listAudit(auditRepo))
|
||||
g.GET("/uploads", uploadsWebHandler)
|
||||
g.GET("/uploads/api", listUploads(uploadRepo))
|
||||
g.GET("/uploads/:id/download", downloadUpload(uploadRepo, uploadStore))
|
||||
g.DELETE("/uploads/:id", deleteUpload(uploadRepo, auditRepo, uploadStore))
|
||||
g.GET("/submissions", submissionsWebHandler)
|
||||
g.GET("/submissions/api", listSubmissions(submissionRepo))
|
||||
g.GET("/docs", DocsHandler)
|
||||
g.GET("/docs/spec", OpenAPISpecHandler)
|
||||
}
|
||||
@@ -170,6 +181,10 @@ func updateCluster(repo *storage.ClusterRepo, auditRepo *audit.Repo) gin.Handler
|
||||
merged.DefaultSubmitArgs = decodeStringSlice(v)
|
||||
}
|
||||
|
||||
if len(merged.DefaultSubmitArgs) > 0 {
|
||||
slog.Default().Info("admin: cluster sets default_submit_args which is deprecated; the spark_submit Tool no longer prepends these args. Configure flags via spark_conf / extra_args in the structured tool call.", "cluster_id", merged.ID)
|
||||
}
|
||||
|
||||
// Re-validate after merging so a missing auth_type on a fresh
|
||||
// cluster (defaulted by DB to "none") doesn't get clobbered.
|
||||
if err := validateClusterUpdate(&merged); err != nil {
|
||||
@@ -270,6 +285,141 @@ func deleteCluster(repo *storage.ClusterRepo, auditRepo *audit.Repo) gin.Handler
|
||||
}
|
||||
}
|
||||
|
||||
func listUploads(repo *storage.UploadRepo) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
search := c.Query("search")
|
||||
limitStr := c.DefaultQuery("limit", "100")
|
||||
limit, err := strconv.Atoi(limitStr)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid limit"})
|
||||
return
|
||||
}
|
||||
|
||||
uploads, err := repo.List(c.Request.Context(), search, limit)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, uploads)
|
||||
}
|
||||
}
|
||||
|
||||
func listSubmissions(repo *storage.SubmissionRepo) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
limitStr := c.DefaultQuery("limit", "100")
|
||||
limit, err := strconv.Atoi(limitStr)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid limit"})
|
||||
return
|
||||
}
|
||||
offsetStr := c.DefaultQuery("offset", "0")
|
||||
offset, err := strconv.Atoi(offsetStr)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid offset"})
|
||||
return
|
||||
}
|
||||
|
||||
subs, err := repo.List(c.Request.Context(), storage.SubmissionFilter{
|
||||
Search: c.Query("search"),
|
||||
FileID: c.Query("file_id"),
|
||||
ClusterID: c.Query("cluster_id"),
|
||||
AppID: c.Query("app_id"),
|
||||
Limit: limit,
|
||||
Offset: offset,
|
||||
})
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, subs)
|
||||
}
|
||||
}
|
||||
|
||||
func downloadUpload(repo *storage.UploadRepo, store *uploads.Store) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
ctx := c.Request.Context()
|
||||
|
||||
meta, err := repo.Get(ctx, id)
|
||||
if err != nil {
|
||||
if errors.Is(err, storage.ErrNotFound) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "upload not found"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
dataPath, err := store.AbsPath(id)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
data, err := os.ReadFile(dataPath)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "upload file missing"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.Header("Content-Disposition", "attachment; filename=\""+escapeDispositionFilename(meta.Name)+"\"")
|
||||
c.Data(http.StatusOK, "application/octet-stream", data)
|
||||
}
|
||||
}
|
||||
|
||||
// escapeDispositionFilename escapes quotes and backslashes for the
|
||||
// Content-Disposition filename parameter. The upload validator already rejects
|
||||
// names containing quotes or newlines; this is defense in depth.
|
||||
func escapeDispositionFilename(name string) string {
|
||||
name = strings.ReplaceAll(name, "\\", "\\\\")
|
||||
name = strings.ReplaceAll(name, "\"", "\\\"")
|
||||
return name
|
||||
}
|
||||
|
||||
func deleteUpload(repo *storage.UploadRepo, auditRepo *audit.Repo, store *uploads.Store) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
ctx := c.Request.Context()
|
||||
|
||||
meta, err := repo.Get(ctx, id)
|
||||
if err != nil && !errors.Is(err, storage.ErrNotFound) {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if err := repo.Delete(ctx, id); err != nil && !errors.Is(err, storage.ErrNotFound) {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// Always attempt to remove the on-disk files; the DB is only an index.
|
||||
if store != nil {
|
||||
dataPath := filepath.Join(store.Root, id)
|
||||
_ = os.Remove(dataPath)
|
||||
_ = os.Remove(dataPath + ".meta.json")
|
||||
}
|
||||
|
||||
if meta.FileID != "" {
|
||||
details, _ := audit.MarshalDetails(map[string]any{
|
||||
"file_id": id,
|
||||
"name": meta.Name,
|
||||
"size": meta.Size,
|
||||
"sha256": meta.SHA256,
|
||||
})
|
||||
_ = auditRepo.Insert(ctx, &audit.Entry{
|
||||
Actor: actor(c),
|
||||
Action: audit.ActionUploadDelete,
|
||||
ClusterID: id,
|
||||
Details: details,
|
||||
})
|
||||
}
|
||||
|
||||
c.Status(http.StatusNoContent)
|
||||
}
|
||||
}
|
||||
func listAudit(auditRepo *audit.Repo) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
limitStr := c.DefaultQuery("limit", "100")
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
@@ -15,6 +16,7 @@ import (
|
||||
"spark-mcp-go/internal/audit"
|
||||
"spark-mcp-go/internal/cluster"
|
||||
"spark-mcp-go/internal/storage"
|
||||
"spark-mcp-go/internal/uploads"
|
||||
)
|
||||
|
||||
func init() {
|
||||
@@ -29,7 +31,11 @@ func newTestAdmin(t *testing.T) (*gin.Engine, *storage.ClusterRepo, *audit.Repo)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
r := gin.New()
|
||||
Mount(r, db.Clusters(), audit.NewRepo(db), []string{"good-token"})
|
||||
uploadStore, err := uploads.New(filepath.Join(t.TempDir(), "uploads"))
|
||||
if err != nil {
|
||||
t.Fatalf("new upload store: %v", err)
|
||||
}
|
||||
Mount(r, db.Clusters(), db.Uploads(), db.Submissions(), audit.NewRepo(db), &uploadStore, []string{"good-token"})
|
||||
return r, db.Clusters(), audit.NewRepo(db)
|
||||
}
|
||||
|
||||
@@ -54,20 +60,20 @@ func doReq(t *testing.T, r *gin.Engine, method, path, token string, body any) *h
|
||||
|
||||
func fullClusterMap(id string) map[string]any {
|
||||
return map[string]any{
|
||||
"id": id,
|
||||
"name": id + "-name",
|
||||
"rm_url": "http://rm.example.com:8088",
|
||||
"shs_url": "http://shs.example.com:18080",
|
||||
"spark_submit_execute_bin": "/usr/bin/spark-submit",
|
||||
"is_active": true,
|
||||
"auth_type": "basic",
|
||||
"auth_username": "admin",
|
||||
"auth_password": "real-secret",
|
||||
"ssl_verify": false,
|
||||
"ssl_ca_bundle": "",
|
||||
"url_allowlist": []string{"*.example.com"},
|
||||
"default_submit_args": []string{"--master", "yarn"},
|
||||
"rate_limit_per_min": 10,
|
||||
"id": id,
|
||||
"name": id + "-name",
|
||||
"rm_url": "http://rm.example.com:8088",
|
||||
"shs_url": "http://shs.example.com:18080",
|
||||
"spark_submit_execute_bin": "/usr/bin/spark-submit",
|
||||
"is_active": true,
|
||||
"auth_type": "basic",
|
||||
"auth_username": "admin",
|
||||
"auth_password": "real-secret",
|
||||
"ssl_verify": false,
|
||||
"ssl_ca_bundle": "",
|
||||
"url_allowlist": []string{"*.example.com"},
|
||||
"default_submit_args": []string{"--master", "yarn"},
|
||||
"rate_limit_per_min": 10,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,155 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
|
||||
"spark-mcp-go/internal/storage"
|
||||
"spark-mcp-go/internal/uploads"
|
||||
)
|
||||
|
||||
func newTestAdminSubmissions(t *testing.T) (*gin.Engine, *storage.SubmissionRepo, *storage.UploadRepo, *uploads.Store) {
|
||||
t.Helper()
|
||||
db, err := storage.Open(":memory:")
|
||||
if err != nil {
|
||||
t.Fatalf("open: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
|
||||
uploadStore, err := uploads.New(filepath.Join(t.TempDir(), "uploads"))
|
||||
if err != nil {
|
||||
t.Fatalf("new upload store: %v", err)
|
||||
}
|
||||
|
||||
r := gin.New()
|
||||
Mount(r, db.Clusters(), db.Uploads(), db.Submissions(), nil, &uploadStore, []string{"good-token"})
|
||||
return r, db.Submissions(), db.Uploads(), &uploadStore
|
||||
}
|
||||
|
||||
func TestSubmissionsWeb_RequiresAuth(t *testing.T) {
|
||||
r, _, _, _ := newTestAdminSubmissions(t)
|
||||
|
||||
w := doReq(t, r, "GET", "/admin/submissions", "", nil)
|
||||
if w.Code != http.StatusUnauthorized {
|
||||
t.Errorf("GET /admin/submissions without auth: got %d, want %d", w.Code, http.StatusUnauthorized)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSubmissionsWeb_WithAuth(t *testing.T) {
|
||||
r, _, _, _ := newTestAdminSubmissions(t)
|
||||
|
||||
w := doReq(t, r, "GET", "/admin/submissions", "good-token", nil)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("GET /admin/submissions with auth: got %d, want %d", w.Code, http.StatusOK)
|
||||
}
|
||||
ct := w.Header().Get("Content-Type")
|
||||
if !strings.Contains(ct, "text/html") {
|
||||
t.Errorf("Content-Type = %q, want text/html", ct)
|
||||
}
|
||||
body := w.Body.String()
|
||||
if !strings.Contains(body, "href=\"/admin/uploads\"") && !strings.Contains(body, "href='/admin/uploads'") {
|
||||
t.Errorf("body missing uploads nav link")
|
||||
}
|
||||
if !strings.Contains(body, "href=\"/admin/submissions\"") && !strings.Contains(body, "href='/admin/submissions'") {
|
||||
t.Errorf("body missing submissions nav link")
|
||||
}
|
||||
if !strings.Contains(body, "submissions-body") {
|
||||
t.Errorf("body missing submissions table")
|
||||
}
|
||||
}
|
||||
|
||||
func TestListSubmissions_SearchByAppID(t *testing.T) {
|
||||
r, subRepo, uploadRepo, _ := newTestAdminSubmissions(t)
|
||||
ctx := context.Background()
|
||||
|
||||
uploadedAt := time.Unix(0, time.Now().UnixNano())
|
||||
if err := uploadRepo.Create(ctx, "00000000000000000000000000000001", "alpha.py", 10, "deadbeef", uploadedAt); err != nil {
|
||||
t.Fatalf("create upload: %v", err)
|
||||
}
|
||||
if err := uploadRepo.Create(ctx, "00000000000000000000000000000002", "beta.py", 20, "cafebabe", uploadedAt.Add(time.Second)); err != nil {
|
||||
t.Fatalf("create upload: %v", err)
|
||||
}
|
||||
|
||||
submissions := []storage.Submission{
|
||||
{FileID: "00000000000000000000000000000001", AppID: "application_111_0001", TrackingURL: "http://a", ClusterID: "c1", ExitCode: 0, DurationMS: 100, SubmittedAt: time.Now()},
|
||||
{FileID: "00000000000000000000000000000002", AppID: "application_222_0001", TrackingURL: "http://b", ClusterID: "c1", ExitCode: 0, DurationMS: 100, SubmittedAt: time.Now().Add(time.Second)},
|
||||
}
|
||||
for _, s := range submissions {
|
||||
if err := subRepo.Create(ctx, s); err != nil {
|
||||
t.Fatalf("create submission: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
w := doReq(t, r, "GET", "/admin/submissions/api?search=application_111", "good-token", nil)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("search: got status %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
var list []map[string]any
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &list); err != nil {
|
||||
t.Fatalf("unmarshal search: %v", err)
|
||||
}
|
||||
if len(list) != 1 {
|
||||
t.Errorf("search result: got %d, want 1", len(list))
|
||||
}
|
||||
if len(list) > 0 && list[0]["app_id"] != "application_111_0001" {
|
||||
t.Errorf("app_id=%v, want application_111_0001", list[0]["app_id"])
|
||||
}
|
||||
|
||||
w = doReq(t, r, "GET", "/admin/submissions/api?search=beta.py", "good-token", nil)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("search by name: got status %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &list); err != nil {
|
||||
t.Fatalf("unmarshal search: %v", err)
|
||||
}
|
||||
if len(list) != 1 {
|
||||
t.Errorf("search by name result: got %d, want 1", len(list))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDownloadUpload(t *testing.T) {
|
||||
r, _, uploadRepo, store := newTestAdminSubmissions(t)
|
||||
ctx := context.Background()
|
||||
|
||||
id := "00000000000000000000000000000003"
|
||||
payload := []byte("hello payload")
|
||||
if err := uploadRepo.Create(ctx, id, "report.csv", int64(len(payload)), "feedface", time.Now()); err != nil {
|
||||
t.Fatalf("create upload: %v", err)
|
||||
}
|
||||
dataPath, err := store.AbsPath(id)
|
||||
if err != nil {
|
||||
t.Fatalf("abspath: %v", err)
|
||||
}
|
||||
if err := os.WriteFile(dataPath, payload, 0o640); err != nil {
|
||||
t.Fatalf("write data: %v", err)
|
||||
}
|
||||
|
||||
w := doReq(t, r, "GET", "/admin/uploads/"+id+"/download", "good-token", nil)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("download: got status %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if string(w.Body.Bytes()) != string(payload) {
|
||||
t.Errorf("body = %q, want %q", w.Body.Bytes(), payload)
|
||||
}
|
||||
cd := w.Header().Get("Content-Disposition")
|
||||
if !strings.Contains(cd, `filename="report.csv"`) {
|
||||
t.Errorf("Content-Disposition = %q, want filename=\"report.csv\"", cd)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDownloadUpload_NotFound(t *testing.T) {
|
||||
r, _, _, _ := newTestAdminSubmissions(t)
|
||||
|
||||
w := doReq(t, r, "GET", "/admin/uploads/00000000000000000000000000000099/download", "good-token", nil)
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Errorf("download unknown: got %d, want %d", w.Code, http.StatusNotFound)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,196 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
|
||||
"spark-mcp-go/internal/audit"
|
||||
"spark-mcp-go/internal/storage"
|
||||
"spark-mcp-go/internal/uploads"
|
||||
)
|
||||
|
||||
func newTestAdminUploads(t *testing.T) (*gin.Engine, *storage.UploadRepo, *storage.SubmissionRepo, *audit.Repo, *uploads.Store) {
|
||||
t.Helper()
|
||||
db, err := storage.Open(":memory:")
|
||||
if err != nil {
|
||||
t.Fatalf("open: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
|
||||
uploadStore, err := uploads.New(filepath.Join(t.TempDir(), "uploads"))
|
||||
if err != nil {
|
||||
t.Fatalf("new upload store: %v", err)
|
||||
}
|
||||
|
||||
r := gin.New()
|
||||
Mount(r, db.Clusters(), db.Uploads(), db.Submissions(), audit.NewRepo(db), &uploadStore, []string{"good-token"})
|
||||
return r, db.Uploads(), db.Submissions(), audit.NewRepo(db), &uploadStore
|
||||
}
|
||||
|
||||
func TestUploadsWeb_RequiresAuth(t *testing.T) {
|
||||
r, _, _, _, _ := newTestAdminUploads(t)
|
||||
|
||||
w := doReq(t, r, "GET", "/admin/uploads", "", nil)
|
||||
if w.Code != http.StatusUnauthorized {
|
||||
t.Errorf("GET /admin/uploads without auth: got %d, want %d", w.Code, http.StatusUnauthorized)
|
||||
}
|
||||
|
||||
w = doReq(t, r, "GET", "/admin/uploads", "good-token", nil)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Errorf("GET /admin/uploads with auth: got %d, want %d", w.Code, http.StatusOK)
|
||||
}
|
||||
ct := w.Header().Get("Content-Type")
|
||||
if !strings.Contains(ct, "text/html") {
|
||||
t.Errorf("Content-Type = %q, want text/html", ct)
|
||||
}
|
||||
body := w.Body.String()
|
||||
if !strings.Contains(body, "href='/admin/uploads'") && !strings.Contains(body, "href=\"/admin/uploads\"") {
|
||||
t.Errorf("body missing uploads nav link")
|
||||
}
|
||||
if !strings.Contains(body, "uploads-body") {
|
||||
t.Errorf("body missing uploads table")
|
||||
}
|
||||
}
|
||||
|
||||
func TestListUploads(t *testing.T) {
|
||||
r, repo, _, _, _ := newTestAdminUploads(t)
|
||||
ctx := context.Background()
|
||||
|
||||
uploadedAt := time.Unix(0, time.Now().UnixNano())
|
||||
if err := repo.Create(ctx, "00000000000000000000000000000001", "alpha.txt", 10, "deadbeef", uploadedAt); err != nil {
|
||||
t.Fatalf("create upload: %v", err)
|
||||
}
|
||||
if err := repo.Create(ctx, "00000000000000000000000000000002", "beta.txt", 20, "cafebabe", uploadedAt.Add(time.Second)); err != nil {
|
||||
t.Fatalf("create upload: %v", err)
|
||||
}
|
||||
|
||||
w := doReq(t, r, "GET", "/admin/uploads/api", "good-token", nil)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("got status %d, want %d", w.Code, http.StatusOK)
|
||||
}
|
||||
var list []map[string]any
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &list); err != nil {
|
||||
t.Fatalf("unmarshal list: %v", err)
|
||||
}
|
||||
if len(list) != 2 {
|
||||
t.Errorf("got %d uploads, want 2", len(list))
|
||||
}
|
||||
|
||||
w = doReq(t, r, "GET", "/admin/uploads/api?search=alpha", "good-token", nil)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("search: got status %d", w.Code)
|
||||
}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &list); err != nil {
|
||||
t.Fatalf("unmarshal search: %v", err)
|
||||
}
|
||||
if len(list) != 1 {
|
||||
t.Errorf("search result: got %d, want 1", len(list))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeleteUpload(t *testing.T) {
|
||||
r, repo, _, auditRepo, store := newTestAdminUploads(t)
|
||||
ctx := context.Background()
|
||||
|
||||
id := "00000000000000000000000000000003"
|
||||
dataPath := filepath.Join(store.Root, id)
|
||||
metaPath := dataPath + ".meta.json"
|
||||
|
||||
if err := os.WriteFile(dataPath, []byte("payload"), 0o640); err != nil {
|
||||
t.Fatalf("create data file: %v", err)
|
||||
}
|
||||
sc := map[string]any{
|
||||
"name": "delete-me.txt",
|
||||
"size": 7,
|
||||
"sha256": "feedface",
|
||||
"uploaded_at": time.Now().Format(time.RFC3339Nano),
|
||||
}
|
||||
b, _ := json.Marshal(sc)
|
||||
if err := os.WriteFile(metaPath, b, 0o600); err != nil {
|
||||
t.Fatalf("create sidecar: %v", err)
|
||||
}
|
||||
if err := repo.Create(ctx, id, "delete-me.txt", 7, "feedface", time.Now()); err != nil {
|
||||
t.Fatalf("create db row: %v", err)
|
||||
}
|
||||
|
||||
w := doReq(t, r, "DELETE", "/admin/uploads/"+id, "good-token", nil)
|
||||
if w.Code != http.StatusNoContent {
|
||||
t.Fatalf("delete: got status %d, want %d: %s", w.Code, http.StatusNoContent, w.Body.String())
|
||||
}
|
||||
|
||||
if _, err := repo.Get(ctx, id); !errors.Is(err, storage.ErrNotFound) {
|
||||
t.Errorf("db row still exists: %v", err)
|
||||
}
|
||||
if _, err := os.Stat(dataPath); !os.IsNotExist(err) {
|
||||
t.Errorf("data file still exists")
|
||||
}
|
||||
if _, err := os.Stat(metaPath); !os.IsNotExist(err) {
|
||||
t.Errorf("sidecar still exists")
|
||||
}
|
||||
|
||||
entries, err := auditRepo.List(ctx, 100)
|
||||
if err != nil {
|
||||
t.Fatalf("list audit: %v", err)
|
||||
}
|
||||
found := false
|
||||
for _, e := range entries {
|
||||
if e.Action == audit.ActionUploadDelete && e.ClusterID == id {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Errorf("missing upload.delete audit entry")
|
||||
}
|
||||
}
|
||||
|
||||
func TestListUploads_SearchByAppID(t *testing.T) {
|
||||
r, uploadRepo, subRepo, _, _ := newTestAdminUploads(t)
|
||||
ctx := context.Background()
|
||||
|
||||
uploadedAt := time.Unix(0, time.Now().UnixNano())
|
||||
if err := uploadRepo.Create(ctx, "00000000000000000000000000000010", "gamma.py", 10, "deadbeef", uploadedAt); err != nil {
|
||||
t.Fatalf("create upload: %v", err)
|
||||
}
|
||||
if err := uploadRepo.Create(ctx, "00000000000000000000000000000011", "delta.py", 20, "cafebabe", uploadedAt.Add(time.Second)); err != nil {
|
||||
t.Fatalf("create upload: %v", err)
|
||||
}
|
||||
|
||||
// Only gamma.py has an associated submission; searching by app_id should
|
||||
// still return it even though the filename doesn't match the query.
|
||||
if err := subRepo.Create(ctx, storage.Submission{
|
||||
FileID: "00000000000000000000000000000010",
|
||||
AppID: "application_gamma_0001",
|
||||
TrackingURL: "http://gamma",
|
||||
ClusterID: "c1",
|
||||
ExitCode: 0,
|
||||
DurationMS: 100,
|
||||
SubmittedAt: time.Now(),
|
||||
}); err != nil {
|
||||
t.Fatalf("create submission: %v", err)
|
||||
}
|
||||
|
||||
w := doReq(t, r, "GET", "/admin/uploads/api?search=application_gamma", "good-token", nil)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("search by app_id: got status %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
var list []map[string]any
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &list); err != nil {
|
||||
t.Fatalf("unmarshal search: %v", err)
|
||||
}
|
||||
if len(list) != 1 {
|
||||
t.Errorf("search by app_id result: got %d, want 1", len(list))
|
||||
}
|
||||
if len(list) > 0 && list[0]["file_id"] != "00000000000000000000000000000010" {
|
||||
t.Errorf("file_id=%v, want upload 10", list[0]["file_id"])
|
||||
}
|
||||
}
|
||||
+19
-1
@@ -7,7 +7,7 @@ import (
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
//go:embed web/cluster.html
|
||||
//go:embed web
|
||||
var webFS embed.FS
|
||||
|
||||
func webHandler(c *gin.Context) {
|
||||
@@ -18,3 +18,21 @@ func webHandler(c *gin.Context) {
|
||||
}
|
||||
c.Data(http.StatusOK, "text/html; charset=utf-8", data)
|
||||
}
|
||||
|
||||
func uploadsWebHandler(c *gin.Context) {
|
||||
data, err := webFS.ReadFile("web/uploads.html")
|
||||
if err != nil {
|
||||
c.String(http.StatusInternalServerError, "uploads.html: %v", err)
|
||||
return
|
||||
}
|
||||
c.Data(http.StatusOK, "text/html; charset=utf-8", data)
|
||||
}
|
||||
|
||||
func submissionsWebHandler(c *gin.Context) {
|
||||
data, err := webFS.ReadFile("web/submissions.html")
|
||||
if err != nil {
|
||||
c.String(http.StatusInternalServerError, "submissions.html: %v", err)
|
||||
return
|
||||
}
|
||||
c.Data(http.StatusOK, "text/html; charset=utf-8", data)
|
||||
}
|
||||
|
||||
@@ -34,6 +34,10 @@
|
||||
flex-wrap: wrap;
|
||||
}
|
||||
header h1 { margin: 0; font-size: 1.1rem; color: var(--accent); }
|
||||
nav { display: flex; gap: .75rem; }
|
||||
nav a { color: var(--muted); text-decoration: none; }
|
||||
nav a:hover { color: var(--text); }
|
||||
nav a.active { color: var(--accent); }
|
||||
.auth { display: flex; align-items: center; gap: .5rem; flex-wrap: wrap; }
|
||||
input[type='text'], input[type='password'], input[type='number'], select, textarea {
|
||||
background: var(--bg);
|
||||
@@ -111,7 +115,14 @@
|
||||
</head>
|
||||
<body>
|
||||
<header>
|
||||
<h1>spark-mcp-go Admin</h1>
|
||||
<div style='display:flex;align-items:center;gap:1rem;'>
|
||||
<h1>spark-mcp-go Admin</h1>
|
||||
<nav>
|
||||
<a href='/admin' class='active'>Clusters</a>
|
||||
<a href='/admin/uploads'>Uploads</a>
|
||||
<a href='/admin/submissions'>Submissions</a>
|
||||
</nav>
|
||||
</div>
|
||||
<div class='auth'>
|
||||
<button id='set-token' class='secondary'>Set Token</button>
|
||||
<button id='logout' class='danger'>Logout</button>
|
||||
|
||||
@@ -0,0 +1,367 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang='en'>
|
||||
<head>
|
||||
<meta charset='utf-8'>
|
||||
<meta name='viewport' content='width=device-width, initial-scale=1'>
|
||||
<title>Submissions - spark-mcp-go Admin</title>
|
||||
<style>
|
||||
:root {
|
||||
--bg: #0d1117;
|
||||
--panel: #161b22;
|
||||
--border: #30363d;
|
||||
--text: #c9d1d9;
|
||||
--muted: #8b949e;
|
||||
--accent: #58a6ff;
|
||||
--danger: #f85149;
|
||||
--ok: #3fb950;
|
||||
}
|
||||
* { box-sizing: border-box; }
|
||||
body {
|
||||
margin: 0;
|
||||
font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, 'Liberation Mono', monospace;
|
||||
background: var(--bg);
|
||||
color: var(--text);
|
||||
line-height: 1.5;
|
||||
}
|
||||
header {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: space-between;
|
||||
gap: 1rem;
|
||||
padding: 1rem;
|
||||
border-bottom: 1px solid var(--border);
|
||||
background: var(--panel);
|
||||
flex-wrap: wrap;
|
||||
}
|
||||
header h1 { margin: 0; font-size: 1.1rem; color: var(--accent); }
|
||||
nav { display: flex; gap: .75rem; }
|
||||
nav a { color: var(--muted); text-decoration: none; }
|
||||
nav a:hover { color: var(--text); }
|
||||
nav a.active { color: var(--accent); }
|
||||
.auth { display: flex; align-items: center; gap: .5rem; flex-wrap: wrap; }
|
||||
input[type='text'], input[type='password'] {
|
||||
background: var(--bg);
|
||||
color: var(--text);
|
||||
border: 1px solid var(--border);
|
||||
border-radius: 4px;
|
||||
padding: .4rem .5rem;
|
||||
font: inherit;
|
||||
}
|
||||
input[type='text'] { min-width: 260px; }
|
||||
button {
|
||||
background: var(--accent);
|
||||
color: #fff;
|
||||
border: none;
|
||||
border-radius: 4px;
|
||||
padding: .4rem .8rem;
|
||||
font: inherit;
|
||||
cursor: pointer;
|
||||
}
|
||||
button:hover { opacity: .9; }
|
||||
button.danger { background: var(--danger); }
|
||||
button.secondary { background: var(--border); color: var(--text); }
|
||||
main { padding: 1rem; max-width: 1400px; margin: 0 auto; }
|
||||
.toolbar { display: flex; gap: .5rem; margin-bottom: 1rem; align-items: center; flex-wrap: wrap; }
|
||||
.toolbar span { color: var(--muted); }
|
||||
.error {
|
||||
background: rgba(248, 81, 73, .15);
|
||||
border: 1px solid var(--danger);
|
||||
color: var(--danger);
|
||||
padding: .75rem;
|
||||
border-radius: 4px;
|
||||
margin-bottom: 1rem;
|
||||
white-space: pre-wrap;
|
||||
}
|
||||
.hidden { display: none; }
|
||||
table { width: 100%; border-collapse: collapse; }
|
||||
th, td { text-align: left; padding: .5rem; border-bottom: 1px solid var(--border); }
|
||||
th { color: var(--muted); font-weight: 600; user-select: none; }
|
||||
th.sortable { cursor: pointer; }
|
||||
th.sortable:hover { color: var(--text); }
|
||||
td { vertical-align: middle; font-size: .9rem; }
|
||||
td.mono { font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, monospace; }
|
||||
td.wrap { word-break: break-all; }
|
||||
.actions { display: flex; gap: .4rem; }
|
||||
dialog.modal {
|
||||
border: 1px solid var(--border);
|
||||
background: var(--panel);
|
||||
color: var(--text);
|
||||
border-radius: 6px;
|
||||
padding: 1.5rem;
|
||||
min-width: 360px;
|
||||
max-width: 90vw;
|
||||
}
|
||||
dialog.modal::backdrop { background: rgba(0, 0, 0, .65); }
|
||||
dialog.modal h2 { margin: 0 0 .5rem; font-size: 1.1rem; color: var(--accent); }
|
||||
dialog.modal form { display: flex; flex-direction: column; gap: .75rem; }
|
||||
dialog.modal label { color: var(--muted); font-size: .9rem; }
|
||||
dialog.modal input { width: 100%; }
|
||||
.modal-error { color: var(--danger); font-size: .9rem; margin: 0; }
|
||||
.modal-actions { display: flex; gap: .5rem; justify-content: flex-end; }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<header>
|
||||
<div style='display:flex;align-items:center;gap:1rem;'>
|
||||
<h1>spark-mcp-go Admin</h1>
|
||||
<nav>
|
||||
<a href='/admin'>Clusters</a>
|
||||
<a href='/admin/uploads'>Uploads</a>
|
||||
<a href='/admin/submissions' class='active'>Submissions</a>
|
||||
</nav>
|
||||
</div>
|
||||
<div class='auth'>
|
||||
<button id='set-token' class='secondary'>Set Token</button>
|
||||
<button id='logout' class='danger'>Logout</button>
|
||||
</div>
|
||||
</header>
|
||||
|
||||
<main>
|
||||
<div id='error' class='error hidden'></div>
|
||||
|
||||
<section class='toolbar'>
|
||||
<input id='search' type='text' placeholder='Search by file name or app id...'>
|
||||
<button id='refresh' class='secondary'>Refresh</button>
|
||||
<span id='count'></span>
|
||||
</section>
|
||||
|
||||
<section>
|
||||
<table>
|
||||
<thead>
|
||||
<tr>
|
||||
<th class='sortable' data-key='id'>ID</th>
|
||||
<th class='sortable' data-key='file_id'>File ID</th>
|
||||
<th class='sortable' data-key='file_name'>File Name</th>
|
||||
<th class='sortable' data-key='app_id'>App ID</th>
|
||||
<th>Tracking URL</th>
|
||||
<th class='sortable' data-key='cluster_id'>Cluster</th>
|
||||
<th class='sortable' data-key='exit_code'>Exit</th>
|
||||
<th class='sortable' data-key='duration_ms'>Duration (ms)</th>
|
||||
<th class='sortable' data-key='submitted_at'>Submitted At</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody id='submissions-body'>
|
||||
<tr><td colspan='9'>Loading...</td></tr>
|
||||
</tbody>
|
||||
</table>
|
||||
</section>
|
||||
</main>
|
||||
|
||||
<dialog id='token-modal' class='modal'>
|
||||
<h2>Admin Token Required</h2>
|
||||
<p id='modal-error' class='modal-error hidden'></p>
|
||||
<form id='token-form'>
|
||||
<label for='modal-token'>Bearer token</label>
|
||||
<input id='modal-token' type='password' autocomplete='off' required>
|
||||
<div class='modal-actions'>
|
||||
<button type='submit' id='modal-save'>Save</button>
|
||||
<button type='button' id='modal-cancel' class='secondary hidden'>Cancel</button>
|
||||
</div>
|
||||
</form>
|
||||
</dialog>
|
||||
|
||||
<script>
|
||||
const $ = (sel) => document.querySelector(sel);
|
||||
const API = '/admin/submissions';
|
||||
|
||||
let submissions = [];
|
||||
let sortKey = 'submitted_at';
|
||||
let sortDir = -1;
|
||||
let pendingRetry = null;
|
||||
let cancelBlocker = null;
|
||||
|
||||
function headers() {
|
||||
return {
|
||||
'Authorization': 'Bearer ' + (localStorage.getItem('adminToken') || ''),
|
||||
'Content-Type': 'application/json'
|
||||
};
|
||||
}
|
||||
|
||||
function showError(msg) {
|
||||
const el = $('#error');
|
||||
el.textContent = msg || '';
|
||||
el.classList.toggle('hidden', !msg);
|
||||
}
|
||||
|
||||
function showModalError(msg) {
|
||||
const el = $('#modal-error');
|
||||
el.textContent = msg || '';
|
||||
el.classList.toggle('hidden', !msg);
|
||||
}
|
||||
|
||||
function openTokenModal(opts = {}) {
|
||||
const dialog = $('#token-modal');
|
||||
const cancelBtn = $('#modal-cancel');
|
||||
showModalError(opts.error || '');
|
||||
cancelBtn.classList.toggle('hidden', !opts.allowCancel);
|
||||
if (cancelBlocker) {
|
||||
dialog.removeEventListener('cancel', cancelBlocker);
|
||||
cancelBlocker = null;
|
||||
}
|
||||
if (!opts.allowCancel) {
|
||||
cancelBlocker = (e) => e.preventDefault();
|
||||
dialog.addEventListener('cancel', cancelBlocker);
|
||||
}
|
||||
dialog.showModal();
|
||||
$('#modal-token').focus();
|
||||
}
|
||||
|
||||
function closeTokenModal() {
|
||||
const dialog = $('#token-modal');
|
||||
showModalError('');
|
||||
if (cancelBlocker) {
|
||||
dialog.removeEventListener('cancel', cancelBlocker);
|
||||
cancelBlocker = null;
|
||||
}
|
||||
dialog.close();
|
||||
$('#modal-token').value = '';
|
||||
}
|
||||
|
||||
function saveTokenFromModal(e) {
|
||||
e.preventDefault();
|
||||
localStorage.setItem('adminToken', $('#modal-token').value.trim());
|
||||
closeTokenModal();
|
||||
if (pendingRetry) {
|
||||
const fn = pendingRetry;
|
||||
pendingRetry = null;
|
||||
fn();
|
||||
} else {
|
||||
loadSubmissions();
|
||||
}
|
||||
}
|
||||
|
||||
async function api(method, path, body) {
|
||||
const opts = { method, headers: headers() };
|
||||
if (body !== undefined) opts.body = JSON.stringify(body);
|
||||
const res = await fetch(API + path, opts);
|
||||
if (!res.ok) {
|
||||
if (res.status === 401) {
|
||||
localStorage.removeItem('adminToken');
|
||||
pendingRetry = () => api(method, path, body);
|
||||
openTokenModal({
|
||||
error: 'Token rejected by server (401). Enter a valid token.',
|
||||
allowCancel: false
|
||||
});
|
||||
throw new Error('401 Unauthorized');
|
||||
}
|
||||
let detail = res.statusText;
|
||||
try {
|
||||
const j = await res.json();
|
||||
if (j.error) detail = j.error;
|
||||
} catch (_) {}
|
||||
throw new Error(`${res.status} ${detail}`);
|
||||
}
|
||||
if (res.status === 204) return null;
|
||||
return res.json();
|
||||
}
|
||||
|
||||
function formatTime(ts) {
|
||||
if (!ts) return '';
|
||||
const d = new Date(ts);
|
||||
return isNaN(d) ? ts : d.toISOString();
|
||||
}
|
||||
|
||||
function truncate(s, n) {
|
||||
if (!s || s.length <= n) return s || '';
|
||||
return s.slice(0, n) + '…';
|
||||
}
|
||||
|
||||
function escapeHtml(s) {
|
||||
return String(s)
|
||||
.replace(/&/g, '&')
|
||||
.replace(/</g, '<')
|
||||
.replace(/>/g, '>')
|
||||
.replace(/"/g, '"')
|
||||
.replace(/'/g, ''');
|
||||
}
|
||||
|
||||
async function loadSubmissions() {
|
||||
showError('');
|
||||
try {
|
||||
const search = $('#search').value.trim();
|
||||
const path = search ? `/api?search=${encodeURIComponent(search)}` : '/api';
|
||||
submissions = await api('GET', path);
|
||||
renderList();
|
||||
} catch (e) {
|
||||
showError('Failed to load submissions: ' + e.message);
|
||||
}
|
||||
}
|
||||
|
||||
function compare(a, b) {
|
||||
let av = a[sortKey];
|
||||
let bv = b[sortKey];
|
||||
if (sortKey === 'id' || sortKey === 'exit_code' || sortKey === 'duration_ms') {
|
||||
av = Number(av); bv = Number(bv);
|
||||
} else if (sortKey === 'submitted_at') {
|
||||
av = new Date(av).getTime(); bv = new Date(bv).getTime();
|
||||
} else {
|
||||
av = String(av).toLowerCase(); bv = String(bv).toLowerCase();
|
||||
}
|
||||
if (av < bv) return -sortDir;
|
||||
if (av > bv) return sortDir;
|
||||
return 0;
|
||||
}
|
||||
|
||||
function renderList() {
|
||||
const tbody = $('#submissions-body');
|
||||
const list = submissions.slice().sort(compare);
|
||||
$('#count').textContent = list.length + ' submission' + (list.length === 1 ? '' : 's');
|
||||
tbody.innerHTML = '';
|
||||
if (!list.length) {
|
||||
tbody.innerHTML = '<tr><td colspan="9">No submissions.</td></tr>';
|
||||
return;
|
||||
}
|
||||
for (const s of list) {
|
||||
const tr = document.createElement('tr');
|
||||
const fileLink = s.file_id ? `<a href='/admin/uploads/${encodeURIComponent(s.file_id)}'>${escapeHtml(truncate(s.file_id, 8))}</a>` : escapeHtml(truncate(s.file_id, 8));
|
||||
tr.innerHTML = `
|
||||
<td class='mono'>${escapeHtml(s.id)}</td>
|
||||
<td class='mono' title='${escapeHtml(s.file_id)}'>${fileLink}</td>
|
||||
<td>${escapeHtml(s.file_name || '')}</td>
|
||||
<td class='mono'>${escapeHtml(s.app_id)}</td>
|
||||
<td class='wrap'>${escapeHtml(s.tracking_url || '')}</td>
|
||||
<td>${escapeHtml(s.cluster_id)}</td>
|
||||
<td>${escapeHtml(s.exit_code)}</td>
|
||||
<td>${escapeHtml(s.duration_ms)}</td>
|
||||
<td>${escapeHtml(formatTime(s.submitted_at))}</td>
|
||||
`;
|
||||
tbody.appendChild(tr);
|
||||
}
|
||||
}
|
||||
|
||||
function handleSort(e) {
|
||||
const key = e.target.dataset.key;
|
||||
if (!key) return;
|
||||
if (sortKey === key) {
|
||||
sortDir = -sortDir;
|
||||
} else {
|
||||
sortKey = key;
|
||||
sortDir = key === 'submitted_at' ? -1 : 1;
|
||||
}
|
||||
document.querySelectorAll('th.sortable').forEach(th => {
|
||||
th.textContent = th.textContent.replace(/ [▲▼]$/, '');
|
||||
});
|
||||
const marker = sortDir > 0 ? ' ▲' : ' ▼';
|
||||
e.target.textContent = e.target.textContent.replace(/ [▲▼]$/, '') + marker;
|
||||
renderList();
|
||||
}
|
||||
|
||||
$('#search').addEventListener('input', loadSubmissions);
|
||||
$('#refresh').addEventListener('click', loadSubmissions);
|
||||
$('#set-token').addEventListener('click', () => openTokenModal({ allowCancel: true }));
|
||||
$('#logout').addEventListener('click', () => {
|
||||
localStorage.removeItem('adminToken');
|
||||
openTokenModal({ error: 'Token cleared. Enter a new token to continue.', allowCancel: false });
|
||||
});
|
||||
$('#token-form').addEventListener('submit', saveTokenFromModal);
|
||||
$('#modal-cancel').addEventListener('click', closeTokenModal);
|
||||
document.querySelectorAll('th.sortable').forEach(th => th.addEventListener('click', handleSort));
|
||||
|
||||
if (!localStorage.getItem('adminToken')) {
|
||||
openTokenModal({ allowCancel: false });
|
||||
} else {
|
||||
loadSubmissions();
|
||||
}
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -0,0 +1,378 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang='en'>
|
||||
<head>
|
||||
<meta charset='utf-8'>
|
||||
<meta name='viewport' content='width=device-width, initial-scale=1'>
|
||||
<title>Uploads - spark-mcp-go Admin</title>
|
||||
<style>
|
||||
:root {
|
||||
--bg: #0d1117;
|
||||
--panel: #161b22;
|
||||
--border: #30363d;
|
||||
--text: #c9d1d9;
|
||||
--muted: #8b949e;
|
||||
--accent: #58a6ff;
|
||||
--danger: #f85149;
|
||||
--ok: #3fb950;
|
||||
}
|
||||
* { box-sizing: border-box; }
|
||||
body {
|
||||
margin: 0;
|
||||
font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, 'Liberation Mono', monospace;
|
||||
background: var(--bg);
|
||||
color: var(--text);
|
||||
line-height: 1.5;
|
||||
}
|
||||
header {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: space-between;
|
||||
gap: 1rem;
|
||||
padding: 1rem;
|
||||
border-bottom: 1px solid var(--border);
|
||||
background: var(--panel);
|
||||
flex-wrap: wrap;
|
||||
}
|
||||
header h1 { margin: 0; font-size: 1.1rem; color: var(--accent); }
|
||||
nav { display: flex; gap: .75rem; }
|
||||
nav a { color: var(--muted); text-decoration: none; }
|
||||
nav a:hover { color: var(--text); }
|
||||
nav a.active { color: var(--accent); }
|
||||
.auth { display: flex; align-items: center; gap: .5rem; flex-wrap: wrap; }
|
||||
input[type='text'], input[type='password'] {
|
||||
background: var(--bg);
|
||||
color: var(--text);
|
||||
border: 1px solid var(--border);
|
||||
border-radius: 4px;
|
||||
padding: .4rem .5rem;
|
||||
font: inherit;
|
||||
}
|
||||
input[type='text'] { min-width: 220px; }
|
||||
button {
|
||||
background: var(--accent);
|
||||
color: #fff;
|
||||
border: none;
|
||||
border-radius: 4px;
|
||||
padding: .4rem .8rem;
|
||||
font: inherit;
|
||||
cursor: pointer;
|
||||
}
|
||||
button:hover { opacity: .9; }
|
||||
button.danger { background: var(--danger); }
|
||||
button.secondary { background: var(--border); color: var(--text); }
|
||||
main { padding: 1rem; max-width: 1200px; margin: 0 auto; }
|
||||
.toolbar { display: flex; gap: .5rem; margin-bottom: 1rem; align-items: center; flex-wrap: wrap; }
|
||||
.toolbar span { color: var(--muted); }
|
||||
.error {
|
||||
background: rgba(248, 81, 73, .15);
|
||||
border: 1px solid var(--danger);
|
||||
color: var(--danger);
|
||||
padding: .75rem;
|
||||
border-radius: 4px;
|
||||
margin-bottom: 1rem;
|
||||
white-space: pre-wrap;
|
||||
}
|
||||
.hidden { display: none; }
|
||||
table { width: 100%; border-collapse: collapse; }
|
||||
th, td { text-align: left; padding: .5rem; border-bottom: 1px solid var(--border); }
|
||||
th { color: var(--muted); font-weight: 600; user-select: none; }
|
||||
th.sortable { cursor: pointer; }
|
||||
th.sortable:hover { color: var(--text); }
|
||||
td { vertical-align: middle; font-size: .9rem; }
|
||||
td.mono { font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, monospace; }
|
||||
.actions { display: flex; gap: .4rem; }
|
||||
dialog.modal {
|
||||
border: 1px solid var(--border);
|
||||
background: var(--panel);
|
||||
color: var(--text);
|
||||
border-radius: 6px;
|
||||
padding: 1.5rem;
|
||||
min-width: 360px;
|
||||
max-width: 90vw;
|
||||
}
|
||||
dialog.modal::backdrop { background: rgba(0, 0, 0, .65); }
|
||||
dialog.modal h2 { margin: 0 0 .5rem; font-size: 1.1rem; color: var(--accent); }
|
||||
dialog.modal form { display: flex; flex-direction: column; gap: .75rem; }
|
||||
dialog.modal label { color: var(--muted); font-size: .9rem; }
|
||||
dialog.modal input { width: 100%; }
|
||||
.modal-error { color: var(--danger); font-size: .9rem; margin: 0; }
|
||||
.modal-actions { display: flex; gap: .5rem; justify-content: flex-end; }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<header>
|
||||
<div style='display:flex;align-items:center;gap:1rem;'>
|
||||
<h1>spark-mcp-go Admin</h1>
|
||||
<nav>
|
||||
<a href='/admin'>Clusters</a>
|
||||
<a href='/admin/uploads' class='active'>Uploads</a>
|
||||
<a href='/admin/submissions'>Submissions</a>
|
||||
</nav>
|
||||
</div>
|
||||
<div class='auth'>
|
||||
<button id='set-token' class='secondary'>Set Token</button>
|
||||
<button id='logout' class='danger'>Logout</button>
|
||||
</div>
|
||||
</header>
|
||||
|
||||
<main>
|
||||
<div id='error' class='error hidden'></div>
|
||||
|
||||
<section class='toolbar'>
|
||||
<input id='search' type='text' placeholder='Search by name...'>
|
||||
<button id='refresh' class='secondary'>Refresh</button>
|
||||
<span id='count'></span>
|
||||
</section>
|
||||
|
||||
<section>
|
||||
<table>
|
||||
<thead>
|
||||
<tr>
|
||||
<th class='sortable' data-key='file_id'>File ID</th>
|
||||
<th class='sortable' data-key='name'>Name</th>
|
||||
<th class='sortable' data-key='size'>Size (KB)</th>
|
||||
<th class='sortable' data-key='uploaded_at'>Uploaded At</th>
|
||||
<th>SHA256</th>
|
||||
<th>Actions</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody id='uploads-body'>
|
||||
<tr><td colspan='6'>Loading...</td></tr>
|
||||
</tbody>
|
||||
</table>
|
||||
</section>
|
||||
</main>
|
||||
|
||||
<dialog id='token-modal' class='modal'>
|
||||
<h2>Admin Token Required</h2>
|
||||
<p id='modal-error' class='modal-error hidden'></p>
|
||||
<form id='token-form'>
|
||||
<label for='modal-token'>Bearer token</label>
|
||||
<input id='modal-token' type='password' autocomplete='off' required>
|
||||
<div class='modal-actions'>
|
||||
<button type='submit' id='modal-save'>Save</button>
|
||||
<button type='button' id='modal-cancel' class='secondary hidden'>Cancel</button>
|
||||
</div>
|
||||
</form>
|
||||
</dialog>
|
||||
|
||||
<script>
|
||||
const $ = (sel) => document.querySelector(sel);
|
||||
const API = '/admin/uploads';
|
||||
|
||||
let uploads = [];
|
||||
let sortKey = 'uploaded_at';
|
||||
let sortDir = -1;
|
||||
let pendingRetry = null;
|
||||
let cancelBlocker = null;
|
||||
|
||||
function headers() {
|
||||
return {
|
||||
'Authorization': 'Bearer ' + (localStorage.getItem('adminToken') || ''),
|
||||
'Content-Type': 'application/json'
|
||||
};
|
||||
}
|
||||
|
||||
function showError(msg) {
|
||||
const el = $('#error');
|
||||
el.textContent = msg || '';
|
||||
el.classList.toggle('hidden', !msg);
|
||||
}
|
||||
|
||||
function showModalError(msg) {
|
||||
const el = $('#modal-error');
|
||||
el.textContent = msg || '';
|
||||
el.classList.toggle('hidden', !msg);
|
||||
}
|
||||
|
||||
function openTokenModal(opts = {}) {
|
||||
const dialog = $('#token-modal');
|
||||
const cancelBtn = $('#modal-cancel');
|
||||
showModalError(opts.error || '');
|
||||
cancelBtn.classList.toggle('hidden', !opts.allowCancel);
|
||||
if (cancelBlocker) {
|
||||
dialog.removeEventListener('cancel', cancelBlocker);
|
||||
cancelBlocker = null;
|
||||
}
|
||||
if (!opts.allowCancel) {
|
||||
cancelBlocker = (e) => e.preventDefault();
|
||||
dialog.addEventListener('cancel', cancelBlocker);
|
||||
}
|
||||
dialog.showModal();
|
||||
$('#modal-token').focus();
|
||||
}
|
||||
|
||||
function closeTokenModal() {
|
||||
const dialog = $('#token-modal');
|
||||
showModalError('');
|
||||
if (cancelBlocker) {
|
||||
dialog.removeEventListener('cancel', cancelBlocker);
|
||||
cancelBlocker = null;
|
||||
}
|
||||
dialog.close();
|
||||
$('#modal-token').value = '';
|
||||
}
|
||||
|
||||
function saveTokenFromModal(e) {
|
||||
e.preventDefault();
|
||||
localStorage.setItem('adminToken', $('#modal-token').value.trim());
|
||||
closeTokenModal();
|
||||
if (pendingRetry) {
|
||||
const fn = pendingRetry;
|
||||
pendingRetry = null;
|
||||
fn();
|
||||
} else {
|
||||
loadUploads();
|
||||
}
|
||||
}
|
||||
|
||||
async function api(method, path, body) {
|
||||
const opts = { method, headers: headers() };
|
||||
if (body !== undefined) opts.body = JSON.stringify(body);
|
||||
const res = await fetch(API + path, opts);
|
||||
if (!res.ok) {
|
||||
if (res.status === 401) {
|
||||
localStorage.removeItem('adminToken');
|
||||
pendingRetry = () => api(method, path, body);
|
||||
openTokenModal({
|
||||
error: 'Token rejected by server (401). Enter a valid token.',
|
||||
allowCancel: false
|
||||
});
|
||||
throw new Error('401 Unauthorized');
|
||||
}
|
||||
let detail = res.statusText;
|
||||
try {
|
||||
const j = await res.json();
|
||||
if (j.error) detail = j.error;
|
||||
} catch (_) {}
|
||||
throw new Error(`${res.status} ${detail}`);
|
||||
}
|
||||
if (res.status === 204) return null;
|
||||
return res.json();
|
||||
}
|
||||
|
||||
function formatKB(bytes) {
|
||||
return (bytes / 1024).toFixed(2);
|
||||
}
|
||||
|
||||
function formatTime(ts) {
|
||||
if (!ts) return '';
|
||||
const d = new Date(ts);
|
||||
return isNaN(d) ? ts : d.toISOString();
|
||||
}
|
||||
|
||||
function truncate(s, n) {
|
||||
if (!s || s.length <= n) return s || '';
|
||||
return s.slice(0, n) + '…';
|
||||
}
|
||||
|
||||
function escapeHtml(s) {
|
||||
return String(s)
|
||||
.replace(/&/g, '&')
|
||||
.replace(/</g, '<')
|
||||
.replace(/>/g, '>')
|
||||
.replace(/"/g, '"')
|
||||
.replace(/'/g, ''');
|
||||
}
|
||||
|
||||
function filteredUploads() {
|
||||
const q = $('#search').value.trim().toLowerCase();
|
||||
if (!q) return uploads.slice();
|
||||
return uploads.filter(u => (u.name || '').toLowerCase().includes(q));
|
||||
}
|
||||
|
||||
function compare(a, b) {
|
||||
let av = a[sortKey];
|
||||
let bv = b[sortKey];
|
||||
if (sortKey === 'size') { av = Number(av); bv = Number(bv); }
|
||||
else if (sortKey === 'uploaded_at') { av = new Date(av).getTime(); bv = new Date(bv).getTime(); }
|
||||
else { av = String(av).toLowerCase(); bv = String(bv).toLowerCase(); }
|
||||
if (av < bv) return -sortDir;
|
||||
if (av > bv) return sortDir;
|
||||
return 0;
|
||||
}
|
||||
|
||||
function renderList() {
|
||||
const tbody = $('#uploads-body');
|
||||
const list = filteredUploads().sort(compare);
|
||||
$('#count').textContent = list.length + ' file' + (list.length === 1 ? '' : 's');
|
||||
tbody.innerHTML = '';
|
||||
if (!list.length) {
|
||||
tbody.innerHTML = '<tr><td colspan="6">No uploads.</td></tr>';
|
||||
return;
|
||||
}
|
||||
for (const u of list) {
|
||||
const tr = document.createElement('tr');
|
||||
tr.innerHTML = `
|
||||
<td class='mono' title='${escapeHtml(u.file_id)}'>${escapeHtml(truncate(u.file_id, 8))}</td>
|
||||
<td>${escapeHtml(u.name)}</td>
|
||||
<td>${escapeHtml(formatKB(u.size))}</td>
|
||||
<td>${escapeHtml(formatTime(u.uploaded_at))}</td>
|
||||
<td class='mono' title='${escapeHtml(u.sha256)}'>${escapeHtml(truncate(u.sha256, 8))}</td>
|
||||
<td class='actions'>
|
||||
<button data-id='${escapeHtml(u.file_id)}' class='danger delete'>Delete</button>
|
||||
</td>
|
||||
`;
|
||||
tbody.appendChild(tr);
|
||||
}
|
||||
tbody.querySelectorAll('.delete').forEach(b => b.addEventListener('click', handleDelete));
|
||||
}
|
||||
|
||||
async function loadUploads() {
|
||||
showError('');
|
||||
try {
|
||||
uploads = await api('GET', '/api');
|
||||
renderList();
|
||||
} catch (e) {
|
||||
showError('Failed to load uploads: ' + e.message);
|
||||
}
|
||||
}
|
||||
|
||||
async function handleDelete(e) {
|
||||
const id = e.target.dataset.id;
|
||||
if (!confirm(`Delete upload ${id}?`)) return;
|
||||
showError('');
|
||||
try {
|
||||
await api('DELETE', `/${encodeURIComponent(id)}`);
|
||||
await loadUploads();
|
||||
} catch (err) {
|
||||
showError('Failed to delete upload: ' + err.message);
|
||||
}
|
||||
}
|
||||
|
||||
function handleSort(e) {
|
||||
const key = e.target.dataset.key;
|
||||
if (!key) return;
|
||||
if (sortKey === key) {
|
||||
sortDir = -sortDir;
|
||||
} else {
|
||||
sortKey = key;
|
||||
sortDir = key === 'uploaded_at' ? -1 : 1;
|
||||
}
|
||||
document.querySelectorAll('th.sortable').forEach(th => {
|
||||
th.textContent = th.textContent.replace(/ [▲▼]$/, '');
|
||||
});
|
||||
const marker = sortDir > 0 ? ' ▲' : ' ▼';
|
||||
e.target.textContent = e.target.textContent.replace(/ [▲▼]$/, '') + marker;
|
||||
renderList();
|
||||
}
|
||||
|
||||
$('#search').addEventListener('input', renderList);
|
||||
$('#refresh').addEventListener('click', loadUploads);
|
||||
$('#set-token').addEventListener('click', () => openTokenModal({ allowCancel: true }));
|
||||
$('#logout').addEventListener('click', () => {
|
||||
localStorage.removeItem('adminToken');
|
||||
openTokenModal({ error: 'Token cleared. Enter a new token to continue.', allowCancel: false });
|
||||
});
|
||||
$('#token-form').addEventListener('submit', saveTokenFromModal);
|
||||
$('#modal-cancel').addEventListener('click', closeTokenModal);
|
||||
document.querySelectorAll('th.sortable').forEach(th => th.addEventListener('click', handleSort));
|
||||
|
||||
if (!localStorage.getItem('adminToken')) {
|
||||
openTokenModal({ allowCancel: false });
|
||||
} else {
|
||||
loadUploads();
|
||||
}
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -18,6 +18,8 @@ const (
|
||||
ActionClusterCreate Action = "cluster.create"
|
||||
ActionClusterUpdate Action = "cluster.update"
|
||||
ActionClusterDelete Action = "cluster.delete"
|
||||
ActionUploadCreate Action = "upload.create"
|
||||
ActionUploadDelete Action = "upload.delete"
|
||||
)
|
||||
|
||||
// Entry is one admin write operation recorded for accountability.
|
||||
|
||||
+22
-17
@@ -34,23 +34,28 @@ const (
|
||||
// password" — keeps the existing encrypted blob intact.
|
||||
// - URLAllowlist: extra glob patterns beyond the RM/SHS hosts; fetch_url
|
||||
// uses this to permit long-tail SHS endpoints the agent may construct.
|
||||
// - DefaultSubmitArgs: prepended to every spark_submit Tool call's args.
|
||||
// - DefaultSubmitArgs: deprecated; no longer prepended by spark_submit
|
||||
// (post-7920dc9 the builder constructs argv from structured fields).
|
||||
// Kept for backward compatibility with the admin API and DB schema;
|
||||
// actual removal is a follow-up. New cluster configurations should
|
||||
// leave this empty.
|
||||
// - RateLimitPerMin: 0 disables the per-cluster limiter.
|
||||
type Cluster struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
RMURL string `json:"rm_url"`
|
||||
SHSURL string `json:"shs_url"`
|
||||
SparkSubmitExecuteBin string `json:"spark_submit_execute_bin"`
|
||||
IsActive bool `json:"is_active"`
|
||||
AuthType AuthType `json:"auth_type"`
|
||||
AuthUsername string `json:"auth_username,omitempty"`
|
||||
AuthPassword string `json:"-"` // never serialized; encrypted at rest in auth_password_enc
|
||||
SSLVerify bool `json:"ssl_verify"`
|
||||
SSLCABundle string `json:"ssl_ca_bundle,omitempty"`
|
||||
URLAllowlist []string `json:"url_allowlist,omitempty"`
|
||||
DefaultSubmitArgs []string `json:"default_submit_args,omitempty"`
|
||||
RateLimitPerMin int `json:"rate_limit_per_min"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
RMURL string `json:"rm_url"`
|
||||
SHSURL string `json:"shs_url"`
|
||||
SparkSubmitExecuteBin string `json:"spark_submit_execute_bin"`
|
||||
IsActive bool `json:"is_active"`
|
||||
AuthType AuthType `json:"auth_type"`
|
||||
AuthUsername string `json:"auth_username,omitempty"`
|
||||
AuthPassword string `json:"-"` // never serialized; encrypted at rest in auth_password_enc
|
||||
SSLVerify bool `json:"ssl_verify"`
|
||||
SSLCABundle string `json:"ssl_ca_bundle,omitempty"`
|
||||
URLAllowlist []string `json:"url_allowlist,omitempty"`
|
||||
// Deprecated: see godoc.
|
||||
DefaultSubmitArgs []string `json:"default_submit_args,omitempty"`
|
||||
RateLimitPerMin int `json:"rate_limit_per_min"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
@@ -29,6 +29,7 @@ type Config struct {
|
||||
AnalyzerDataSkewRatio float64
|
||||
AnalyzerGCPressureRatio float64
|
||||
AnalyzerBottleneckShuffleGB float64
|
||||
UploadTTL time.Duration
|
||||
}
|
||||
|
||||
// Load reads configuration from environment variables and returns a populated Config.
|
||||
@@ -44,6 +45,7 @@ func Load() (*Config, error) {
|
||||
HTTPClientTimeout: 30 * time.Second,
|
||||
MaxResponseBytes: 1 << 20,
|
||||
SparkSubmitTimeout: 60 * time.Second,
|
||||
UploadTTL: 168 * time.Hour,
|
||||
LogDir: "./data/logs",
|
||||
LogLevel: "info",
|
||||
LogFormat: "text",
|
||||
@@ -80,6 +82,11 @@ func Load() (*Config, error) {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
cfg.UploadTTL, err = parseDuration("UPLOAD_TTL", cfg.UploadTTL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
cfg.LogDir = envString("LOG_DIR", cfg.LogDir)
|
||||
cfg.LogLevel = envString("LOG_LEVEL", cfg.LogLevel)
|
||||
cfg.LogFormat = envString("LOG_FORMAT", cfg.LogFormat)
|
||||
@@ -110,6 +117,10 @@ func Load() (*Config, error) {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if err := validateUploadTTL(cfg.UploadTTL); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
cfg.SQLitePath = filepath.Join(cfg.DataDir, "spark-mcp.db")
|
||||
|
||||
return cfg, nil
|
||||
@@ -212,6 +223,16 @@ func validateGinMode(mode string) error {
|
||||
return fmt.Errorf("config: invalid GIN_MODE %q, want debug/release/test", mode)
|
||||
}
|
||||
|
||||
func validateUploadTTL(d time.Duration) error {
|
||||
if d <= 0 {
|
||||
return fmt.Errorf("config: UPLOAD_TTL must be positive")
|
||||
}
|
||||
if d > 8760*time.Hour {
|
||||
return fmt.Errorf("config: UPLOAD_TTL must be at most 8760h (1 year)")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// loadDotenv reads KEY=VALUE pairs from path and sets them via os.Setenv only
|
||||
// when the variable is not already defined. Missing files are ignored.
|
||||
func loadDotenv(path string) error {
|
||||
@@ -280,7 +301,7 @@ func (c *Config) String() string {
|
||||
|
||||
return fmt.Sprintf(
|
||||
"ListenAddr=%s DataDir=%s SQLitePath=%s AdminTokens=%s AgentToken=%s "+
|
||||
"HTTPClientTimeout=%s MaxResponseBytes=%d SparkSubmitTimeout=%s "+
|
||||
"HTTPClientTimeout=%s MaxResponseBytes=%d SparkSubmitTimeout=%s UploadTTL=%s "+
|
||||
"LogDir=%s LogLevel=%s LogFormat=%s AnalyzerDataSkewRatio=%.1f "+
|
||||
"AnalyzerGCPressureRatio=%.1f AnalyzerBottleneckShuffleGB=%.1f GinMode=%s",
|
||||
c.ListenAddr,
|
||||
@@ -291,6 +312,7 @@ func (c *Config) String() string {
|
||||
c.HTTPClientTimeout,
|
||||
c.MaxResponseBytes,
|
||||
c.SparkSubmitTimeout,
|
||||
c.UploadTTL,
|
||||
c.LogDir,
|
||||
c.LogLevel,
|
||||
c.LogFormat,
|
||||
|
||||
@@ -23,10 +23,13 @@ import (
|
||||
// application application_xxx" line and a few progress lines.
|
||||
const stdoutCap = 10 << 20
|
||||
|
||||
// appIDPattern matches the YARN line printed on successful submission:
|
||||
//
|
||||
// "Submitted application application_12345_0001"
|
||||
var appIDPattern = regexp.MustCompile(`application_\d+_\d+`)
|
||||
// submittedAppPattern matches the explicit success line, including the
|
||||
// application id captured in group 1. This mirrors the Python reference.
|
||||
var submittedAppPattern = regexp.MustCompile(`Submitted application (\S+)`)
|
||||
|
||||
// trackingURLPattern matches "tracking URL: <url>" on stderr. The URL tail
|
||||
// is used as a fallback source for the application id.
|
||||
var trackingURLPattern = regexp.MustCompile(`tracking URL:\s+(\S+)`)
|
||||
|
||||
// SparkSubmitOpts is the resolved command line for a single spark-submit run.
|
||||
//
|
||||
@@ -43,11 +46,12 @@ type SparkSubmitOpts struct {
|
||||
|
||||
// Result is what spark_submit returns to the LLM.
|
||||
type Result struct {
|
||||
AppID string `json:"app_id,omitempty"`
|
||||
ExitCode int `json:"exit_code"`
|
||||
StdoutTail string `json:"stdout_tail"`
|
||||
StderrTail string `json:"stderr_tail"`
|
||||
DurationMS int64 `json:"duration_ms"`
|
||||
AppID string `json:"app_id,omitempty"`
|
||||
TrackingURL string `json:"tracking_url,omitempty"`
|
||||
ExitCode int `json:"exit_code"`
|
||||
StdoutTail string `json:"stdout_tail"`
|
||||
StderrTail string `json:"stderr_tail"`
|
||||
DurationMS int64 `json:"duration_ms"`
|
||||
}
|
||||
|
||||
// ErrBinaryEmpty is returned when the cluster has no spark-submit binary
|
||||
@@ -90,27 +94,70 @@ func Run(ctx context.Context, opts SparkSubmitOpts) (*Result, error) {
|
||||
DurationMS: dur.Milliseconds(),
|
||||
}
|
||||
|
||||
// app_id: prefer stdout, fall back to stderr. Both are searched only on
|
||||
// the tail (cap-bounded) so we don't pay for the full output.
|
||||
if id := matchAppID(res.StdoutTail); id != "" {
|
||||
res.AppID = id
|
||||
} else if id := matchAppID(res.StderrTail); id != "" {
|
||||
res.AppID = id
|
||||
}
|
||||
appID, trackingURL, ok := parseSparkSubmitOutput(res.StdoutTail, res.StderrTail)
|
||||
res.AppID = appID
|
||||
res.TrackingURL = trackingURL
|
||||
|
||||
if !ok {
|
||||
return res, fmt.Errorf("executor: spark-submit ran but no application_id in output (exit_code=%d)", res.ExitCode)
|
||||
}
|
||||
if err != nil {
|
||||
// Wrap the underlying error but keep the parsed result so the LLM
|
||||
// can see exit_code and stderr even on failure.
|
||||
return res, fmt.Errorf("executor: spark-submit failed: %w", err)
|
||||
}
|
||||
return res, nil
|
||||
}
|
||||
|
||||
func matchAppID(s string) string {
|
||||
if s == "" {
|
||||
// parseSparkSubmitOutput extracts the application id and tracking URL from
|
||||
// spark-submit output. It mirrors the Python reference implementation: first
|
||||
// look for "Submitted application <id>", then fall back to "tracking URL:",
|
||||
// deriving the id from the last path segment. The returned ok is false when
|
||||
// no application id could be determined.
|
||||
func parseSparkSubmitOutput(stdout, stderr string) (appID, trackingURL string, ok bool) {
|
||||
// Stage 1: explicit success line, prefer stdout then stderr.
|
||||
var trackingFound bool
|
||||
if m := submittedAppPattern.FindStringSubmatch(stdout); m != nil {
|
||||
appID = m[1]
|
||||
ok = true
|
||||
} else if m := submittedAppPattern.FindStringSubmatch(stderr); m != nil {
|
||||
appID = m[1]
|
||||
ok = true
|
||||
}
|
||||
|
||||
// Stage 2: tracking URL fallback (stderr only, like the Python reference).
|
||||
if !ok {
|
||||
if m := trackingURLPattern.FindStringSubmatch(stderr); m != nil {
|
||||
trackingURL = m[1]
|
||||
trackingFound = true
|
||||
appID = extractAppIDFromTrackingURL(trackingURL)
|
||||
ok = appID != ""
|
||||
if !ok {
|
||||
trackingURL = ""
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Tracking URL is independent: if present anywhere on stderr, capture it.
|
||||
if trackingURL == "" && !trackingFound {
|
||||
if m := trackingURLPattern.FindStringSubmatch(stderr); m != nil {
|
||||
trackingURL = m[1]
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// extractAppIDFromTrackingURL returns the last path segment if it starts with
|
||||
// "application_", otherwise an empty string.
|
||||
func extractAppIDFromTrackingURL(rawURL string) string {
|
||||
trimmed := strings.TrimRight(rawURL, "/")
|
||||
idx := strings.LastIndex(trimmed, "/")
|
||||
if idx < 0 {
|
||||
return ""
|
||||
}
|
||||
return appIDPattern.FindString(s)
|
||||
tail := trimmed[idx+1:]
|
||||
if strings.HasPrefix(tail, "application_") {
|
||||
return tail
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// exitCodeFromErr returns 0 on nil, the process's exit code on *exec.ExitError,
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
package executor
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestParseAppID_FromStdout(t *testing.T) {
|
||||
stdout := "some prefix\nSubmitted application application_123_456\n"
|
||||
stderr := "random log\n"
|
||||
appID, trackingURL, ok := parseSparkSubmitOutput(stdout, stderr)
|
||||
if !ok {
|
||||
t.Fatal("expected ok")
|
||||
}
|
||||
if appID != "application_123_456" {
|
||||
t.Errorf("appID=%q, want application_123_456", appID)
|
||||
}
|
||||
if trackingURL != "" {
|
||||
t.Errorf("trackingURL=%q, want empty", trackingURL)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseAppID_FromStderr(t *testing.T) {
|
||||
stdout := ""
|
||||
stderr := "Submitted application application_123_456\n"
|
||||
appID, trackingURL, ok := parseSparkSubmitOutput(stdout, stderr)
|
||||
if !ok {
|
||||
t.Fatal("expected ok")
|
||||
}
|
||||
if appID != "application_123_456" {
|
||||
t.Errorf("appID=%q, want application_123_456", appID)
|
||||
}
|
||||
if trackingURL != "" {
|
||||
t.Errorf("trackingURL=%q, want empty", trackingURL)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseAppID_FromTrackingURL(t *testing.T) {
|
||||
stdout := ""
|
||||
stderr := "Log line\ntracking URL: http://rm.example.com:8088/proxy/application_123_456/\n"
|
||||
appID, trackingURL, ok := parseSparkSubmitOutput(stdout, stderr)
|
||||
if !ok {
|
||||
t.Fatal("expected ok")
|
||||
}
|
||||
if appID != "application_123_456" {
|
||||
t.Errorf("appID=%q, want application_123_456", appID)
|
||||
}
|
||||
wantURL := "http://rm.example.com:8088/proxy/application_123_456/"
|
||||
if trackingURL != wantURL {
|
||||
t.Errorf("trackingURL=%q, want %q", trackingURL, wantURL)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseAppID_NoAppID(t *testing.T) {
|
||||
stdout := ""
|
||||
stderr := "some random log"
|
||||
appID, trackingURL, ok := parseSparkSubmitOutput(stdout, stderr)
|
||||
if ok {
|
||||
t.Fatal("expected not ok")
|
||||
}
|
||||
if appID != "" {
|
||||
t.Errorf("appID=%q, want empty", appID)
|
||||
}
|
||||
if trackingURL != "" {
|
||||
t.Errorf("trackingURL=%q, want empty", trackingURL)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseAppID_BadTrackingURL(t *testing.T) {
|
||||
stdout := ""
|
||||
stderr := "tracking URL: http://rm.example.com:8088/proxy/not-an-app/\n"
|
||||
appID, trackingURL, ok := parseSparkSubmitOutput(stdout, stderr)
|
||||
if ok {
|
||||
t.Fatal("expected not ok")
|
||||
}
|
||||
if appID != "" {
|
||||
t.Errorf("appID=%q, want empty", appID)
|
||||
}
|
||||
if trackingURL != "" {
|
||||
t.Errorf("trackingURL=%q, want empty", trackingURL)
|
||||
}
|
||||
}
|
||||
@@ -29,8 +29,8 @@ func TestRun_FakeBinary(t *testing.T) {
|
||||
defer cancel()
|
||||
|
||||
res, err := Run(ctx, SparkSubmitOpts{
|
||||
Binary: "/bin/echo",
|
||||
Args: []string{"--master", "yarn; rm -rf /tmp/this-should-not-exist"},
|
||||
Binary: "/bin/echo",
|
||||
Args: []string{"--master", "yarn; rm -rf /tmp/this-should-not-exist", "Submitted application application_12345_0001"},
|
||||
Timeout: 3 * time.Second,
|
||||
})
|
||||
if err != nil {
|
||||
@@ -43,8 +43,8 @@ func TestRun_FakeBinary(t *testing.T) {
|
||||
if !strings.Contains(res.StdoutTail, "yarn; rm -rf /tmp/this-should-not-exist") {
|
||||
t.Errorf("stdout missing literal arg, got: %q", res.StdoutTail)
|
||||
}
|
||||
if res.AppID != "" {
|
||||
t.Errorf("AppID should be empty for echo, got %q", res.AppID)
|
||||
if res.AppID != "application_12345_0001" {
|
||||
t.Errorf("AppID = %q, want application_12345_0001", res.AppID)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -54,8 +54,8 @@ func TestRun_NonZeroExit(t *testing.T) {
|
||||
t.Skip("uses unix binary")
|
||||
}
|
||||
res, err := Run(context.Background(), SparkSubmitOpts{
|
||||
Binary: "/bin/sh",
|
||||
Args: []string{"-c", "exit 7"},
|
||||
Binary: "/bin/sh",
|
||||
Args: []string{"-c", "exit 7"},
|
||||
Timeout: 3 * time.Second,
|
||||
})
|
||||
if err == nil {
|
||||
@@ -66,22 +66,26 @@ func TestRun_NonZeroExit(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestMatchAppID locks down the regex so a YARN output tweak doesn't
|
||||
// TestSubmittedAppPattern locks down the regex so a YARN output tweak doesn't
|
||||
// silently break app_id extraction.
|
||||
func TestMatchAppID(t *testing.T) {
|
||||
func TestSubmittedAppPattern(t *testing.T) {
|
||||
cases := []struct {
|
||||
in, want string
|
||||
}{
|
||||
{"Submitted application application_12345_0001", "application_12345_0001"},
|
||||
{"... some prefix application_999_42 ... tail", "application_999_42"},
|
||||
{"Submitted application application_999_42 extra", "application_999_42"},
|
||||
{"no app id here", ""},
|
||||
{"", ""},
|
||||
{"application_1_2 application_3_4", "application_1_2"}, // first match
|
||||
{"Submitted application app_1_2 Submitted application app_3_4", "app_1_2"}, // first match
|
||||
}
|
||||
for _, tc := range cases {
|
||||
got := matchAppID(tc.in)
|
||||
m := submittedAppPattern.FindStringSubmatch(tc.in)
|
||||
var got string
|
||||
if m != nil {
|
||||
got = m[1]
|
||||
}
|
||||
if got != tc.want {
|
||||
t.Errorf("matchAppID(%q) = %q, want %q", tc.in, got, tc.want)
|
||||
t.Errorf("submittedAppPattern(%q) = %q, want %q", tc.in, got, tc.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -56,6 +56,15 @@ func (c *Client) DoWithRedirect(ctx context.Context, req *http.Request, maxRedir
|
||||
}
|
||||
// **保留** header (Authorization 等)
|
||||
newReq.Header = current.Header.Clone()
|
||||
// TODO(redirect): per RFC 9110 §15.4, 303 See Other must switch the
|
||||
// method to GET and drop the body; 307 Temporary Redirect and
|
||||
// 308 Permanent Redirect must preserve both method AND body. This
|
||||
// implementation always preserves the method and discards the body,
|
||||
// which is correct for 307/308 GET (YARN amContainerLogs) but
|
||||
// non-compliant for 303 and any 307/308 POST. Current YARN / SHS
|
||||
// API surface only triggers 307 on GET, so the gap is unfired.
|
||||
// TestDoWithRedirect_DeferredRFC9110Gaps locks the current behavior
|
||||
// so a future fix must update that test deliberately.
|
||||
newReq.Body = nil
|
||||
current = newReq
|
||||
}
|
||||
|
||||
@@ -140,3 +140,87 @@ func TestDoWithRedirect_ZeroDisallows(t *testing.T) {
|
||||
t.Fatalf("status=%d, want 302", resp.StatusCode)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDoWithRedirect_DeferredRFC9110Gaps documents two known
|
||||
// non-conformances with RFC 9110 §15.4 that we have NOT fixed because
|
||||
// the current upstream YARN / SHS API surface only triggers 307 on
|
||||
// GET amContainerLogs, where neither gap is fired.
|
||||
//
|
||||
// The assertions below verify the CURRENT (gap) behavior. Any future
|
||||
// fix must update these assertions deliberately — a silent flip would
|
||||
// mean the gap fix passed without anyone thinking about 303 method
|
||||
// switch or 307-POST body resend. See the TODO comment in
|
||||
// redirect.go for the full discussion.
|
||||
//
|
||||
// (1) 303 See Other: we keep the original method (RFC requires switch
|
||||
// to GET). The body IS correctly dropped — that's RFC-compliant
|
||||
// even with the wrong method.
|
||||
// (2) 307 Temporary Redirect / 308 Permanent Redirect on a POST: we
|
||||
// keep the method (correct) but drop the body (RFC requires
|
||||
// resend with the same method).
|
||||
func TestDoWithRedirect_DeferredRFC9110Gaps(t *testing.T) {
|
||||
t.Run("303 keeps method, should switch to GET", func(t *testing.T) {
|
||||
var hitMethod, hitBody string
|
||||
target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
hitMethod = r.Method
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
hitBody = string(b)
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
defer target.Close()
|
||||
|
||||
origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Location", target.URL+"/see-other")
|
||||
w.WriteHeader(http.StatusSeeOther)
|
||||
}))
|
||||
defer origin.Close()
|
||||
|
||||
c := New(Config{Timeout: 5 * time.Second, MaxResponseBytes: 1024})
|
||||
req, _ := http.NewRequest(http.MethodPost, origin.URL+"/start", strings.NewReader("payload-303"))
|
||||
if _, err := c.DoWithRedirect(context.Background(), req, 5, "127.0.0.1"); err != nil {
|
||||
t.Fatalf("DoWithRedirect: %v", err)
|
||||
}
|
||||
|
||||
// CURRENT (gap) behavior. RFC-correct behavior: hitMethod == "GET".
|
||||
if hitMethod != http.MethodPost {
|
||||
t.Errorf("303 arrived as %s, want POST (current gap; RFC says GET)", hitMethod)
|
||||
}
|
||||
// Body must be empty — that part is RFC-compliant (303 always
|
||||
// drops body, even when the method is wrong).
|
||||
if hitBody != "" {
|
||||
t.Errorf("303 body=%q, want empty (RFC requires drop regardless of method)", hitBody)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("307 on POST drops body, should resend", func(t *testing.T) {
|
||||
var hitMethod, hitBody string
|
||||
target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
hitMethod = r.Method
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
hitBody = string(b)
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
defer target.Close()
|
||||
|
||||
origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Location", target.URL+"/temp")
|
||||
w.WriteHeader(http.StatusTemporaryRedirect)
|
||||
}))
|
||||
defer origin.Close()
|
||||
|
||||
c := New(Config{Timeout: 5 * time.Second, MaxResponseBytes: 1024})
|
||||
req, _ := http.NewRequest(http.MethodPost, origin.URL+"/start", strings.NewReader("payload-307"))
|
||||
if _, err := c.DoWithRedirect(context.Background(), req, 5, "127.0.0.1"); err != nil {
|
||||
t.Fatalf("DoWithRedirect: %v", err)
|
||||
}
|
||||
|
||||
// Method preserved (correct).
|
||||
if hitMethod != http.MethodPost {
|
||||
t.Errorf("307 arrived as %s, want POST", hitMethod)
|
||||
}
|
||||
// CURRENT (gap) behavior. RFC-correct behavior: hitBody == "payload-307".
|
||||
if hitBody != "" {
|
||||
t.Errorf("307 body=%q, want empty (current gap; RFC says resend payload-307)", hitBody)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
@@ -116,12 +116,8 @@ func matchAllowedHost(host, pattern string) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
if strings.HasPrefix(pattern, "*.") {
|
||||
pattern = pattern[2:]
|
||||
}
|
||||
if strings.HasPrefix(pattern, ".") {
|
||||
pattern = pattern[1:]
|
||||
}
|
||||
pattern = strings.TrimPrefix(pattern, "*.")
|
||||
pattern = strings.TrimPrefix(pattern, ".")
|
||||
|
||||
return host == pattern || strings.HasSuffix(host, "."+pattern)
|
||||
}
|
||||
|
||||
@@ -5,18 +5,23 @@ import (
|
||||
"time"
|
||||
|
||||
"spark-mcp-go/internal/analyzer"
|
||||
"spark-mcp-go/internal/audit"
|
||||
"spark-mcp-go/internal/httpclient"
|
||||
"spark-mcp-go/internal/storage"
|
||||
"spark-mcp-go/internal/uploads"
|
||||
)
|
||||
|
||||
// Deps bundles the dependencies shared by all MCP Tool handlers.
|
||||
type Deps struct {
|
||||
Logger *slog.Logger
|
||||
AuditRepo *audit.Repo
|
||||
ClusterRepo *storage.ClusterRepo
|
||||
SubmissionRepo *storage.SubmissionRepo
|
||||
SparkSubmitTimeout time.Duration
|
||||
HTTPClient *httpclient.Client
|
||||
MaxResponseBytes int64
|
||||
DataDir string // upload_file writes to DataDir/uploads
|
||||
UploadStore uploads.Store
|
||||
|
||||
AnalyzerThresholds analyzer.Thresholds
|
||||
}
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
@@ -15,14 +16,16 @@ import (
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
|
||||
"spark-mcp-go/internal/audit"
|
||||
"spark-mcp-go/internal/cluster"
|
||||
"spark-mcp-go/internal/httpclient"
|
||||
"spark-mcp-go/internal/storage"
|
||||
"spark-mcp-go/internal/uploads"
|
||||
)
|
||||
|
||||
// testDepsWithDataDir returns dependencies backed by an in-memory DB and a
|
||||
// temporary data directory.
|
||||
func testDepsWithDataDir(t *testing.T) (*Deps, *storage.ClusterRepo) {
|
||||
func testDepsWithDataDir(t *testing.T) (*Deps, *storage.ClusterRepo, *storage.SubmissionRepo, *audit.Repo) {
|
||||
t.Helper()
|
||||
db, err := storage.Open(":memory:")
|
||||
if err != nil {
|
||||
@@ -30,15 +33,24 @@ func testDepsWithDataDir(t *testing.T) (*Deps, *storage.ClusterRepo) {
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
|
||||
dataDir := t.TempDir()
|
||||
uploadStore, err := uploads.New(filepath.Join(dataDir, "uploads"))
|
||||
if err != nil {
|
||||
t.Fatalf("new upload store: %v", err)
|
||||
}
|
||||
|
||||
return &Deps{
|
||||
HTTPClient: httpclient.New(httpclient.Config{
|
||||
Timeout: 5 * time.Second,
|
||||
MaxResponseBytes: 1 << 20,
|
||||
}),
|
||||
AuditRepo: audit.NewRepo(db),
|
||||
ClusterRepo: db.Clusters(),
|
||||
SubmissionRepo: db.Submissions(),
|
||||
MaxResponseBytes: 1 << 20,
|
||||
DataDir: t.TempDir(),
|
||||
}, db.Clusters()
|
||||
DataDir: dataDir,
|
||||
UploadStore: uploadStore,
|
||||
}, db.Clusters(), db.Submissions(), audit.NewRepo(db)
|
||||
}
|
||||
|
||||
// createCluster creates a cluster in the repository with the given fields.
|
||||
@@ -89,7 +101,7 @@ func TestFetchURL(t *testing.T) {
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
deps, repo := testDepsWithDataDir(t)
|
||||
deps, repo, _, _ := testDepsWithDataDir(t)
|
||||
createCluster(t, repo, &cluster.Cluster{
|
||||
ID: "cluster-a",
|
||||
Name: "Cluster A",
|
||||
@@ -283,7 +295,7 @@ func TestFetchURL_RedirectPreservesAuth(t *testing.T) {
|
||||
}))
|
||||
defer server1.Close()
|
||||
|
||||
deps, repo := testDepsWithDataDir(t)
|
||||
deps, repo, _, _ := testDepsWithDataDir(t)
|
||||
createCluster(t, repo, &cluster.Cluster{
|
||||
ID: "cluster-redirect",
|
||||
Name: "Redirect",
|
||||
@@ -312,7 +324,7 @@ func TestFetchURL_RedirectPreservesAuth(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestUploadFile(t *testing.T) {
|
||||
deps, _ := testDepsWithDataDir(t)
|
||||
deps, _, _, auditRepo := testDepsWithDataDir(t)
|
||||
|
||||
req := newToolRequest(UploadFileName, map[string]any{
|
||||
"filename": "hello.txt",
|
||||
@@ -324,15 +336,26 @@ func TestUploadFile(t *testing.T) {
|
||||
}
|
||||
|
||||
payload := resultJSON(t, res)
|
||||
if payload["path"] != "uploads/hello.txt" {
|
||||
t.Errorf("path=%v, want uploads/hello.txt", payload["path"])
|
||||
fileID, ok := payload["file_id"].(string)
|
||||
if !ok || !regexp.MustCompile(`^[0-9a-f]{32}$`).MatchString(fileID) {
|
||||
t.Errorf("file_id=%v, want 32 lowercase hex chars", payload["file_id"])
|
||||
}
|
||||
path, ok := payload["path"].(string)
|
||||
if !ok || !filepath.IsAbs(path) || !strings.HasSuffix(path, "/"+fileID) {
|
||||
t.Errorf("path=%v, want absolute path ending in /%s", payload["path"], fileID)
|
||||
}
|
||||
if payload["name"] != "hello.txt" {
|
||||
t.Errorf("name=%v, want hello.txt", payload["name"])
|
||||
}
|
||||
if payload["size"] != float64(len("hello world")) {
|
||||
t.Errorf("size=%v, want %d", payload["size"], len("hello world"))
|
||||
}
|
||||
sha, ok := payload["sha256"].(string)
|
||||
if !ok || !regexp.MustCompile(`^[0-9a-f]{64}$`).MatchString(sha) {
|
||||
t.Errorf("sha256=%v, want 64-char hex", payload["sha256"])
|
||||
}
|
||||
|
||||
final := filepath.Join(deps.DataDir, "uploads", "hello.txt")
|
||||
got, err := os.ReadFile(final)
|
||||
got, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatalf("read uploaded file: %v", err)
|
||||
}
|
||||
@@ -340,17 +363,32 @@ func TestUploadFile(t *testing.T) {
|
||||
t.Errorf("content=%q, want %q", got, "hello world")
|
||||
}
|
||||
|
||||
info, err := os.Stat(final)
|
||||
info, err := os.Stat(path)
|
||||
if err != nil {
|
||||
t.Fatalf("stat uploaded file: %v", err)
|
||||
}
|
||||
if 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) {
|
||||
deps, _ := testDepsWithDataDir(t)
|
||||
deps, _, _, _ := testDepsWithDataDir(t)
|
||||
|
||||
cases := []string{
|
||||
"../../../etc/passwd",
|
||||
@@ -381,7 +419,7 @@ func TestUploadFile_PathTraversal(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestUploadFile_Base64(t *testing.T) {
|
||||
deps, _ := testDepsWithDataDir(t)
|
||||
deps, _, _, _ := testDepsWithDataDir(t)
|
||||
|
||||
req := newToolRequest(UploadFileName, map[string]any{
|
||||
"filename": "hello.bin",
|
||||
@@ -393,9 +431,19 @@ func TestUploadFile_Base64(t *testing.T) {
|
||||
t.Fatalf("handler error: %v", err)
|
||||
}
|
||||
|
||||
_ = resultJSON(t, res)
|
||||
final := filepath.Join(deps.DataDir, "uploads", "hello.bin")
|
||||
got, err := os.ReadFile(final)
|
||||
payload := resultJSON(t, res)
|
||||
path, ok := payload["path"].(string)
|
||||
if !ok || !filepath.IsAbs(path) {
|
||||
t.Errorf("path=%v, want absolute path", payload["path"])
|
||||
}
|
||||
if payload["name"] != "hello.bin" {
|
||||
t.Errorf("name=%v, want hello.bin", payload["name"])
|
||||
}
|
||||
if payload["size"] != float64(len("hello")) {
|
||||
t.Errorf("size=%v, want %d", payload["size"], len("hello"))
|
||||
}
|
||||
|
||||
got, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatalf("read uploaded file: %v", err)
|
||||
}
|
||||
|
||||
@@ -209,12 +209,8 @@ func hostAllowed(host string, allowed []string) bool {
|
||||
if pattern == host {
|
||||
return true
|
||||
}
|
||||
if strings.HasPrefix(pattern, "*.") {
|
||||
pattern = pattern[2:]
|
||||
}
|
||||
if strings.HasPrefix(pattern, ".") {
|
||||
pattern = pattern[1:]
|
||||
}
|
||||
pattern = strings.TrimPrefix(pattern, "*.")
|
||||
pattern = strings.TrimPrefix(pattern, ".")
|
||||
if host == pattern || strings.HasSuffix(host, "."+pattern) {
|
||||
return true
|
||||
}
|
||||
|
||||
@@ -3,10 +3,15 @@ package tools
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"sort"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
|
||||
"spark-mcp-go/internal/executor"
|
||||
"spark-mcp-go/internal/storage"
|
||||
)
|
||||
|
||||
const SparkSubmitName = "spark_submit"
|
||||
@@ -14,45 +19,131 @@ const SparkSubmitName = "spark_submit"
|
||||
// NewSparkSubmitTool returns the schema for the spark_submit MCP Tool.
|
||||
func NewSparkSubmitTool() mcp.Tool {
|
||||
return mcp.NewTool(SparkSubmitName,
|
||||
mcp.WithDescription("Submit a Spark application to a configured cluster using the cluster's local spark-submit binary. Returns {app_id, exit_code, stdout_tail, stderr_tail, duration_ms}. The cluster's default_submit_args are prepended to your args. Binary is invoked as a child process — no shell, no command injection."),
|
||||
mcp.WithDescription("Submit a Spark application to a configured cluster. The command is built in this order: --master, --queue, --executor-memory, --executor-cores, --num-executors, then --conf per spark_conf entry, then --flag value per extra_args entry, then script_path last. For script_path, pass the `path` field returned by upload_file (do not construct your own path). Returns {app_id, exit_code, stdout_tail, stderr_tail, duration_ms}. Binary is invoked as a child process — no shell, no command injection."),
|
||||
mcp.WithString("cluster_id",
|
||||
mcp.Required(),
|
||||
mcp.Description("ID of the configured cluster (from list_clusters)"),
|
||||
),
|
||||
// mcp-go v0.56.0 uses PropertyOption for array item schema.
|
||||
mcp.WithArray("args",
|
||||
mcp.Description("Spark-submit CLI args (after default_submit_args). Each element becomes one argv entry — no shell interpretation."),
|
||||
mcp.Items(map[string]any{"type": "string"}),
|
||||
mcp.WithString("master",
|
||||
mcp.Required(),
|
||||
mcp.Description("Spark master URL, e.g. yarn, k8s://https://..."),
|
||||
),
|
||||
mcp.WithString("deploy_mode",
|
||||
mcp.Required(),
|
||||
mcp.Description("Required for forward compatibility; not currently emitted in the command by the builder (matches Python reference)."),
|
||||
mcp.Enum("client", "cluster"),
|
||||
),
|
||||
mcp.WithString("script_path",
|
||||
mcp.Required(),
|
||||
mcp.Description("Absolute path to the script. Pass the `path` field returned by upload_file — do not construct your own path. The script is always the last argv element. Paths not minted by upload_file on this server are rejected."),
|
||||
),
|
||||
mcp.WithString("queue",
|
||||
mcp.Required(),
|
||||
mcp.Description("YARN queue name"),
|
||||
),
|
||||
mcp.WithString("executor_memory",
|
||||
mcp.Required(),
|
||||
mcp.Description("e.g. 4G"),
|
||||
),
|
||||
mcp.WithNumber("executor_cores",
|
||||
mcp.Required(),
|
||||
mcp.Description("Cores per executor (integer)"),
|
||||
),
|
||||
mcp.WithNumber("num_executors",
|
||||
mcp.Required(),
|
||||
mcp.Description("Static executor count (integer)"),
|
||||
),
|
||||
mcp.WithObject("spark_conf",
|
||||
mcp.Description("Map of key=value pairs, each emitted as --conf key=value"),
|
||||
mcp.AdditionalProperties(map[string]any{"type": "string"}),
|
||||
),
|
||||
mcp.WithObject("extra_args",
|
||||
mcp.Description("Map of flag -> value, each emitted as --flag value"),
|
||||
mcp.AdditionalProperties(map[string]any{"type": "string"}),
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
// SparkSubmitHandler runs spark-submit against the requested cluster.
|
||||
func (d *Deps) SparkSubmitHandler(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
args := req.GetArguments()
|
||||
|
||||
clusterID, err := req.RequireString("cluster_id")
|
||||
if err != nil {
|
||||
return errResult("spark_submit: " + err.Error()), nil
|
||||
}
|
||||
master, err := requireNonEmptyString(args, "master")
|
||||
if err != nil {
|
||||
return errResult("spark_submit: " + err.Error()), nil
|
||||
}
|
||||
deployMode, err := req.RequireString("deploy_mode")
|
||||
if err != nil {
|
||||
return errResult("spark_submit: " + err.Error()), nil
|
||||
}
|
||||
scriptPath, err := req.RequireString("script_path")
|
||||
if err != nil {
|
||||
return errResult("spark_submit: " + err.Error()), nil
|
||||
}
|
||||
|
||||
args := req.GetArguments()
|
||||
var userArgs []string
|
||||
if rawArgs, ok := args["args"]; ok && rawArgs != nil {
|
||||
arr, ok := rawArgs.([]any)
|
||||
if !ok {
|
||||
return errResult("spark_submit: args is not an array"), nil
|
||||
}
|
||||
for i, item := range arr {
|
||||
s, ok := item.(string)
|
||||
if !ok {
|
||||
return errResult(fmt.Sprintf("spark_submit: args[%d] is not a string", i)), nil
|
||||
}
|
||||
userArgs = append(userArgs, s)
|
||||
queue, err := requireNonEmptyString(args, "queue")
|
||||
if err != nil {
|
||||
return errResult("spark_submit: " + err.Error()), nil
|
||||
}
|
||||
executorMemory, err := requireNonEmptyString(args, "executor_memory")
|
||||
if err != nil {
|
||||
return errResult("spark_submit: " + err.Error()), nil
|
||||
}
|
||||
|
||||
// Reject paths that didn't come through upload_file on this server.
|
||||
// Cluster-local paths and arbitrary host paths are unreachable from the MCP
|
||||
// server's spark-submit, so we fail fast with a clear message rather than
|
||||
// letting spark-submit produce a confusing FileNotFoundException later.
|
||||
// Placed after all required-string parses so error messages surface in
|
||||
// field order (cluster_id, master, ..., script_path, queue, ...).
|
||||
fileID, err := d.UploadStore.Validate(scriptPath)
|
||||
if err != nil {
|
||||
if d.Logger != nil {
|
||||
d.Logger.Warn("spark_submit.unminted_path_rejected", slog.String("path", scriptPath), slog.String("err", err.Error()))
|
||||
}
|
||||
return errResult(fmt.Sprintf("spark_submit: script_path %q is not from upload_file on this server; upload the file first via upload_file and pass back the path field", scriptPath)), nil
|
||||
}
|
||||
|
||||
// deploy_mode is required for forward compatibility but the builder does not
|
||||
// emit --deploy-mode, matching the Python reference's actual behavior. If the
|
||||
// Go side needs --deploy-mode later, add it here and document the divergence.
|
||||
if deployMode != "client" && deployMode != "cluster" {
|
||||
return errResult("spark_submit: deploy_mode must be client or cluster"), nil
|
||||
}
|
||||
|
||||
sparkConf, err := parseStringMap(args, "spark_conf")
|
||||
if err != nil {
|
||||
return errResult("spark_submit: " + err.Error()), nil
|
||||
}
|
||||
extraArgs, err := parseStringMap(args, "extra_args")
|
||||
if err != nil {
|
||||
return errResult("spark_submit: " + err.Error()), nil
|
||||
}
|
||||
|
||||
executorCores, err := requireNonNegativeInt(args, "executor_cores")
|
||||
if err != nil {
|
||||
return errResult("spark_submit: " + err.Error()), nil
|
||||
}
|
||||
numExecutors, err := requireNonNegativeInt(args, "num_executors")
|
||||
if err != nil {
|
||||
return errResult("spark_submit: " + err.Error()), nil
|
||||
}
|
||||
|
||||
callLog := startToolCall(ctx, d.Logger, SparkSubmitName, map[string]any{
|
||||
"cluster_id": clusterID,
|
||||
"args": userArgs,
|
||||
"cluster_id": clusterID,
|
||||
"master": master,
|
||||
"deploy_mode": deployMode,
|
||||
"script_path": scriptPath,
|
||||
"queue": queue,
|
||||
"executor_memory": executorMemory,
|
||||
"executor_cores": executorCores,
|
||||
"num_executors": numExecutors,
|
||||
"spark_conf": sparkConf,
|
||||
"extra_args": extraArgs,
|
||||
})
|
||||
defer callLog.End()
|
||||
|
||||
@@ -62,39 +153,182 @@ func (d *Deps) SparkSubmitHandler(ctx context.Context, req mcp.CallToolRequest)
|
||||
return errResult("spark_submit: cluster " + clusterID + ": " + err.Error()), nil
|
||||
}
|
||||
|
||||
fullArgs := append([]string{}, cl.DefaultSubmitArgs...)
|
||||
fullArgs = append(fullArgs, userArgs...)
|
||||
binary := cl.SparkSubmitExecuteBin
|
||||
cmd := buildSparkSubmitCommand(SparkSubmitCommandOpts{
|
||||
Master: master,
|
||||
Queue: queue,
|
||||
ExecutorMemory: executorMemory,
|
||||
ExecutorCores: executorCores,
|
||||
NumExecutors: numExecutors,
|
||||
SparkConf: sparkConf,
|
||||
ExtraArgs: extraArgs,
|
||||
ScriptPath: scriptPath,
|
||||
})
|
||||
|
||||
if d.Logger != nil {
|
||||
d.Logger.Debug("spark_submit.built", "cluster_id", clusterID, "argv", cmd)
|
||||
}
|
||||
|
||||
result, err := executor.Run(ctx, executor.SparkSubmitOpts{
|
||||
Binary: cl.SparkSubmitExecuteBin,
|
||||
Args: fullArgs,
|
||||
Binary: binary,
|
||||
Args: cmd,
|
||||
Timeout: d.SparkSubmitTimeout,
|
||||
})
|
||||
if err != nil {
|
||||
callLog.WithError(err)
|
||||
// Even on error, result carries ExitCode/Stderr for the LLM.
|
||||
if result != nil {
|
||||
callLog.WithResult(map[string]any{
|
||||
"app_id": result.AppID,
|
||||
"exit_code": result.ExitCode,
|
||||
"duration_ms": result.DurationMS,
|
||||
})
|
||||
callLog.WithResult(sparkSubmitResultLog(binary, cmd, result))
|
||||
return textResult(encodeJSON(map[string]any{
|
||||
"app_id": result.AppID,
|
||||
"exit_code": result.ExitCode,
|
||||
"stdout_tail": result.StdoutTail,
|
||||
"stderr_tail": result.StderrTail,
|
||||
"duration_ms": result.DurationMS,
|
||||
"error": err.Error(),
|
||||
"app_id": result.AppID,
|
||||
"tracking_url": result.TrackingURL,
|
||||
"exit_code": result.ExitCode,
|
||||
"stdout_tail": result.StdoutTail,
|
||||
"stderr_tail": result.StderrTail,
|
||||
"duration_ms": result.DurationMS,
|
||||
"error": err.Error(),
|
||||
})), nil
|
||||
}
|
||||
return errResult("spark_submit: " + err.Error()), nil
|
||||
}
|
||||
|
||||
callLog.WithResult(map[string]any{
|
||||
if result.AppID != "" && d.SubmissionRepo != nil {
|
||||
sub := storage.Submission{
|
||||
FileID: fileID,
|
||||
AppID: result.AppID,
|
||||
TrackingURL: result.TrackingURL,
|
||||
ClusterID: clusterID,
|
||||
ExitCode: result.ExitCode,
|
||||
DurationMS: result.DurationMS,
|
||||
SubmittedAt: time.Now(),
|
||||
}
|
||||
if createErr := d.SubmissionRepo.Create(ctx, sub); createErr != nil {
|
||||
if d.Logger != nil {
|
||||
d.Logger.Warn("spark_submit.persist_failed", slog.String("file_id", fileID), slog.String("app_id", result.AppID), slog.String("err", createErr.Error()))
|
||||
} else {
|
||||
slog.Default().Warn("spark_submit.persist_failed", slog.String("file_id", fileID), slog.String("app_id", result.AppID), slog.String("err", createErr.Error()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
callLog.WithResult(sparkSubmitResultLog(binary, cmd, result))
|
||||
return textResult(encodeJSON(result)), nil
|
||||
}
|
||||
|
||||
// sparkSubmitResultLog builds the standard result map used for both logging
|
||||
// and the admin audit trail.
|
||||
func sparkSubmitResultLog(binary string, cmd []string, result *executor.Result) map[string]any {
|
||||
m := map[string]any{
|
||||
"binary": binary,
|
||||
"argv": cmd,
|
||||
"app_id": result.AppID,
|
||||
"exit_code": result.ExitCode,
|
||||
"duration_ms": result.DurationMS,
|
||||
})
|
||||
return textResult(encodeJSON(result)), nil
|
||||
}
|
||||
if result.TrackingURL != "" {
|
||||
m["tracking_url"] = result.TrackingURL
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
// SparkSubmitCommandOpts holds the structured arguments used to build the
|
||||
// spark-submit argv. The returned argv does NOT include the binary itself.
|
||||
type SparkSubmitCommandOpts struct {
|
||||
Master string
|
||||
Queue string
|
||||
ExecutorMemory string
|
||||
ExecutorCores int
|
||||
NumExecutors int
|
||||
SparkConf map[string]string
|
||||
ExtraArgs map[string]string
|
||||
ScriptPath string
|
||||
}
|
||||
|
||||
// buildSparkSubmitCommand builds the post-binary argv for spark-submit.
|
||||
// DefaultSubmitArgs is no longer prepended; the field is currently unused at
|
||||
// runtime but kept for backward compatibility with the admin API. Removal is
|
||||
// a separate follow-up.
|
||||
func buildSparkSubmitCommand(opts SparkSubmitCommandOpts) []string {
|
||||
cmd := []string{"--master", opts.Master}
|
||||
cmd = append(cmd, "--queue", opts.Queue)
|
||||
cmd = append(cmd, "--executor-memory", opts.ExecutorMemory)
|
||||
cmd = append(cmd, "--executor-cores", strconv.Itoa(opts.ExecutorCores))
|
||||
cmd = append(cmd, "--num-executors", strconv.Itoa(opts.NumExecutors))
|
||||
|
||||
for _, k := range sortedStringKeys(opts.SparkConf) {
|
||||
cmd = append(cmd, "--conf", k+"="+opts.SparkConf[k])
|
||||
}
|
||||
for _, k := range sortedStringKeys(opts.ExtraArgs) {
|
||||
cmd = append(cmd, "--"+k, opts.ExtraArgs[k])
|
||||
}
|
||||
cmd = append(cmd, opts.ScriptPath)
|
||||
return cmd
|
||||
}
|
||||
|
||||
// requireNonNegativeInt extracts an integer parameter from the request and
|
||||
// rejects negative or non-integer values.
|
||||
func requireNonNegativeInt(args map[string]any, key string) (int, error) {
|
||||
v, ok := args[key]
|
||||
if !ok || v == nil {
|
||||
return 0, fmt.Errorf("%s is required", key)
|
||||
}
|
||||
switch n := v.(type) {
|
||||
case int:
|
||||
if n < 0 {
|
||||
return 0, fmt.Errorf("%s must be a non-negative integer", key)
|
||||
}
|
||||
return n, nil
|
||||
case float64:
|
||||
ni := int(n)
|
||||
if n < 0 || float64(ni) != n {
|
||||
return 0, fmt.Errorf("%s must be a non-negative integer", key)
|
||||
}
|
||||
return ni, nil
|
||||
default:
|
||||
return 0, fmt.Errorf("%s must be a number", key)
|
||||
}
|
||||
}
|
||||
|
||||
// parseStringMap extracts an optional object whose values are strings.
|
||||
func parseStringMap(args map[string]any, key string) (map[string]string, error) {
|
||||
out := make(map[string]string)
|
||||
v, ok := args[key]
|
||||
if !ok || v == nil {
|
||||
return out, nil
|
||||
}
|
||||
m, ok := v.(map[string]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("%s must be an object", key)
|
||||
}
|
||||
for k, val := range m {
|
||||
s, ok := val.(string)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("%s[%q] must be a string", key, k)
|
||||
}
|
||||
out[k] = s
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// sortedStringKeys returns the keys of m in deterministic order.
|
||||
func sortedStringKeys(m map[string]string) []string {
|
||||
keys := make([]string, 0, len(m))
|
||||
for k := range m {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
return keys
|
||||
}
|
||||
|
||||
// requireNonEmptyString extracts a required string parameter and rejects empty
|
||||
// values. MCP's RequireString already enforces presence/type; this enforces
|
||||
// that the field is not blank.
|
||||
func requireNonEmptyString(args map[string]any, key string) (string, error) {
|
||||
s, ok := args[key].(string)
|
||||
if !ok {
|
||||
return "", fmt.Errorf("%s is required", key)
|
||||
}
|
||||
if s == "" {
|
||||
return "", fmt.Errorf("%s is required", key)
|
||||
}
|
||||
return s, nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,337 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
|
||||
"spark-mcp-go/internal/audit"
|
||||
"spark-mcp-go/internal/cluster"
|
||||
"spark-mcp-go/internal/httpclient"
|
||||
"spark-mcp-go/internal/storage"
|
||||
"spark-mcp-go/internal/uploads"
|
||||
)
|
||||
|
||||
func TestSparkSubmit_StructuredCommand(t *testing.T) {
|
||||
deps, repo, _, _ := testDepsWithDataDir(t)
|
||||
deps.SparkSubmitTimeout = 5 * time.Second
|
||||
store := deps.UploadStore
|
||||
|
||||
fileID, _, _, _, absPath, err := store.Save([]byte("print('hi')\n"), "hello.py")
|
||||
if err != nil {
|
||||
t.Fatalf("save upload: %v", err)
|
||||
}
|
||||
|
||||
echoed := filepath.Join(t.TempDir(), "echoed")
|
||||
echoScript := filepath.Join(t.TempDir(), "echo-args.sh")
|
||||
if err := os.WriteFile(echoScript, []byte("#!/bin/sh\nfor arg do\n echo \"$arg\" >> \"$ECHO_FILE\"\ndone\necho 'Submitted application application_test_0001'\n"), 0o755); err != nil {
|
||||
t.Fatalf("write echo script: %v", err)
|
||||
}
|
||||
|
||||
t.Setenv("ECHO_FILE", echoed)
|
||||
|
||||
createCluster(t, repo, &cluster.Cluster{
|
||||
ID: "cluster-echo",
|
||||
Name: "Echo",
|
||||
IsActive: true,
|
||||
AuthType: cluster.AuthNone,
|
||||
SparkSubmitExecuteBin: echoScript,
|
||||
})
|
||||
|
||||
req := newToolRequest(SparkSubmitName, map[string]any{
|
||||
"cluster_id": "cluster-echo",
|
||||
"master": "yarn",
|
||||
"deploy_mode": "cluster",
|
||||
"script_path": absPath,
|
||||
"queue": "default",
|
||||
"executor_memory": "2G",
|
||||
"executor_cores": 1,
|
||||
"num_executors": 4,
|
||||
"spark_conf": map[string]any{"a": "1", "b": "2"},
|
||||
"extra_args": map[string]any{"name": "wordcount-job"},
|
||||
})
|
||||
res, err := deps.SparkSubmitHandler(context.Background(), req)
|
||||
if err != nil {
|
||||
t.Fatalf("handler error: %v", err)
|
||||
}
|
||||
if res.IsError {
|
||||
t.Fatalf("unexpected error result: %v", res.Content)
|
||||
}
|
||||
|
||||
got, err := os.ReadFile(echoed)
|
||||
if err != nil {
|
||||
t.Fatalf("read echoed args: %v", err)
|
||||
}
|
||||
lines := strings.Split(strings.TrimSpace(string(got)), "\n")
|
||||
|
||||
want := []string{
|
||||
"--master", "yarn",
|
||||
"--queue", "default",
|
||||
"--executor-memory", "2G",
|
||||
"--executor-cores", "1",
|
||||
"--num-executors", "4",
|
||||
"--conf", "a=1",
|
||||
"--conf", "b=2",
|
||||
"--name", "wordcount-job",
|
||||
absPath,
|
||||
}
|
||||
if len(lines) != len(want) {
|
||||
t.Fatalf("lines=%v\nwant=%v", lines, want)
|
||||
}
|
||||
if lines[0] != "--master" {
|
||||
t.Errorf("line[0]=%q, want \"--master\"", lines[0])
|
||||
}
|
||||
for i, l := range lines {
|
||||
if l != want[i] {
|
||||
t.Errorf("line[%d]=%q, want %q", i, l, want[i])
|
||||
}
|
||||
}
|
||||
|
||||
if lines[len(lines)-1] != absPath {
|
||||
t.Errorf("last line=%q, want script_path %q", lines[len(lines)-1], absPath)
|
||||
}
|
||||
|
||||
if !strings.HasSuffix(absPath, "/"+fileID) {
|
||||
t.Errorf("absPath=%q does not end with fileID %q", absPath, fileID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSparkSubmit_PersistsSubmission(t *testing.T) {
|
||||
// Open a dedicated DB so we can insert upload_files metadata for the join.
|
||||
db, err := storage.Open(":memory:")
|
||||
if err != nil {
|
||||
t.Fatalf("open db: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
|
||||
dataDir := t.TempDir()
|
||||
store, err := uploads.New(filepath.Join(dataDir, "uploads"))
|
||||
if err != nil {
|
||||
t.Fatalf("new upload store: %v", err)
|
||||
}
|
||||
|
||||
deps := &Deps{
|
||||
HTTPClient: httpclient.New(httpclient.Config{Timeout: 5 * time.Second, MaxResponseBytes: 1 << 20}),
|
||||
AuditRepo: audit.NewRepo(db),
|
||||
ClusterRepo: db.Clusters(),
|
||||
SubmissionRepo: db.Submissions(),
|
||||
MaxResponseBytes: 1 << 20,
|
||||
DataDir: dataDir,
|
||||
UploadStore: store,
|
||||
SparkSubmitTimeout: 5 * time.Second,
|
||||
}
|
||||
repo := db.Clusters()
|
||||
subRepo := db.Submissions()
|
||||
|
||||
fileID, _, _, _, absPath, err := store.Save([]byte("print('hi')\n"), "hello.py")
|
||||
if err != nil {
|
||||
t.Fatalf("save upload: %v", err)
|
||||
}
|
||||
if err := db.Uploads().Create(context.Background(), fileID, "hello.py", 14, "dummy-sha", time.Now()); err != nil {
|
||||
t.Fatalf("create upload record: %v", err)
|
||||
}
|
||||
|
||||
echoScript := filepath.Join(t.TempDir(), "echo-submit.sh")
|
||||
if err := os.WriteFile(echoScript, []byte("#!/bin/sh\necho 'tracking URL: http://rm.example.com:8088/proxy/application_persist_0001/' >&2\necho 'Submitted application application_persist_0001'\n"), 0o755); err != nil {
|
||||
t.Fatalf("write echo script: %v", err)
|
||||
}
|
||||
|
||||
createCluster(t, repo, &cluster.Cluster{
|
||||
ID: "cluster-persist",
|
||||
Name: "Persist",
|
||||
IsActive: true,
|
||||
AuthType: cluster.AuthNone,
|
||||
SparkSubmitExecuteBin: echoScript,
|
||||
})
|
||||
|
||||
req := newToolRequest(SparkSubmitName, map[string]any{
|
||||
"cluster_id": "cluster-persist",
|
||||
"master": "yarn",
|
||||
"deploy_mode": "cluster",
|
||||
"script_path": absPath,
|
||||
"queue": "default",
|
||||
"executor_memory": "2G",
|
||||
"executor_cores": 1,
|
||||
"num_executors": 4,
|
||||
})
|
||||
res, err := deps.SparkSubmitHandler(context.Background(), req)
|
||||
if err != nil {
|
||||
t.Fatalf("handler error: %v", err)
|
||||
}
|
||||
if res.IsError {
|
||||
t.Fatalf("unexpected error result: %v", res.Content)
|
||||
}
|
||||
|
||||
payload := resultJSON(t, res)
|
||||
if payload["app_id"] != "application_persist_0001" {
|
||||
t.Errorf("app_id=%v, want application_persist_0001", payload["app_id"])
|
||||
}
|
||||
if payload["tracking_url"] != "http://rm.example.com:8088/proxy/application_persist_0001/" {
|
||||
t.Errorf("tracking_url=%v, want tracking URL", payload["tracking_url"])
|
||||
}
|
||||
|
||||
subs, err := subRepo.List(context.Background(), storage.SubmissionFilter{Limit: 10})
|
||||
if err != nil {
|
||||
t.Fatalf("list submissions: %v", err)
|
||||
}
|
||||
if len(subs) != 1 {
|
||||
t.Fatalf("got %d submissions, want 1", len(subs))
|
||||
}
|
||||
s := subs[0]
|
||||
if s.FileID != fileID {
|
||||
t.Errorf("file_id=%q, want %q", s.FileID, fileID)
|
||||
}
|
||||
if s.AppID != "application_persist_0001" {
|
||||
t.Errorf("app_id=%q, want application_persist_0001", s.AppID)
|
||||
}
|
||||
if s.TrackingURL != "http://rm.example.com:8088/proxy/application_persist_0001/" {
|
||||
t.Errorf("tracking_url=%q, want tracking URL", s.TrackingURL)
|
||||
}
|
||||
if s.ClusterID != "cluster-persist" {
|
||||
t.Errorf("cluster_id=%q, want cluster-persist", s.ClusterID)
|
||||
}
|
||||
if s.FileName != "hello.py" {
|
||||
t.Errorf("file_name=%q, want hello.py", s.FileName)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSparkSubmit_MissingRequiredField(t *testing.T) {
|
||||
deps, _, _, _ := testDepsWithDataDir(t)
|
||||
store := deps.UploadStore
|
||||
_, _, _, _, mintedPath, err := store.Save([]byte("# dummy\n"), "dummy.py")
|
||||
if err != nil {
|
||||
t.Fatalf("save upload: %v", err)
|
||||
}
|
||||
|
||||
req := newToolRequest(SparkSubmitName, map[string]any{
|
||||
"cluster_id": "cluster-echo",
|
||||
"deploy_mode": "cluster",
|
||||
"script_path": mintedPath,
|
||||
"queue": "default",
|
||||
"executor_memory": "2G",
|
||||
"executor_cores": 1,
|
||||
"num_executors": 4,
|
||||
})
|
||||
res, err := deps.SparkSubmitHandler(context.Background(), req)
|
||||
if err != nil {
|
||||
t.Fatalf("handler error: %v", err)
|
||||
}
|
||||
if !res.IsError {
|
||||
t.Fatalf("expected error result, got: %v", res.Content)
|
||||
}
|
||||
text, ok := mcp.AsTextContent(res.Content[0])
|
||||
if !ok {
|
||||
t.Fatalf("content is not text: %T", res.Content[0])
|
||||
}
|
||||
if !strings.Contains(text.Text, "master") {
|
||||
t.Errorf("error text=%q, want mention of master", text.Text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSparkSubmit_BadSparkConfValue(t *testing.T) {
|
||||
deps, _, _, _ := testDepsWithDataDir(t)
|
||||
store := deps.UploadStore
|
||||
_, _, _, _, mintedPath, err := store.Save([]byte("# dummy\n"), "dummy.py")
|
||||
if err != nil {
|
||||
t.Fatalf("save upload: %v", err)
|
||||
}
|
||||
|
||||
req := newToolRequest(SparkSubmitName, map[string]any{
|
||||
"cluster_id": "cluster-echo",
|
||||
"master": "yarn",
|
||||
"deploy_mode": "cluster",
|
||||
"script_path": mintedPath,
|
||||
"queue": "default",
|
||||
"executor_memory": "2G",
|
||||
"executor_cores": 1,
|
||||
"num_executors": 4,
|
||||
"spark_conf": map[string]any{"a": 1},
|
||||
})
|
||||
res, err := deps.SparkSubmitHandler(context.Background(), req)
|
||||
if err != nil {
|
||||
t.Fatalf("handler error: %v", err)
|
||||
}
|
||||
if !res.IsError {
|
||||
t.Fatalf("expected error result, got: %v", res.Content)
|
||||
}
|
||||
text, ok := mcp.AsTextContent(res.Content[0])
|
||||
if !ok {
|
||||
t.Fatalf("content is not text: %T", res.Content[0])
|
||||
}
|
||||
if !strings.Contains(text.Text, "spark_conf") {
|
||||
t.Errorf("error text=%q, want mention of spark_conf", text.Text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSparkSubmit_RejectsNonMintedPath(t *testing.T) {
|
||||
deps, _, _, _ := testDepsWithDataDir(t)
|
||||
|
||||
unminted := filepath.Join(t.TempDir(), "unminted.py")
|
||||
if err := os.WriteFile(unminted, []byte("print('not from upload_file')\n"), 0o644); err != nil {
|
||||
t.Fatalf("write unminted file: %v", err)
|
||||
}
|
||||
|
||||
req := newToolRequest(SparkSubmitName, map[string]any{
|
||||
"cluster_id": "cluster-echo",
|
||||
"master": "yarn",
|
||||
"deploy_mode": "cluster",
|
||||
"script_path": unminted,
|
||||
"queue": "default",
|
||||
"executor_memory": "2G",
|
||||
"executor_cores": 1,
|
||||
"num_executors": 4,
|
||||
})
|
||||
res, err := deps.SparkSubmitHandler(context.Background(), req)
|
||||
if err != nil {
|
||||
t.Fatalf("handler error: %v", err)
|
||||
}
|
||||
if !res.IsError {
|
||||
t.Fatalf("expected error result, got: %v", res.Content)
|
||||
}
|
||||
text, ok := mcp.AsTextContent(res.Content[0])
|
||||
if !ok {
|
||||
t.Fatalf("content is not text: %T", res.Content[0])
|
||||
}
|
||||
if !strings.Contains(text.Text, "not from upload_file") {
|
||||
t.Errorf("error text=%q, want mention of not from upload_file", text.Text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSparkSubmit_EmptyMaster(t *testing.T) {
|
||||
deps, _, _, _ := testDepsWithDataDir(t)
|
||||
store := deps.UploadStore
|
||||
_, _, _, _, mintedPath, err := store.Save([]byte("# dummy\n"), "dummy.py")
|
||||
if err != nil {
|
||||
t.Fatalf("save upload: %v", err)
|
||||
}
|
||||
|
||||
req := newToolRequest(SparkSubmitName, map[string]any{
|
||||
"cluster_id": "cluster-echo",
|
||||
"master": "",
|
||||
"deploy_mode": "cluster",
|
||||
"script_path": mintedPath,
|
||||
"queue": "default",
|
||||
"executor_memory": "2G",
|
||||
"executor_cores": 1,
|
||||
"num_executors": 4,
|
||||
})
|
||||
res, err := deps.SparkSubmitHandler(context.Background(), req)
|
||||
if err != nil {
|
||||
t.Fatalf("handler error: %v", err)
|
||||
}
|
||||
if !res.IsError {
|
||||
t.Fatalf("expected error result, got: %v", res.Content)
|
||||
}
|
||||
text, ok := mcp.AsTextContent(res.Content[0])
|
||||
if !ok {
|
||||
t.Fatalf("content is not text: %T", res.Content[0])
|
||||
}
|
||||
if !strings.Contains(text.Text, "master is required") && !strings.Contains(text.Text, "master") {
|
||||
t.Errorf("error text=%q, want mention of master", text.Text)
|
||||
}
|
||||
}
|
||||
@@ -4,11 +4,12 @@ import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
|
||||
"spark-mcp-go/internal/audit"
|
||||
)
|
||||
|
||||
const UploadFileName = "upload_file"
|
||||
@@ -19,7 +20,7 @@ var filenameRegex = regexp.MustCompile(`^[a-zA-Z0-9._-]+$`)
|
||||
// NewUploadFileTool returns the schema for the upload_file MCP Tool.
|
||||
func NewUploadFileTool() mcp.Tool {
|
||||
return mcp.NewTool(UploadFileName,
|
||||
mcp.WithDescription("Upload a script or config file to ./data/uploads/ so it can be referenced by spark_submit later."),
|
||||
mcp.WithDescription("Upload a script or config file to the server's data/uploads/ directory. Returns {file_id, path, name, size, sha256}. To submit it, pass the returned `path` as the `script_path` field in spark_submit, along with the structured fields (master, deploy_mode, queue, executor_memory, executor_cores, num_executors; optional spark_conf, extra_args)."),
|
||||
mcp.WithString("filename",
|
||||
mcp.Required(),
|
||||
mcp.Description("Plain file name without path separators (1-128 chars, [a-zA-Z0-9._-])"),
|
||||
@@ -72,17 +73,32 @@ func (d *Deps) UploadFileHandler(ctx context.Context, req mcp.CallToolRequest) (
|
||||
data = []byte(content)
|
||||
}
|
||||
|
||||
final := filepath.Join(d.DataDir, "uploads", filename)
|
||||
if err := os.MkdirAll(filepath.Dir(final), 0o750); err != nil {
|
||||
return errResult("upload_file: create uploads dir: " + err.Error()), nil
|
||||
fileID, _, size, sha256Hex, absPath, err := d.UploadStore.Save(data, filename)
|
||||
if err != nil {
|
||||
return errResult("upload_file: save upload: " + err.Error()), nil
|
||||
}
|
||||
if err := os.WriteFile(final, data, 0o640); err != nil {
|
||||
return errResult("upload_file: write file: " + 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{
|
||||
"path": fmt.Sprintf("uploads/%s", filename),
|
||||
"size": len(data),
|
||||
"file_id": fileID,
|
||||
"path": absPath,
|
||||
"name": filename,
|
||||
"size": size,
|
||||
"sha256": sha256Hex,
|
||||
}
|
||||
return textResult(encodeJSON(result)), nil
|
||||
}
|
||||
|
||||
@@ -38,6 +38,28 @@ CREATE TABLE IF NOT EXISTS audit_log (
|
||||
details TEXT
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_audit_ts ON audit_log(timestamp DESC);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS upload_files (
|
||||
file_id TEXT PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
size INTEGER NOT NULL,
|
||||
sha256 TEXT NOT NULL,
|
||||
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()
|
||||
}
|
||||
@@ -0,0 +1,149 @@
|
||||
package storage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// UploadMeta is the database index record for an uploaded file.
|
||||
//
|
||||
// The .meta.json sidecar remains the source of truth; this struct is the
|
||||
// queryable mirror used by the admin UI.
|
||||
type UploadMeta struct {
|
||||
FileID string `json:"file_id"`
|
||||
Name string `json:"name"`
|
||||
Size int64 `json:"size"`
|
||||
SHA256 string `json:"sha256"`
|
||||
UploadedAt time.Time `json:"uploaded_at"`
|
||||
}
|
||||
|
||||
// UploadRepo provides CRUD operations for upload file index records.
|
||||
type UploadRepo struct {
|
||||
db *DB
|
||||
}
|
||||
|
||||
// Uploads returns a repository bound to this DB handle.
|
||||
func (d *DB) Uploads() *UploadRepo {
|
||||
return &UploadRepo{db: d}
|
||||
}
|
||||
|
||||
// Create inserts a new upload index record. Timestamps are stored as Unix
|
||||
// nanoseconds for stable ORDER BY semantics.
|
||||
func (r *UploadRepo) Create(ctx context.Context, fileID, name string, size int64, sha256Hex string, uploadedAt time.Time) error {
|
||||
_, err := r.db.sqlDB.ExecContext(ctx, `
|
||||
INSERT INTO upload_files (file_id, name, size, sha256, uploaded_at)
|
||||
VALUES (?, ?, ?, ?, ?)`,
|
||||
fileID, name, size, sha256Hex, uploadedAt.UnixNano(),
|
||||
)
|
||||
if err != nil {
|
||||
if strings.Contains(err.Error(), "UNIQUE constraint failed") {
|
||||
return fmt.Errorf("storage: create upload %s: primary key conflict: %w", fileID, err)
|
||||
}
|
||||
return fmt.Errorf("storage: create upload %s: %w", fileID, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Get retrieves an upload index record by file_id. If the record does not
|
||||
// exist, it returns ErrNotFound.
|
||||
func (r *UploadRepo) Get(ctx context.Context, fileID string) (UploadMeta, error) {
|
||||
row := r.db.sqlDB.QueryRowContext(ctx, `
|
||||
SELECT file_id, name, size, sha256, uploaded_at
|
||||
FROM upload_files
|
||||
WHERE file_id = ?`, fileID)
|
||||
|
||||
m, err := scanUpload(row.Scan)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return UploadMeta{}, ErrNotFound
|
||||
}
|
||||
if err != nil {
|
||||
return UploadMeta{}, fmt.Errorf("storage: get upload %s: %w", fileID, err)
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
// List returns upload index records ordered by uploaded_at descending.
|
||||
//
|
||||
// If search is non-empty, file_name or any associated spark_submit app_id is
|
||||
// filtered with a case-insensitive LIKE.
|
||||
// A limit of zero or less falls back to 100; values above 500 are capped.
|
||||
func (r *UploadRepo) List(ctx context.Context, search string, limit int) ([]UploadMeta, error) {
|
||||
if limit <= 0 {
|
||||
limit = 100
|
||||
}
|
||||
if limit > 500 {
|
||||
limit = 500
|
||||
}
|
||||
|
||||
var rows *sql.Rows
|
||||
var err error
|
||||
if search != "" {
|
||||
rows, err = r.db.sqlDB.QueryContext(ctx, `
|
||||
SELECT DISTINCT u.file_id, u.name, u.size, u.sha256, u.uploaded_at
|
||||
FROM upload_files u
|
||||
LEFT JOIN spark_submissions s ON u.file_id = s.file_id
|
||||
WHERE u.name LIKE ? OR s.app_id LIKE ?
|
||||
ORDER BY u.uploaded_at DESC
|
||||
LIMIT ?`,
|
||||
"%"+search+"%", "%"+search+"%", limit,
|
||||
)
|
||||
} else {
|
||||
rows, err = r.db.sqlDB.QueryContext(ctx, `
|
||||
SELECT u.file_id, u.name, u.size, u.sha256, u.uploaded_at
|
||||
FROM upload_files u
|
||||
ORDER BY u.uploaded_at DESC
|
||||
LIMIT ?`, limit)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("storage: list uploads: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var out []UploadMeta
|
||||
for rows.Next() {
|
||||
m, err := scanUpload(rows.Scan)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("storage: list uploads: %w", err)
|
||||
}
|
||||
out = append(out, m)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, fmt.Errorf("storage: list uploads: %w", err)
|
||||
}
|
||||
if out == nil {
|
||||
// Force empty slice (not null) in JSON.
|
||||
out = []UploadMeta{}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// Delete removes an upload index record by file_id. If the record does not
|
||||
// exist, it returns ErrNotFound.
|
||||
func (r *UploadRepo) Delete(ctx context.Context, fileID string) error {
|
||||
res, err := r.db.sqlDB.ExecContext(ctx, `DELETE FROM upload_files WHERE file_id = ?`, fileID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("storage: delete upload %s: %w", fileID, err)
|
||||
}
|
||||
n, err := res.RowsAffected()
|
||||
if err != nil {
|
||||
return fmt.Errorf("storage: delete upload %s: rows affected: %w", fileID, err)
|
||||
}
|
||||
if n == 0 {
|
||||
return ErrNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func scanUpload(scan scanFunc) (UploadMeta, error) {
|
||||
var m UploadMeta
|
||||
var uploadedAt int64
|
||||
if err := scan(&m.FileID, &m.Name, &m.Size, &m.SHA256, &uploadedAt); err != nil {
|
||||
return UploadMeta{}, err
|
||||
}
|
||||
m.UploadedAt = time.Unix(0, uploadedAt)
|
||||
return m, nil
|
||||
}
|
||||
@@ -0,0 +1,162 @@
|
||||
package storage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func newUploadRepo(t *testing.T) *UploadRepo {
|
||||
t.Helper()
|
||||
db, err := Open(":memory:")
|
||||
if err != nil {
|
||||
t.Fatalf("Open in-memory db: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
return db.Uploads()
|
||||
}
|
||||
|
||||
func assertUploadEqual(t *testing.T, got, want UploadMeta) {
|
||||
t.Helper()
|
||||
if got.FileID != want.FileID {
|
||||
t.Errorf("FileID: got %q, want %q", got.FileID, want.FileID)
|
||||
}
|
||||
if got.Name != want.Name {
|
||||
t.Errorf("Name: got %q, want %q", got.Name, want.Name)
|
||||
}
|
||||
if got.Size != want.Size {
|
||||
t.Errorf("Size: got %d, want %d", got.Size, want.Size)
|
||||
}
|
||||
if got.SHA256 != want.SHA256 {
|
||||
t.Errorf("SHA256: got %q, want %q", got.SHA256, want.SHA256)
|
||||
}
|
||||
if !got.UploadedAt.Equal(want.UploadedAt) {
|
||||
t.Errorf("UploadedAt: got %v, want %v", got.UploadedAt, want.UploadedAt)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadRepo_CreateGetListDelete(t *testing.T) {
|
||||
repo := newUploadRepo(t)
|
||||
ctx := context.Background()
|
||||
uploadedAt := time.Unix(0, time.Now().UnixNano())
|
||||
want := UploadMeta{
|
||||
FileID: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
|
||||
Name: "report.csv",
|
||||
Size: 42,
|
||||
SHA256: "deadbeef",
|
||||
UploadedAt: uploadedAt,
|
||||
}
|
||||
if err := repo.Create(ctx, want.FileID, want.Name, want.Size, want.SHA256, want.UploadedAt); err != nil {
|
||||
t.Fatalf("Create: %v", err)
|
||||
}
|
||||
got, err := repo.Get(ctx, want.FileID)
|
||||
if err != nil {
|
||||
t.Fatalf("Get: %v", err)
|
||||
}
|
||||
assertUploadEqual(t, got, want)
|
||||
if err := repo.Create(ctx, want.FileID, "other", 1, "abcd", time.Now()); err == nil {
|
||||
t.Errorf("duplicate Create succeeded, want error")
|
||||
}
|
||||
list, err := repo.List(ctx, "", 100)
|
||||
if err != nil {
|
||||
t.Fatalf("List: %v", err)
|
||||
}
|
||||
if len(list) != 1 {
|
||||
t.Fatalf("List len: got %d, want 1", len(list))
|
||||
}
|
||||
assertUploadEqual(t, list[0], want)
|
||||
if err := repo.Delete(ctx, want.FileID); err != nil {
|
||||
t.Fatalf("Delete: %v", err)
|
||||
}
|
||||
if _, err := repo.Get(ctx, want.FileID); !errors.Is(err, ErrNotFound) {
|
||||
t.Fatalf("Get after Delete: got %v, want ErrNotFound", err)
|
||||
}
|
||||
if err := repo.Delete(ctx, want.FileID); !errors.Is(err, ErrNotFound) {
|
||||
t.Fatalf("Delete again: got %v, want ErrNotFound", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadRepo_ListOrderAndSearch(t *testing.T) {
|
||||
repo := newUploadRepo(t)
|
||||
ctx := context.Background()
|
||||
base := time.Unix(0, time.Now().UnixNano())
|
||||
records := []UploadMeta{
|
||||
{FileID: "00000000000000000000000000000001", Name: "alpha.csv", Size: 1, SHA256: "a", UploadedAt: base.Add(-2 * time.Hour)},
|
||||
{FileID: "00000000000000000000000000000002", Name: "beta.log", Size: 2, SHA256: "b", UploadedAt: base.Add(-1 * time.Hour)},
|
||||
{FileID: "00000000000000000000000000000003", Name: "gamma.csv", Size: 3, SHA256: "c", UploadedAt: base},
|
||||
}
|
||||
for _, r := range records {
|
||||
if err := repo.Create(ctx, r.FileID, r.Name, r.Size, r.SHA256, r.UploadedAt); err != nil {
|
||||
t.Fatalf("Create %s: %v", r.FileID, err)
|
||||
}
|
||||
}
|
||||
all, err := repo.List(ctx, "", 100)
|
||||
if err != nil {
|
||||
t.Fatalf("List: %v", err)
|
||||
}
|
||||
if len(all) != 3 {
|
||||
t.Fatalf("List len: got %d, want 3", len(all))
|
||||
}
|
||||
wantOrder := []string{"00000000000000000000000000000003", "00000000000000000000000000000002", "00000000000000000000000000000001"}
|
||||
for i, id := range wantOrder {
|
||||
if all[i].FileID != id {
|
||||
t.Errorf("List order[%d]: got %q, want %q", i, all[i].FileID, id)
|
||||
}
|
||||
}
|
||||
csv, err := repo.List(ctx, "csv", 100)
|
||||
if err != nil {
|
||||
t.Fatalf("List search: %v", err)
|
||||
}
|
||||
if len(csv) != 2 {
|
||||
t.Fatalf("search csv len: got %d, want 2", len(csv))
|
||||
}
|
||||
if csv[0].Name != "gamma.csv" || csv[1].Name != "alpha.csv" {
|
||||
t.Errorf("search csv order: got %v", []string{csv[0].Name, csv[1].Name})
|
||||
}
|
||||
beta, err := repo.List(ctx, "beta", 100)
|
||||
if err != nil {
|
||||
t.Fatalf("List search beta: %v", err)
|
||||
}
|
||||
if len(beta) != 1 || beta[0].Name != "beta.log" {
|
||||
t.Errorf("search beta: got %+v", beta)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadRepo_ListLimitCap(t *testing.T) {
|
||||
repo := newUploadRepo(t)
|
||||
ctx := context.Background()
|
||||
base := time.Unix(0, time.Now().UnixNano())
|
||||
for i := 0; i < 10; i++ {
|
||||
id := fmt.Sprintf("%032d", i)
|
||||
if err := repo.Create(ctx, id, "x", int64(i), "h", base.Add(time.Duration(i)*time.Second)); err != nil {
|
||||
t.Fatalf("Create %d: %v", i, err)
|
||||
}
|
||||
}
|
||||
list, err := repo.List(ctx, "", 0)
|
||||
if err != nil {
|
||||
t.Fatalf("List default limit: %v", err)
|
||||
}
|
||||
// Default limit is 100 but only 10 rows exist.
|
||||
if len(list) != 10 {
|
||||
t.Errorf("default limit: got %d, want 10", len(list))
|
||||
}
|
||||
|
||||
list, err = repo.List(ctx, "", 501)
|
||||
if err != nil {
|
||||
t.Fatalf("List cap: %v", err)
|
||||
}
|
||||
if len(list) != 10 {
|
||||
t.Errorf("cap limit: got %d, want 10", len(list))
|
||||
}
|
||||
|
||||
// Explicit small limit is respected.
|
||||
list, err = repo.List(ctx, "", 3)
|
||||
if err != nil {
|
||||
t.Fatalf("List small limit: %v", err)
|
||||
}
|
||||
if len(list) != 3 {
|
||||
t.Errorf("small limit: got %d, want 3", len(list))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,345 @@
|
||||
// Package uploads stores files uploaded through the upload_file MCP Tool on
|
||||
// the server's local filesystem. Files are named by a server-minted file ID
|
||||
// rather than the user's original filename, which is kept only in a sidecar.
|
||||
package uploads
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"spark-mcp-go/internal/storage"
|
||||
)
|
||||
|
||||
var fileIDRegex = regexp.MustCompile(`^[0-9a-f]{32}$`)
|
||||
|
||||
// Store is a local file store rooted at Root.
|
||||
type Store struct {
|
||||
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
|
||||
// does not exist. The root must be an absolute path.
|
||||
func New(root string) (Store, error) {
|
||||
if root == "" {
|
||||
return Store{}, fmt.Errorf("uploads: root must not be empty")
|
||||
}
|
||||
absRoot, err := filepath.Abs(root)
|
||||
if err != nil {
|
||||
return Store{}, fmt.Errorf("uploads: resolve root: %w", err)
|
||||
}
|
||||
if !filepath.IsAbs(absRoot) {
|
||||
return Store{}, fmt.Errorf("uploads: root %q is not an absolute path", absRoot)
|
||||
}
|
||||
|
||||
if err := os.MkdirAll(absRoot, 0o750); err != nil {
|
||||
return Store{}, fmt.Errorf("uploads: create root: %w", err)
|
||||
}
|
||||
|
||||
info, err := os.Stat(absRoot)
|
||||
if err != nil {
|
||||
return Store{}, fmt.Errorf("uploads: stat root: %w", err)
|
||||
}
|
||||
if !info.IsDir() {
|
||||
return Store{}, fmt.Errorf("uploads: root %q is not a directory", absRoot)
|
||||
}
|
||||
|
||||
return Store{Root: absRoot}, nil
|
||||
}
|
||||
|
||||
type sidecar struct {
|
||||
Name string `json:"name"`
|
||||
Size int64 `json:"size"`
|
||||
Sha256 string `json:"sha256"`
|
||||
UploadedAt time.Time `json:"uploaded_at"`
|
||||
}
|
||||
|
||||
// Save writes data to Root/<fileID> and a sidecar to Root/<fileID>.meta.json.
|
||||
// The original filename is recorded in the sidecar and never appears in the
|
||||
// on-disk name. It returns the minted file ID, the original name, the size,
|
||||
// the hex-encoded SHA-256 digest, and the absolute path to the data file.
|
||||
func (s *Store) Save(data []byte, originalName string) (fileID, name string, size int64, sha256Hex string, absPath string, err error) {
|
||||
if originalName == "" {
|
||||
return "", "", 0, "", "", fmt.Errorf("uploads: original name must not be empty")
|
||||
}
|
||||
|
||||
sum := sha256.Sum256(data)
|
||||
sha256Hex = hex.EncodeToString(sum[:])
|
||||
|
||||
for {
|
||||
fileID, err = newFileID()
|
||||
if err != nil {
|
||||
return "", "", 0, "", "", err
|
||||
}
|
||||
absPath = filepath.Join(s.Root, fileID)
|
||||
_, err = os.Stat(absPath)
|
||||
if os.IsNotExist(err) {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return "", "", 0, "", "", fmt.Errorf("uploads: check fileID collision: %w", err)
|
||||
}
|
||||
// Collision is astronomically unlikely; regenerate just in case.
|
||||
}
|
||||
|
||||
if err := os.WriteFile(absPath, data, 0o640); err != nil {
|
||||
return "", "", 0, "", "", fmt.Errorf("uploads: write data file: %w", err)
|
||||
}
|
||||
|
||||
meta := sidecar{
|
||||
Name: originalName,
|
||||
Size: int64(len(data)),
|
||||
Sha256: sha256Hex,
|
||||
UploadedAt: time.Now(),
|
||||
}
|
||||
metaBytes, err := json.Marshal(meta)
|
||||
if err != nil {
|
||||
// Best-effort cleanup of the orphaned data file.
|
||||
_ = os.Remove(absPath)
|
||||
return "", "", 0, "", "", fmt.Errorf("uploads: marshal sidecar: %w", err)
|
||||
}
|
||||
|
||||
metaPath := absPath + ".meta.json"
|
||||
if err := os.WriteFile(metaPath, metaBytes, 0o600); err != nil {
|
||||
_ = os.Remove(absPath)
|
||||
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
|
||||
}
|
||||
|
||||
// AbsPath returns the absolute on-disk path for a well-formed fileID.
|
||||
func (s *Store) AbsPath(fileID string) (string, error) {
|
||||
if !fileIDRegex.MatchString(fileID) {
|
||||
return "", fmt.Errorf("uploads: malformed fileID %q", fileID)
|
||||
}
|
||||
return filepath.Join(s.Root, fileID), nil
|
||||
}
|
||||
|
||||
// Validate checks that absPath is inside Root and that its sidecar exists.
|
||||
// It returns the fileID so callers can canonicalize the path.
|
||||
func (s *Store) Validate(absPath string) (string, error) {
|
||||
clean, err := filepath.Abs(absPath)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("uploads: resolve path: %w", err)
|
||||
}
|
||||
|
||||
rel, err := filepath.Rel(s.Root, clean)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("uploads: path %q is not under Root", absPath)
|
||||
}
|
||||
if rel == "." || strings.HasPrefix(rel, "..") || filepath.IsAbs(rel) {
|
||||
return "", fmt.Errorf("uploads: path %q is not under Root", absPath)
|
||||
}
|
||||
|
||||
if strings.HasSuffix(rel, ".meta.json") {
|
||||
return "", fmt.Errorf("uploads: path %q is a sidecar, pass the data file path", absPath)
|
||||
}
|
||||
|
||||
fileID := strings.TrimSuffix(rel, ".meta.json")
|
||||
if !fileIDRegex.MatchString(fileID) {
|
||||
return "", fmt.Errorf("uploads: path %q does not denote a valid upload", absPath)
|
||||
}
|
||||
|
||||
metaPath := filepath.Join(s.Root, fileID+".meta.json")
|
||||
if _, err := os.Stat(metaPath); err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return "", fmt.Errorf("uploads: sidecar missing for %s", fileID)
|
||||
}
|
||||
return "", fmt.Errorf("uploads: stat sidecar for %s: %w", fileID, err)
|
||||
}
|
||||
|
||||
return fileID, nil
|
||||
}
|
||||
|
||||
// Sweep deletes data files (and their sidecars) whose uploaded_at is older
|
||||
// than ttl, as well as orphan sidecars that have no matching data file. For
|
||||
// data files whose sidecar is missing, the file's mtime is used as a fallback.
|
||||
func (s *Store) Sweep(ctx context.Context, ttl time.Duration) (deleted int, err error) {
|
||||
entries, err := os.ReadDir(s.Root)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("uploads: read root: %w", err)
|
||||
}
|
||||
|
||||
type item struct {
|
||||
data bool
|
||||
meta bool
|
||||
}
|
||||
items := make(map[string]*item)
|
||||
for _, e := range entries {
|
||||
name := e.Name()
|
||||
base := strings.TrimSuffix(name, ".meta.json")
|
||||
if base == name && fileIDRegex.MatchString(name) {
|
||||
it := items[name]
|
||||
if it == nil {
|
||||
it = &item{}
|
||||
items[name] = it
|
||||
}
|
||||
it.data = true
|
||||
continue
|
||||
}
|
||||
if fileIDRegex.MatchString(base) {
|
||||
it := items[base]
|
||||
if it == nil {
|
||||
it = &item{}
|
||||
items[base] = it
|
||||
}
|
||||
it.meta = true
|
||||
}
|
||||
// Ignore unrelated files silently.
|
||||
}
|
||||
|
||||
cutoff := time.Now().Add(-ttl)
|
||||
|
||||
for fileID, it := range items {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return deleted, err
|
||||
}
|
||||
dataPath := filepath.Join(s.Root, fileID)
|
||||
metaPath := dataPath + ".meta.json"
|
||||
|
||||
if it.data {
|
||||
expired := false
|
||||
corruptSidecar := false
|
||||
if it.meta {
|
||||
uploadedAt, parseErr := readUploadedAt(metaPath)
|
||||
if parseErr == nil && uploadedAt.Before(cutoff) {
|
||||
expired = true
|
||||
} else if parseErr != nil {
|
||||
corruptSidecar = true
|
||||
}
|
||||
}
|
||||
if !expired && (corruptSidecar || !it.meta) {
|
||||
info, statErr := os.Stat(dataPath)
|
||||
if statErr == nil && info.ModTime().Before(cutoff) {
|
||||
expired = true
|
||||
}
|
||||
}
|
||||
if expired {
|
||||
if it.meta {
|
||||
_ = os.Remove(metaPath)
|
||||
}
|
||||
if rmErr := os.Remove(dataPath); rmErr == nil {
|
||||
deleted++
|
||||
}
|
||||
}
|
||||
} else if it.meta {
|
||||
// Orphan sidecar with no data file: reap it.
|
||||
if rmErr := os.Remove(metaPath); rmErr == nil {
|
||||
deleted++
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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) {
|
||||
var b [16]byte
|
||||
if _, err := rand.Read(b[:]); err != nil {
|
||||
return "", fmt.Errorf("uploads: generate fileID: %w", err)
|
||||
}
|
||||
return hex.EncodeToString(b[:]), nil
|
||||
}
|
||||
|
||||
func readUploadedAt(path string) (time.Time, error) {
|
||||
sc, err := readSidecar(path)
|
||||
if err != nil {
|
||||
return time.Time{}, err
|
||||
}
|
||||
if sc.UploadedAt.IsZero() {
|
||||
return time.Time{}, fmt.Errorf("uploads: missing uploaded_at")
|
||||
}
|
||||
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
|
||||
}
|
||||
@@ -0,0 +1,533 @@
|
||||
package uploads
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"spark-mcp-go/internal/storage"
|
||||
)
|
||||
|
||||
func TestStore_SaveAndRetrieve(t *testing.T) {
|
||||
store, err := New(t.TempDir())
|
||||
if err != nil {
|
||||
t.Fatalf("new store: %v", err)
|
||||
}
|
||||
|
||||
data := []byte("hello uploads")
|
||||
sum := sha256.Sum256(data)
|
||||
wantSHA := hex.EncodeToString(sum[:])
|
||||
|
||||
fileID, name, size, sha256Hex, absPath, err := store.Save(data, "hello.txt")
|
||||
if err != nil {
|
||||
t.Fatalf("save: %v", err)
|
||||
}
|
||||
|
||||
if fileID == "" || !regexp.MustCompile(`^[0-9a-f]{32}$`).MatchString(fileID) {
|
||||
t.Errorf("fileID=%q, want 32 lowercase hex chars", fileID)
|
||||
}
|
||||
if name != "hello.txt" {
|
||||
t.Errorf("name=%q, want hello.txt", name)
|
||||
}
|
||||
if size != int64(len(data)) {
|
||||
t.Errorf("size=%d, want %d", size, len(data))
|
||||
}
|
||||
if sha256Hex != wantSHA {
|
||||
t.Errorf("sha256=%q, want %q", sha256Hex, wantSHA)
|
||||
}
|
||||
if !filepath.IsAbs(absPath) {
|
||||
t.Errorf("absPath=%q is not absolute", absPath)
|
||||
}
|
||||
if filepath.Base(absPath) != fileID {
|
||||
t.Errorf("absPath base=%q, want fileID %q", filepath.Base(absPath), fileID)
|
||||
}
|
||||
|
||||
gotData, err := os.ReadFile(absPath)
|
||||
if err != nil {
|
||||
t.Fatalf("read data file: %v", err)
|
||||
}
|
||||
if string(gotData) != string(data) {
|
||||
t.Errorf("data=%q, want %q", gotData, data)
|
||||
}
|
||||
|
||||
info, err := os.Stat(absPath)
|
||||
if err != nil {
|
||||
t.Fatalf("stat data file: %v", err)
|
||||
}
|
||||
if info.Mode().Perm() != 0o640 {
|
||||
t.Errorf("data mode=%o, want %o", info.Mode().Perm(), 0o640)
|
||||
}
|
||||
|
||||
metaPath := absPath + ".meta.json"
|
||||
metaBytes, err := os.ReadFile(metaPath)
|
||||
if err != nil {
|
||||
t.Fatalf("read sidecar: %v", err)
|
||||
}
|
||||
var sc sidecar
|
||||
if err := json.Unmarshal(metaBytes, &sc); err != nil {
|
||||
t.Fatalf("unmarshal sidecar: %v", err)
|
||||
}
|
||||
if sc.Name != "hello.txt" {
|
||||
t.Errorf("sidecar name=%q, want hello.txt", sc.Name)
|
||||
}
|
||||
if sc.Size != int64(len(data)) {
|
||||
t.Errorf("sidecar size=%d, want %d", sc.Size, len(data))
|
||||
}
|
||||
if sc.Sha256 != wantSHA {
|
||||
t.Errorf("sidecar sha256=%q, want %q", sc.Sha256, wantSHA)
|
||||
}
|
||||
if sc.UploadedAt.IsZero() {
|
||||
t.Errorf("sidecar uploaded_at is zero")
|
||||
}
|
||||
|
||||
metaInfo, err := os.Stat(metaPath)
|
||||
if err != nil {
|
||||
t.Fatalf("stat sidecar: %v", err)
|
||||
}
|
||||
if metaInfo.Mode().Perm() != 0o600 {
|
||||
t.Errorf("sidecar mode=%o, want %o", metaInfo.Mode().Perm(), 0o600)
|
||||
}
|
||||
|
||||
validatedID, err := store.Validate(absPath)
|
||||
if err != nil {
|
||||
t.Fatalf("validate: %v", err)
|
||||
}
|
||||
if validatedID != fileID {
|
||||
t.Errorf("validatedID=%q, want %q", validatedID, fileID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_AbsPath_RejectsMalformed(t *testing.T) {
|
||||
store, err := New(t.TempDir())
|
||||
if err != nil {
|
||||
t.Fatalf("new store: %v", err)
|
||||
}
|
||||
|
||||
validID := "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
|
||||
if _, err := store.AbsPath(validID); err != nil {
|
||||
t.Errorf("valid fileID rejected: %v", err)
|
||||
}
|
||||
|
||||
cases := []string{
|
||||
"",
|
||||
"abc",
|
||||
"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
|
||||
"0123456789abcdef0123456789abcdefg",
|
||||
"0123456789abcdef0123456789abcde ",
|
||||
}
|
||||
for _, id := range cases {
|
||||
t.Run(fmt.Sprintf("id=%q", id), func(t *testing.T) {
|
||||
_, err := store.AbsPath(id)
|
||||
if err == nil {
|
||||
t.Errorf("AbsPath(%q) succeeded, want error", id)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_Validate_RejectsOutsideRoot(t *testing.T) {
|
||||
if os.PathSeparator != '/' {
|
||||
t.Skip("Unix-style path test")
|
||||
}
|
||||
|
||||
store, err := New(t.TempDir())
|
||||
if err != nil {
|
||||
t.Fatalf("new store: %v", err)
|
||||
}
|
||||
|
||||
fileID := "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
|
||||
if err := os.WriteFile(filepath.Join(store.Root, fileID), []byte("x"), 0o640); err != nil {
|
||||
t.Fatalf("create data file: %v", err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(store.Root, fileID+".meta.json"), []byte(`{"uploaded_at":"`+time.Now().Format(time.RFC3339Nano)+`"}`), 0o600); err != nil {
|
||||
t.Fatalf("create sidecar: %v", err)
|
||||
}
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
path string
|
||||
}{
|
||||
{"parent", filepath.Join(store.Root, "..", "other", fileID)},
|
||||
{"dotdot", filepath.Join(store.Root, "..", "..", "tmp", fileID)},
|
||||
{"different_volume", "/other/volume/" + fileID},
|
||||
}
|
||||
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
_, err := store.Validate(c.path)
|
||||
if err == nil {
|
||||
t.Errorf("Validate(%q) succeeded, want error", c.path)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_Sweep_DeletesExpiredKeepsRecent(t *testing.T) {
|
||||
store, err := New(t.TempDir())
|
||||
if err != nil {
|
||||
t.Fatalf("new store: %v", err)
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
expiredID := "00000000000000000000000000000001"
|
||||
freshID := "00000000000000000000000000000002"
|
||||
noMetaID := "00000000000000000000000000000003"
|
||||
|
||||
create := func(id string, uploadedAt *time.Time, withData bool) {
|
||||
if withData {
|
||||
if err := os.WriteFile(filepath.Join(store.Root, id), []byte("x"), 0o640); err != nil {
|
||||
t.Fatalf("create data file %s: %v", id, err)
|
||||
}
|
||||
}
|
||||
if uploadedAt != nil {
|
||||
sc := map[string]any{
|
||||
"name": "x",
|
||||
"size": 1,
|
||||
"sha256": "abc",
|
||||
"uploaded_at": uploadedAt.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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
create(expiredID, timePtr(now.Add(-2*time.Hour)), true)
|
||||
create(freshID, timePtr(now), true)
|
||||
create(noMetaID, nil, true)
|
||||
|
||||
deleted, err := store.Sweep(context.Background(), 1*time.Hour)
|
||||
if err != nil {
|
||||
t.Fatalf("sweep: %v", err)
|
||||
}
|
||||
if deleted != 1 {
|
||||
t.Errorf("deleted=%d, want 1", deleted)
|
||||
}
|
||||
|
||||
if _, err := os.Stat(filepath.Join(store.Root, expiredID)); !os.IsNotExist(err) {
|
||||
t.Errorf("expired data file still exists")
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(store.Root, expiredID+".meta.json")); !os.IsNotExist(err) {
|
||||
t.Errorf("expired sidecar still exists")
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(store.Root, freshID)); err != nil {
|
||||
t.Errorf("fresh data file missing: %v", err)
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(store.Root, noMetaID)); err != nil {
|
||||
t.Errorf("no-sidecar data file missing: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_Sweep_OrphanSidecarReaped(t *testing.T) {
|
||||
store, err := New(t.TempDir())
|
||||
if err != nil {
|
||||
t.Fatalf("new store: %v", err)
|
||||
}
|
||||
|
||||
id := "00000000000000000000000000000004"
|
||||
sc := map[string]any{
|
||||
"name": "orphan",
|
||||
"size": 1,
|
||||
"sha256": "abc",
|
||||
"uploaded_at": time.Now().Add(-2 * time.Hour).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 orphan sidecar: %v", err)
|
||||
}
|
||||
|
||||
deleted, err := store.Sweep(context.Background(), 1*time.Hour)
|
||||
if err != nil {
|
||||
t.Fatalf("sweep: %v", err)
|
||||
}
|
||||
if deleted != 1 {
|
||||
t.Errorf("deleted=%d, want 1", deleted)
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(store.Root, id+".meta.json")); !os.IsNotExist(err) {
|
||||
t.Errorf("orphan sidecar still exists")
|
||||
}
|
||||
}
|
||||
|
||||
func timePtr(t time.Time) *time.Time {
|
||||
return &t
|
||||
}
|
||||
|
||||
func TestStore_Validate_RejectsSidecarPath(t *testing.T) {
|
||||
store, err := New(t.TempDir())
|
||||
if err != nil {
|
||||
t.Fatalf("new store: %v", err)
|
||||
}
|
||||
|
||||
fileID := "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
|
||||
if err := os.WriteFile(filepath.Join(store.Root, fileID), []byte("x"), 0o640); err != nil {
|
||||
t.Fatalf("create data file: %v", err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(store.Root, fileID+".meta.json"), []byte(`{"uploaded_at":"`+time.Now().Format(time.RFC3339Nano)+`"}`), 0o600); err != nil {
|
||||
t.Fatalf("create sidecar: %v", err)
|
||||
}
|
||||
|
||||
_, err = store.Validate(filepath.Join(store.Root, fileID+".meta.json"))
|
||||
if err == nil {
|
||||
t.Errorf("Validate(sidecar) succeeded, want error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_Sweep_HonorsContextCancel(t *testing.T) {
|
||||
store, err := New(t.TempDir())
|
||||
if err != nil {
|
||||
t.Fatalf("new store: %v", err)
|
||||
}
|
||||
|
||||
id := "00000000000000000000000000000006"
|
||||
if err := os.WriteFile(filepath.Join(store.Root, id), []byte("x"), 0o640); err != nil {
|
||||
t.Fatalf("create data file: %v", err)
|
||||
}
|
||||
sc := map[string]any{
|
||||
"name": "x",
|
||||
"size": 1,
|
||||
"sha256": "abc",
|
||||
"uploaded_at": time.Now().Add(-2 * time.Hour).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)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
deleted, err := store.Sweep(ctx, 1*time.Hour)
|
||||
if err != ctx.Err() {
|
||||
t.Fatalf("sweep err=%v, want %v", err, ctx.Err())
|
||||
}
|
||||
if deleted != 0 {
|
||||
t.Errorf("deleted=%d, want 0", deleted)
|
||||
}
|
||||
if _, existErr := os.Stat(filepath.Join(store.Root, id)); os.IsNotExist(existErr) {
|
||||
t.Logf("cancelled sweep managed to delete the file before checking ctx; this is timing-dependent and acceptable")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_Sweep_DeletesDataWithCorruptSidecar(t *testing.T) {
|
||||
store, err := New(t.TempDir())
|
||||
if err != nil {
|
||||
t.Fatalf("new store: %v", err)
|
||||
}
|
||||
|
||||
id := "00000000000000000000000000000005"
|
||||
dataPath := filepath.Join(store.Root, id)
|
||||
metaPath := dataPath + ".meta.json"
|
||||
|
||||
if err := os.WriteFile(dataPath, []byte("stale data"), 0o640); err != nil {
|
||||
t.Fatalf("create data file: %v", err)
|
||||
}
|
||||
// Corrupt sidecar: not valid JSON.
|
||||
if err := os.WriteFile(metaPath, []byte("not valid json"), 0o600); err != nil {
|
||||
t.Fatalf("create sidecar: %v", err)
|
||||
}
|
||||
// Set mtime well in the past so the fallback triggers.
|
||||
past := time.Now().Add(-2 * time.Hour)
|
||||
if err := os.Chtimes(dataPath, past, past); err != nil {
|
||||
t.Fatalf("set mtime: %v", err)
|
||||
}
|
||||
|
||||
deleted, err := store.Sweep(context.Background(), 1*time.Hour)
|
||||
if err != nil {
|
||||
t.Fatalf("sweep: %v", err)
|
||||
}
|
||||
if deleted != 1 {
|
||||
t.Errorf("deleted=%d, want 1", deleted)
|
||||
}
|
||||
if _, err := os.Stat(dataPath); !os.IsNotExist(err) {
|
||||
t.Errorf("data file with corrupt sidecar still exists")
|
||||
}
|
||||
if _, err := os.Stat(metaPath); !os.IsNotExist(err) {
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"net/http"
|
||||
"os"
|
||||
"os/signal"
|
||||
"path/filepath"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
@@ -26,8 +27,11 @@ import (
|
||||
"spark-mcp-go/internal/mcp/tools"
|
||||
"spark-mcp-go/internal/middleware"
|
||||
"spark-mcp-go/internal/storage"
|
||||
"spark-mcp-go/internal/uploads"
|
||||
)
|
||||
|
||||
const startupSweepTimeout = 5 * time.Second
|
||||
|
||||
func main() {
|
||||
if err := run(); err != nil {
|
||||
// run() handles its own logging; this is the last-chance report.
|
||||
@@ -65,6 +69,34 @@ func run() error {
|
||||
defer db.Close()
|
||||
logger.Info("storage.open", "path", cfg.SQLitePath)
|
||||
|
||||
uploadStore, err := uploads.New(filepath.Join(cfg.DataDir, "uploads"))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
sweepCtx, sweepCancel := context.WithTimeout(context.Background(), startupSweepTimeout)
|
||||
n, err := uploadStore.Sweep(sweepCtx, cfg.UploadTTL)
|
||||
sweepCancel()
|
||||
if errors.Is(err, context.DeadlineExceeded) {
|
||||
logger.Warn("uploads.sweep.timeout", "deleted", n, "err", err)
|
||||
} else if err != nil {
|
||||
logger.Error("uploads.sweep", "err", err)
|
||||
} else {
|
||||
logger.Info("uploads.sweep", "deleted", n)
|
||||
}
|
||||
|
||||
backfillCtx, backfillCancel := context.WithTimeout(context.Background(), startupSweepTimeout)
|
||||
n, err = uploadStore.Backfill(backfillCtx, db.Uploads())
|
||||
backfillCancel()
|
||||
if errors.Is(err, context.DeadlineExceeded) {
|
||||
logger.Warn("uploads.backfill.timeout", "inserted", n, "err", err)
|
||||
} else if err != nil {
|
||||
logger.Error("uploads.backfill", "err", err)
|
||||
} else {
|
||||
logger.Info("uploads.backfill", "inserted", n)
|
||||
}
|
||||
|
||||
uploadStore.SetRepo(db.Uploads())
|
||||
|
||||
gin.SetMode(cfg.GinMode)
|
||||
r := gin.New()
|
||||
r.Use(gin.Recovery())
|
||||
@@ -73,10 +105,11 @@ func run() error {
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "version": "0.0.0"})
|
||||
})
|
||||
|
||||
admin.Mount(r, db.Clusters(), audit.NewRepo(db), cfg.AdminTokens)
|
||||
admin.Mount(r, db.Clusters(), db.Uploads(), db.Submissions(), audit.NewRepo(db), &uploadStore, cfg.AdminTokens)
|
||||
|
||||
mcpHandler, err := mcpsrv.Handler(&tools.Deps{
|
||||
Logger: logger.With("component", "mcp"),
|
||||
AuditRepo: audit.NewRepo(db),
|
||||
ClusterRepo: db.Clusters(),
|
||||
SparkSubmitTimeout: cfg.SparkSubmitTimeout,
|
||||
HTTPClient: httpclient.New(httpclient.Config{
|
||||
@@ -85,6 +118,8 @@ func run() error {
|
||||
}),
|
||||
MaxResponseBytes: cfg.MaxResponseBytes,
|
||||
DataDir: cfg.DataDir,
|
||||
UploadStore: uploadStore,
|
||||
SubmissionRepo: db.Submissions(),
|
||||
AnalyzerThresholds: analyzer.Thresholds{
|
||||
DataSkewRatio: cfg.AnalyzerDataSkewRatio,
|
||||
GCPressureRatio: cfg.AnalyzerGCPressureRatio,
|
||||
|
||||
Reference in New Issue
Block a user