tools: 仿照 Python build_spark_submit_command 将 spark_submit 改为结构化参数
将 spark_submit 的自由 args[] 替换为与 Python 服务一致的 schema:cluster_id、 master、deploy_mode、script_path、queue、executor_memory、executor_cores、 num_executors、spark_conf、extra_args。命令行顺序严格对齐 Python 实现: binary → --master → --queue → --executor-memory → --executor-cores → --num-executors → 按 key 排序的 --conf key=value → 按 key 排序的 --flag value → script_path 永远在最后。 deploy_mode 是 Python 参考函数的入参,但该函数实际并不输出 --deploy-mode; Go 端保留这个字段并校验其值,但同样不发射它,以保持行为一致。如果后续需要 --deploy-mode,再在此处添加并记录差异。 cluster.DefaultSubmitArgs 不再被拼接到 argv 中;该字段仍保留用于 admin API 向后兼容,运行时移除它的工作是后续独立的 follow-up。 删除了 translateArgs、pathLikeArg 以及 spark_submit.unminted_path 警告逻辑 (本次提交取代 2570713),因为自由参数已不存在,警告无从触发。新增 parseStringMap、sortedStringKeys 和 requireNonNegativeInt 辅助函数; requireNonNegativeInt 同时接受 int 与 float64,以兼容测试直接构造的整数 字面量和 JSON 解码后的 float64。 测试更新: - TestSparkSubmit_StructuredCommand:验证完整 argv 顺序,并断言 script_path 是最后一项。 - TestSparkSubmit_MissingRequiredField:校验缺少 master 时返回错误。 - TestSparkSubmit_BadSparkConfValue:校验 spark_conf 的值不是字符串时报错。 Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -1,20 +1,19 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"log/slog"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
|
||||
"spark-mcp-go/internal/cluster"
|
||||
)
|
||||
|
||||
func TestSparkSubmit_PathPassthrough(t *testing.T) {
|
||||
func TestSparkSubmit_StructuredCommand(t *testing.T) {
|
||||
deps, repo := testDepsWithDataDir(t)
|
||||
deps.SparkSubmitTimeout = 5 * time.Second
|
||||
store := deps.UploadStore
|
||||
@@ -41,8 +40,16 @@ func TestSparkSubmit_PathPassthrough(t *testing.T) {
|
||||
})
|
||||
|
||||
req := newToolRequest(SparkSubmitName, map[string]any{
|
||||
"cluster_id": "cluster-echo",
|
||||
"args": []any{absPath, "--class", "Main"},
|
||||
"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 {
|
||||
@@ -57,9 +64,21 @@ func TestSparkSubmit_PathPassthrough(t *testing.T) {
|
||||
t.Fatalf("read echoed args: %v", err)
|
||||
}
|
||||
lines := strings.Split(strings.TrimSpace(string(got)), "\n")
|
||||
want := []string{absPath, "--class", "Main"}
|
||||
|
||||
want := []string{
|
||||
echoScript,
|
||||
"--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("echoed lines=%v, want %v", lines, want)
|
||||
t.Fatalf("lines=%v\nwant=%v", lines, want)
|
||||
}
|
||||
for i, l := range lines {
|
||||
if l != want[i] {
|
||||
@@ -67,148 +86,69 @@ func TestSparkSubmit_PathPassthrough(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// Sanity: the canonical path still contains the minted file ID.
|
||||
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 TestTranslateArgs_Logging(t *testing.T) {
|
||||
func TestSparkSubmit_MissingRequiredField(t *testing.T) {
|
||||
deps, _ := testDepsWithDataDir(t)
|
||||
|
||||
t.Run("minted_path_no_warn", func(t *testing.T) {
|
||||
store := deps.UploadStore
|
||||
fileID, _, _, _, absPath, err := store.Save([]byte("print('hi')\n"), "hello.py")
|
||||
if err != nil {
|
||||
t.Fatalf("save upload: %v", err)
|
||||
}
|
||||
|
||||
logger, buf := capturingLogger(t)
|
||||
out := translateArgs([]string{absPath, "--class", "Main"}, deps.UploadStore, logger)
|
||||
if out[0] != absPath {
|
||||
t.Errorf("out[0]=%q, want %q", out[0], absPath)
|
||||
}
|
||||
|
||||
for _, rec := range parseLogRecords(t, buf) {
|
||||
if rec["msg"] == "spark_submit.unminted_path" {
|
||||
t.Errorf("unexpected warn for minted path: %v", rec)
|
||||
}
|
||||
}
|
||||
|
||||
if !strings.HasSuffix(absPath, "/"+fileID) {
|
||||
t.Errorf("absPath=%q does not end with fileID %q", absPath, fileID)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("cluster_local_path_warns", func(t *testing.T) {
|
||||
logger, buf := capturingLogger(t)
|
||||
out := translateArgs([]string{"/opt/spark/examples/pi.py"}, deps.UploadStore, logger)
|
||||
if len(out) != 1 || out[0] != "/opt/spark/examples/pi.py" {
|
||||
t.Errorf("out=%v, want [/opt/spark/examples/pi.py]", out)
|
||||
}
|
||||
|
||||
var warns int
|
||||
for _, rec := range parseLogRecords(t, buf) {
|
||||
if rec["msg"] == "spark_submit.unminted_path" {
|
||||
warns++
|
||||
if rec["path"] != "/opt/spark/examples/pi.py" {
|
||||
t.Errorf("warn path=%v, want /opt/spark/examples/pi.py", rec["path"])
|
||||
}
|
||||
}
|
||||
}
|
||||
if warns != 1 {
|
||||
t.Errorf("warns=%d, want 1; log=%s", warns, buf.String())
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("non_path_args_never_warn", func(t *testing.T) {
|
||||
logger, buf := capturingLogger(t)
|
||||
out := translateArgs([]string{"--class", "Main", "--master", "yarn"}, deps.UploadStore, logger)
|
||||
want := []string{"--class", "Main", "--master", "yarn"}
|
||||
if len(out) != len(want) {
|
||||
t.Fatalf("out=%v, want %v", out, want)
|
||||
}
|
||||
for i := range out {
|
||||
if out[i] != want[i] {
|
||||
t.Errorf("out[%d]=%q, want %q", i, out[i], want[i])
|
||||
}
|
||||
}
|
||||
|
||||
for _, rec := range parseLogRecords(t, buf) {
|
||||
if rec["msg"] == "spark_submit.unminted_path" {
|
||||
t.Errorf("unexpected warn for non-path args: %v", rec)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func capturingLogger(t *testing.T) (*slog.Logger, *bytes.Buffer) {
|
||||
t.Helper()
|
||||
buf := &bytes.Buffer{}
|
||||
return slog.New(slog.NewJSONHandler(buf, &slog.HandlerOptions{Level: slog.LevelDebug})), buf
|
||||
}
|
||||
|
||||
func parseLogRecords(t *testing.T, buf *bytes.Buffer) []map[string]any {
|
||||
t.Helper()
|
||||
var recs []map[string]any
|
||||
for _, line := range strings.Split(strings.TrimSpace(buf.String()), "\n") {
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
var rec map[string]any
|
||||
if err := json.Unmarshal([]byte(line), &rec); err != nil {
|
||||
t.Fatalf("invalid log JSON: %q (err=%v)", line, err)
|
||||
}
|
||||
recs = append(recs, rec)
|
||||
}
|
||||
return recs
|
||||
}
|
||||
|
||||
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"},
|
||||
"cluster_id": "cluster-echo",
|
||||
"deploy_mode": "cluster",
|
||||
"script_path": "/tmp/script.py",
|
||||
"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)
|
||||
if !res.IsError {
|
||||
t.Fatalf("expected error result, got: %v", res.Content)
|
||||
}
|
||||
|
||||
got, err := os.ReadFile(echoed)
|
||||
if err != nil {
|
||||
t.Fatalf("read echoed args: %v", err)
|
||||
text, ok := mcp.AsTextContent(res.Content[0])
|
||||
if !ok {
|
||||
t.Fatalf("content is not text: %T", res.Content[0])
|
||||
}
|
||||
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])
|
||||
}
|
||||
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)
|
||||
|
||||
req := newToolRequest(SparkSubmitName, map[string]any{
|
||||
"cluster_id": "cluster-echo",
|
||||
"master": "yarn",
|
||||
"deploy_mode": "cluster",
|
||||
"script_path": "/tmp/script.py",
|
||||
"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)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user