uploads,tools,config,main: 将 upload_file 返回路径改为服务器端绝对路径并透传给 spark_submit
之前 upload_file 把文件写到 DataDir/uploads/<filename>,返回相对路径
uploads/hello.txt。MCP 服务器与 agent 不在同一台机器,agent 无法构造服务
器本地路径,因此 spark_submit 无法定位脚本。
改为由新的 internal/uploads 包 mint 一个 32 字符十六进制 file_id,文件落盘为
DataDir/uploads/<file_id>,元数据写入 <file_id>.meta.json(原始文件名只存在
sidecar 里)。upload_file 返回 {file_id, path, name, size, sha256},其中 path
是绝对路径。spark_submit 的 description 明确要求 LLM 直接把 upload_file 返回
的 path 放进 args,不要自己构造路径。
为什么只在描述里约束而不在代码里拒绝非 mint 的绝对路径:集群本地已有路径
(如 /opt/spark/examples/pi.py)是合法的 spark-submit 参数,代码不能替
LLM 拒绝。
测试锁定:
- internal/uploads: Save 往返、AbsPath/Validate 非法路径、Sweep 过期/未过期/
孤立 sidecar
- internal/mcp/tools: upload_file 新响应字段、spark_submit 透传 mint 路径与
集群本地路径
刻意未做:S3/HDFS 上传、给 spark_submit 新增 file_id 参数、在代码层面拒绝
集群本地绝对路径。
破坏性变更:upload_file 响应从相对 path 改为绝对 path,并新增 file_id/name/
sha256 字段。该工具尚无外部调用者。
Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -22,6 +22,9 @@ MAX_RESPONSE_BYTES=1048576
|
|||||||
# Timeout for spark-submit process execution (default: 60s)
|
# Timeout for spark-submit process execution (default: 60s)
|
||||||
SPARK_SUBMIT_TIMEOUT=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 directory (default: ./data/logs)
|
||||||
LOG_DIR=./data/logs
|
LOG_DIR=./data/logs
|
||||||
|
|
||||||
|
|||||||
@@ -29,6 +29,7 @@ type Config struct {
|
|||||||
AnalyzerDataSkewRatio float64
|
AnalyzerDataSkewRatio float64
|
||||||
AnalyzerGCPressureRatio float64
|
AnalyzerGCPressureRatio float64
|
||||||
AnalyzerBottleneckShuffleGB float64
|
AnalyzerBottleneckShuffleGB float64
|
||||||
|
UploadTTL time.Duration
|
||||||
}
|
}
|
||||||
|
|
||||||
// Load reads configuration from environment variables and returns a populated Config.
|
// Load reads configuration from environment variables and returns a populated Config.
|
||||||
@@ -44,6 +45,7 @@ func Load() (*Config, error) {
|
|||||||
HTTPClientTimeout: 30 * time.Second,
|
HTTPClientTimeout: 30 * time.Second,
|
||||||
MaxResponseBytes: 1 << 20,
|
MaxResponseBytes: 1 << 20,
|
||||||
SparkSubmitTimeout: 60 * time.Second,
|
SparkSubmitTimeout: 60 * time.Second,
|
||||||
|
UploadTTL: 168 * time.Hour,
|
||||||
LogDir: "./data/logs",
|
LogDir: "./data/logs",
|
||||||
LogLevel: "info",
|
LogLevel: "info",
|
||||||
LogFormat: "text",
|
LogFormat: "text",
|
||||||
@@ -80,6 +82,11 @@ func Load() (*Config, error) {
|
|||||||
return nil, err
|
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.LogDir = envString("LOG_DIR", cfg.LogDir)
|
||||||
cfg.LogLevel = envString("LOG_LEVEL", cfg.LogLevel)
|
cfg.LogLevel = envString("LOG_LEVEL", cfg.LogLevel)
|
||||||
cfg.LogFormat = envString("LOG_FORMAT", cfg.LogFormat)
|
cfg.LogFormat = envString("LOG_FORMAT", cfg.LogFormat)
|
||||||
@@ -110,6 +117,10 @@ func Load() (*Config, error) {
|
|||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if err := validateUploadTTL(cfg.UploadTTL); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
cfg.SQLitePath = filepath.Join(cfg.DataDir, "spark-mcp.db")
|
cfg.SQLitePath = filepath.Join(cfg.DataDir, "spark-mcp.db")
|
||||||
|
|
||||||
return cfg, nil
|
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)
|
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
|
// 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.
|
// when the variable is not already defined. Missing files are ignored.
|
||||||
func loadDotenv(path string) error {
|
func loadDotenv(path string) error {
|
||||||
@@ -280,7 +301,7 @@ func (c *Config) String() string {
|
|||||||
|
|
||||||
return fmt.Sprintf(
|
return fmt.Sprintf(
|
||||||
"ListenAddr=%s DataDir=%s SQLitePath=%s AdminTokens=%s AgentToken=%s "+
|
"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 "+
|
"LogDir=%s LogLevel=%s LogFormat=%s AnalyzerDataSkewRatio=%.1f "+
|
||||||
"AnalyzerGCPressureRatio=%.1f AnalyzerBottleneckShuffleGB=%.1f GinMode=%s",
|
"AnalyzerGCPressureRatio=%.1f AnalyzerBottleneckShuffleGB=%.1f GinMode=%s",
|
||||||
c.ListenAddr,
|
c.ListenAddr,
|
||||||
@@ -291,6 +312,7 @@ func (c *Config) String() string {
|
|||||||
c.HTTPClientTimeout,
|
c.HTTPClientTimeout,
|
||||||
c.MaxResponseBytes,
|
c.MaxResponseBytes,
|
||||||
c.SparkSubmitTimeout,
|
c.SparkSubmitTimeout,
|
||||||
|
c.UploadTTL,
|
||||||
c.LogDir,
|
c.LogDir,
|
||||||
c.LogLevel,
|
c.LogLevel,
|
||||||
c.LogFormat,
|
c.LogFormat,
|
||||||
|
|||||||
@@ -7,6 +7,7 @@ import (
|
|||||||
"spark-mcp-go/internal/analyzer"
|
"spark-mcp-go/internal/analyzer"
|
||||||
"spark-mcp-go/internal/httpclient"
|
"spark-mcp-go/internal/httpclient"
|
||||||
"spark-mcp-go/internal/storage"
|
"spark-mcp-go/internal/storage"
|
||||||
|
"spark-mcp-go/internal/uploads"
|
||||||
)
|
)
|
||||||
|
|
||||||
// Deps bundles the dependencies shared by all MCP Tool handlers.
|
// Deps bundles the dependencies shared by all MCP Tool handlers.
|
||||||
@@ -17,6 +18,7 @@ type Deps struct {
|
|||||||
HTTPClient *httpclient.Client
|
HTTPClient *httpclient.Client
|
||||||
MaxResponseBytes int64
|
MaxResponseBytes int64
|
||||||
DataDir string // upload_file writes to DataDir/uploads
|
DataDir string // upload_file writes to DataDir/uploads
|
||||||
|
UploadStore *uploads.Store
|
||||||
|
|
||||||
AnalyzerThresholds analyzer.Thresholds
|
AnalyzerThresholds analyzer.Thresholds
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -8,6 +8,7 @@ import (
|
|||||||
"net/http/httptest"
|
"net/http/httptest"
|
||||||
"os"
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
|
"regexp"
|
||||||
"strconv"
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
@@ -18,6 +19,7 @@ import (
|
|||||||
"spark-mcp-go/internal/cluster"
|
"spark-mcp-go/internal/cluster"
|
||||||
"spark-mcp-go/internal/httpclient"
|
"spark-mcp-go/internal/httpclient"
|
||||||
"spark-mcp-go/internal/storage"
|
"spark-mcp-go/internal/storage"
|
||||||
|
"spark-mcp-go/internal/uploads"
|
||||||
)
|
)
|
||||||
|
|
||||||
// testDepsWithDataDir returns dependencies backed by an in-memory DB and a
|
// testDepsWithDataDir returns dependencies backed by an in-memory DB and a
|
||||||
@@ -30,6 +32,12 @@ func testDepsWithDataDir(t *testing.T) (*Deps, *storage.ClusterRepo) {
|
|||||||
}
|
}
|
||||||
t.Cleanup(func() { _ = db.Close() })
|
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{
|
return &Deps{
|
||||||
HTTPClient: httpclient.New(httpclient.Config{
|
HTTPClient: httpclient.New(httpclient.Config{
|
||||||
Timeout: 5 * time.Second,
|
Timeout: 5 * time.Second,
|
||||||
@@ -37,7 +45,8 @@ func testDepsWithDataDir(t *testing.T) (*Deps, *storage.ClusterRepo) {
|
|||||||
}),
|
}),
|
||||||
ClusterRepo: db.Clusters(),
|
ClusterRepo: db.Clusters(),
|
||||||
MaxResponseBytes: 1 << 20,
|
MaxResponseBytes: 1 << 20,
|
||||||
DataDir: t.TempDir(),
|
DataDir: dataDir,
|
||||||
|
UploadStore: uploadStore,
|
||||||
}, db.Clusters()
|
}, db.Clusters()
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -324,15 +333,26 @@ func TestUploadFile(t *testing.T) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
payload := resultJSON(t, res)
|
payload := resultJSON(t, res)
|
||||||
if payload["path"] != "uploads/hello.txt" {
|
fileID, ok := payload["file_id"].(string)
|
||||||
t.Errorf("path=%v, want uploads/hello.txt", payload["path"])
|
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")) {
|
if payload["size"] != float64(len("hello world")) {
|
||||||
t.Errorf("size=%v, want %d", payload["size"], 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(path)
|
||||||
got, err := os.ReadFile(final)
|
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatalf("read uploaded file: %v", err)
|
t.Fatalf("read uploaded file: %v", err)
|
||||||
}
|
}
|
||||||
@@ -340,7 +360,7 @@ func TestUploadFile(t *testing.T) {
|
|||||||
t.Errorf("content=%q, want %q", got, "hello world")
|
t.Errorf("content=%q, want %q", got, "hello world")
|
||||||
}
|
}
|
||||||
|
|
||||||
info, err := os.Stat(final)
|
info, err := os.Stat(path)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatalf("stat uploaded file: %v", err)
|
t.Fatalf("stat uploaded file: %v", err)
|
||||||
}
|
}
|
||||||
@@ -393,9 +413,19 @@ func TestUploadFile_Base64(t *testing.T) {
|
|||||||
t.Fatalf("handler error: %v", err)
|
t.Fatalf("handler error: %v", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
_ = resultJSON(t, res)
|
payload := resultJSON(t, res)
|
||||||
final := filepath.Join(deps.DataDir, "uploads", "hello.bin")
|
path, ok := payload["path"].(string)
|
||||||
got, err := os.ReadFile(final)
|
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 {
|
if err != nil {
|
||||||
t.Fatalf("read uploaded file: %v", err)
|
t.Fatalf("read uploaded file: %v", err)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -7,6 +7,7 @@ import (
|
|||||||
"github.com/mark3labs/mcp-go/mcp"
|
"github.com/mark3labs/mcp-go/mcp"
|
||||||
|
|
||||||
"spark-mcp-go/internal/executor"
|
"spark-mcp-go/internal/executor"
|
||||||
|
"spark-mcp-go/internal/uploads"
|
||||||
)
|
)
|
||||||
|
|
||||||
const SparkSubmitName = "spark_submit"
|
const SparkSubmitName = "spark_submit"
|
||||||
@@ -14,7 +15,7 @@ const SparkSubmitName = "spark_submit"
|
|||||||
// NewSparkSubmitTool returns the schema for the spark_submit MCP Tool.
|
// NewSparkSubmitTool returns the schema for the spark_submit MCP Tool.
|
||||||
func NewSparkSubmitTool() mcp.Tool {
|
func NewSparkSubmitTool() mcp.Tool {
|
||||||
return mcp.NewTool(SparkSubmitName,
|
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 using the cluster's local spark-submit binary. The `args` array is appended to the cluster's default_submit_args; pass one argv element per CLI flag. If the script to run was just uploaded via upload_file, put the `path` field from that response as one of the args — do NOT construct your own path. Cluster-local paths (e.g. /opt/spark/examples/pi.py) are also accepted. 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.WithString("cluster_id",
|
||||||
mcp.Required(),
|
mcp.Required(),
|
||||||
mcp.Description("ID of the configured cluster (from list_clusters)"),
|
mcp.Description("ID of the configured cluster (from list_clusters)"),
|
||||||
@@ -63,7 +64,7 @@ func (d *Deps) SparkSubmitHandler(ctx context.Context, req mcp.CallToolRequest)
|
|||||||
}
|
}
|
||||||
|
|
||||||
fullArgs := append([]string{}, cl.DefaultSubmitArgs...)
|
fullArgs := append([]string{}, cl.DefaultSubmitArgs...)
|
||||||
fullArgs = append(fullArgs, userArgs...)
|
fullArgs = append(fullArgs, translateArgs(userArgs, d.UploadStore)...)
|
||||||
|
|
||||||
result, err := executor.Run(ctx, executor.SparkSubmitOpts{
|
result, err := executor.Run(ctx, executor.SparkSubmitOpts{
|
||||||
Binary: cl.SparkSubmitExecuteBin,
|
Binary: cl.SparkSubmitExecuteBin,
|
||||||
@@ -98,3 +99,23 @@ func (d *Deps) SparkSubmitHandler(ctx context.Context, req mcp.CallToolRequest)
|
|||||||
})
|
})
|
||||||
return textResult(encodeJSON(result)), nil
|
return textResult(encodeJSON(result)), nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// translateArgs canonicalizes server-minted upload paths and leaves
|
||||||
|
// cluster-local paths untouched. It is a code-level guardrail, not a
|
||||||
|
// validator: non-minted absolute paths are intentionally accepted.
|
||||||
|
func translateArgs(userArgs []string, store *uploads.Store) []string {
|
||||||
|
if store == nil {
|
||||||
|
return userArgs
|
||||||
|
}
|
||||||
|
out := make([]string, len(userArgs))
|
||||||
|
for i, arg := range userArgs {
|
||||||
|
if fileID, err := store.Validate(arg); err == nil {
|
||||||
|
if canonical, err := store.AbsPath(fileID); err == nil {
|
||||||
|
out[i] = canonical
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
}
|
||||||
|
out[i] = arg
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,120 @@
|
|||||||
|
package tools
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"spark-mcp-go/internal/cluster"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestSparkSubmit_PathPassthrough(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\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",
|
||||||
|
"args": []any{absPath, "--class", "Main"},
|
||||||
|
})
|
||||||
|
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{absPath, "--class", "Main"}
|
||||||
|
if len(lines) != len(want) {
|
||||||
|
t.Fatalf("echoed lines=%v, want %v", lines, want)
|
||||||
|
}
|
||||||
|
for i, l := range lines {
|
||||||
|
if l != want[i] {
|
||||||
|
t.Errorf("line[%d]=%q, want %q", i, l, want[i])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Sanity: the canonical path still contains the minted file ID.
|
||||||
|
if !strings.HasSuffix(absPath, "/"+fileID) {
|
||||||
|
t.Errorf("absPath=%q does not end with fileID %q", absPath, fileID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSparkSubmit_ClusterLocalPathPassthrough(t *testing.T) {
|
||||||
|
deps, repo := testDepsWithDataDir(t)
|
||||||
|
deps.SparkSubmitTimeout = 5 * time.Second
|
||||||
|
|
||||||
|
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\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,
|
||||||
|
})
|
||||||
|
|
||||||
|
clusterLocalPath := "/usr/bin/env"
|
||||||
|
req := newToolRequest(SparkSubmitName, map[string]any{
|
||||||
|
"cluster_id": "cluster-echo",
|
||||||
|
"args": []any{clusterLocalPath, "--class", "Main"},
|
||||||
|
})
|
||||||
|
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{clusterLocalPath, "--class", "Main"}
|
||||||
|
if len(lines) != len(want) {
|
||||||
|
t.Fatalf("echoed lines=%v, want %v", lines, want)
|
||||||
|
}
|
||||||
|
for i, l := range lines {
|
||||||
|
if l != want[i] {
|
||||||
|
t.Errorf("line[%d]=%q, want %q", i, l, want[i])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -4,7 +4,6 @@ import (
|
|||||||
"context"
|
"context"
|
||||||
"encoding/base64"
|
"encoding/base64"
|
||||||
"fmt"
|
"fmt"
|
||||||
"os"
|
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
"regexp"
|
"regexp"
|
||||||
|
|
||||||
@@ -19,7 +18,7 @@ var filenameRegex = regexp.MustCompile(`^[a-zA-Z0-9._-]+$`)
|
|||||||
// NewUploadFileTool returns the schema for the upload_file MCP Tool.
|
// NewUploadFileTool returns the schema for the upload_file MCP Tool.
|
||||||
func NewUploadFileTool() mcp.Tool {
|
func NewUploadFileTool() mcp.Tool {
|
||||||
return mcp.NewTool(UploadFileName,
|
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 one of the args in spark_submit (e.g. args=[<path>, '--class', 'Main']). The path is server-side; the agent does not need to construct one."),
|
||||||
mcp.WithString("filename",
|
mcp.WithString("filename",
|
||||||
mcp.Required(),
|
mcp.Required(),
|
||||||
mcp.Description("Plain file name without path separators (1-128 chars, [a-zA-Z0-9._-])"),
|
mcp.Description("Plain file name without path separators (1-128 chars, [a-zA-Z0-9._-])"),
|
||||||
@@ -60,6 +59,10 @@ func (d *Deps) UploadFileHandler(ctx context.Context, req mcp.CallToolRequest) (
|
|||||||
return errResult("upload_file: " + err.Error()), nil
|
return errResult("upload_file: " + err.Error()), nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if d.UploadStore == nil {
|
||||||
|
return errResult("upload_file: upload store not configured"), nil
|
||||||
|
}
|
||||||
|
|
||||||
var data []byte
|
var data []byte
|
||||||
switch encoding {
|
switch encoding {
|
||||||
case "base64":
|
case "base64":
|
||||||
@@ -72,17 +75,17 @@ func (d *Deps) UploadFileHandler(ctx context.Context, req mcp.CallToolRequest) (
|
|||||||
data = []byte(content)
|
data = []byte(content)
|
||||||
}
|
}
|
||||||
|
|
||||||
final := filepath.Join(d.DataDir, "uploads", filename)
|
fileID, _, size, sha256Hex, absPath, err := d.UploadStore.Save(data, filename)
|
||||||
if err := os.MkdirAll(filepath.Dir(final), 0o750); err != nil {
|
if err != nil {
|
||||||
return errResult("upload_file: create uploads dir: " + err.Error()), 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
|
|
||||||
}
|
}
|
||||||
|
|
||||||
result := map[string]any{
|
result := map[string]any{
|
||||||
"path": fmt.Sprintf("uploads/%s", filename),
|
"file_id": fileID,
|
||||||
"size": len(data),
|
"path": absPath,
|
||||||
|
"name": filename,
|
||||||
|
"size": size,
|
||||||
|
"sha256": sha256Hex,
|
||||||
}
|
}
|
||||||
return textResult(encodeJSON(result)), nil
|
return textResult(encodeJSON(result)), nil
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,253 @@
|
|||||||
|
// 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 (
|
||||||
|
"crypto/rand"
|
||||||
|
"crypto/sha256"
|
||||||
|
"encoding/hex"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"regexp"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
var fileIDRegex = regexp.MustCompile(`^[0-9a-f]{32}$`)
|
||||||
|
|
||||||
|
// Store is a local file store rooted at Root.
|
||||||
|
type Store struct {
|
||||||
|
Root string
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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 nil, fmt.Errorf("uploads: root must not be empty")
|
||||||
|
}
|
||||||
|
absRoot, err := filepath.Abs(root)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("uploads: resolve root: %w", err)
|
||||||
|
}
|
||||||
|
if !filepath.IsAbs(absRoot) {
|
||||||
|
return nil, fmt.Errorf("uploads: root %q is not an absolute path", absRoot)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := os.MkdirAll(absRoot, 0o750); err != nil {
|
||||||
|
return nil, fmt.Errorf("uploads: create root: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
info, err := os.Stat(absRoot)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("uploads: stat root: %w", err)
|
||||||
|
}
|
||||||
|
if !info.IsDir() {
|
||||||
|
return nil, 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)
|
||||||
|
}
|
||||||
|
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
|
||||||
|
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(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 {
|
||||||
|
dataPath := filepath.Join(s.Root, fileID)
|
||||||
|
metaPath := dataPath + ".meta.json"
|
||||||
|
|
||||||
|
if it.data {
|
||||||
|
expired := false
|
||||||
|
if it.meta {
|
||||||
|
uploadedAt, parseErr := readUploadedAt(metaPath)
|
||||||
|
if parseErr == nil && uploadedAt.Before(cutoff) {
|
||||||
|
expired = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !expired && !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
|
||||||
|
}
|
||||||
|
|
||||||
|
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) {
|
||||||
|
data, err := os.ReadFile(path)
|
||||||
|
if err != nil {
|
||||||
|
return time.Time{}, err
|
||||||
|
}
|
||||||
|
var sc sidecar
|
||||||
|
if err := json.Unmarshal(data, &sc); err != nil {
|
||||||
|
return time.Time{}, err
|
||||||
|
}
|
||||||
|
if sc.UploadedAt.IsZero() {
|
||||||
|
return time.Time{}, fmt.Errorf("uploads: missing uploaded_at")
|
||||||
|
}
|
||||||
|
return sc.UploadedAt, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,258 @@
|
|||||||
|
package uploads
|
||||||
|
|
||||||
|
import (
|
||||||
|
"crypto/sha256"
|
||||||
|
"encoding/hex"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"regexp"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
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(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(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
|
||||||
|
}
|
||||||
@@ -11,6 +11,7 @@ import (
|
|||||||
"net/http"
|
"net/http"
|
||||||
"os"
|
"os"
|
||||||
"os/signal"
|
"os/signal"
|
||||||
|
"path/filepath"
|
||||||
"syscall"
|
"syscall"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
@@ -26,6 +27,7 @@ import (
|
|||||||
"spark-mcp-go/internal/mcp/tools"
|
"spark-mcp-go/internal/mcp/tools"
|
||||||
"spark-mcp-go/internal/middleware"
|
"spark-mcp-go/internal/middleware"
|
||||||
"spark-mcp-go/internal/storage"
|
"spark-mcp-go/internal/storage"
|
||||||
|
"spark-mcp-go/internal/uploads"
|
||||||
)
|
)
|
||||||
|
|
||||||
func main() {
|
func main() {
|
||||||
@@ -65,6 +67,17 @@ func run() error {
|
|||||||
defer db.Close()
|
defer db.Close()
|
||||||
logger.Info("storage.open", "path", cfg.SQLitePath)
|
logger.Info("storage.open", "path", cfg.SQLitePath)
|
||||||
|
|
||||||
|
uploadStore, err := uploads.New(filepath.Join(cfg.DataDir, "uploads"))
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
n, err := uploadStore.Sweep(cfg.UploadTTL)
|
||||||
|
if err != nil {
|
||||||
|
logger.Error("uploads.sweep", "err", err)
|
||||||
|
} else {
|
||||||
|
logger.Info("uploads.sweep", "deleted", n)
|
||||||
|
}
|
||||||
|
|
||||||
gin.SetMode(cfg.GinMode)
|
gin.SetMode(cfg.GinMode)
|
||||||
r := gin.New()
|
r := gin.New()
|
||||||
r.Use(gin.Recovery())
|
r.Use(gin.Recovery())
|
||||||
@@ -85,6 +98,7 @@ func run() error {
|
|||||||
}),
|
}),
|
||||||
MaxResponseBytes: cfg.MaxResponseBytes,
|
MaxResponseBytes: cfg.MaxResponseBytes,
|
||||||
DataDir: cfg.DataDir,
|
DataDir: cfg.DataDir,
|
||||||
|
UploadStore: uploadStore,
|
||||||
AnalyzerThresholds: analyzer.Thresholds{
|
AnalyzerThresholds: analyzer.Thresholds{
|
||||||
DataSkewRatio: cfg.AnalyzerDataSkewRatio,
|
DataSkewRatio: cfg.AnalyzerDataSkewRatio,
|
||||||
GCPressureRatio: cfg.AnalyzerGCPressureRatio,
|
GCPressureRatio: cfg.AnalyzerGCPressureRatio,
|
||||||
|
|||||||
Reference in New Issue
Block a user