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:
@@ -3,14 +3,12 @@ package tools
|
|||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"fmt"
|
"fmt"
|
||||||
"log/slog"
|
"sort"
|
||||||
"path/filepath"
|
"strconv"
|
||||||
"strings"
|
|
||||||
|
|
||||||
"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"
|
||||||
@@ -18,15 +16,47 @@ 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. 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 accepted but will be logged as a warning. To avoid the warning, always use the `path` field returned by upload_file. Returns {app_id, exit_code, stdout_tail, stderr_tail, duration_ms}. 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.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)"),
|
||||||
),
|
),
|
||||||
// mcp-go v0.56.0 uses PropertyOption for array item schema.
|
mcp.WithString("master",
|
||||||
mcp.WithArray("args",
|
mcp.Required(),
|
||||||
mcp.Description("Spark-submit CLI args (after default_submit_args). Each element becomes one argv entry — no shell interpretation."),
|
mcp.Description("Spark master URL, e.g. yarn, k8s://https://..."),
|
||||||
mcp.Items(map[string]any{"type": "string"}),
|
),
|
||||||
|
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."),
|
||||||
|
),
|
||||||
|
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"}),
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
@@ -37,26 +67,65 @@ func (d *Deps) SparkSubmitHandler(ctx context.Context, req mcp.CallToolRequest)
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return errResult("spark_submit: " + err.Error()), nil
|
return errResult("spark_submit: " + err.Error()), nil
|
||||||
}
|
}
|
||||||
|
master, err := req.RequireString("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
|
||||||
|
}
|
||||||
|
queue, err := req.RequireString("queue")
|
||||||
|
if err != nil {
|
||||||
|
return errResult("spark_submit: " + err.Error()), nil
|
||||||
|
}
|
||||||
|
executorMemory, err := req.RequireString("executor_memory")
|
||||||
|
if err != nil {
|
||||||
|
return errResult("spark_submit: " + err.Error()), nil
|
||||||
|
}
|
||||||
|
|
||||||
args := req.GetArguments()
|
args := req.GetArguments()
|
||||||
var userArgs []string
|
|
||||||
if rawArgs, ok := args["args"]; ok && rawArgs != nil {
|
// deploy_mode is required for forward compatibility but the builder does not
|
||||||
arr, ok := rawArgs.([]any)
|
// emit --deploy-mode, matching the Python reference's actual behavior. If the
|
||||||
if !ok {
|
// Go side needs --deploy-mode later, add it here and document the divergence.
|
||||||
return errResult("spark_submit: args is not an array"), nil
|
if deployMode != "client" && deployMode != "cluster" {
|
||||||
|
return errResult("spark_submit: deploy_mode must be client or cluster"), nil
|
||||||
}
|
}
|
||||||
for i, item := range arr {
|
|
||||||
s, ok := item.(string)
|
sparkConf, err := parseStringMap(args, "spark_conf")
|
||||||
if !ok {
|
if err != nil {
|
||||||
return errResult(fmt.Sprintf("spark_submit: args[%d] is not a string", i)), nil
|
return errResult("spark_submit: " + err.Error()), nil
|
||||||
}
|
}
|
||||||
userArgs = append(userArgs, s)
|
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{
|
callLog := startToolCall(ctx, d.Logger, SparkSubmitName, map[string]any{
|
||||||
"cluster_id": clusterID,
|
"cluster_id": clusterID,
|
||||||
"args": userArgs,
|
"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()
|
defer callLog.End()
|
||||||
|
|
||||||
@@ -66,12 +135,34 @@ func (d *Deps) SparkSubmitHandler(ctx context.Context, req mcp.CallToolRequest)
|
|||||||
return errResult("spark_submit: cluster " + clusterID + ": " + err.Error()), nil
|
return errResult("spark_submit: cluster " + clusterID + ": " + err.Error()), nil
|
||||||
}
|
}
|
||||||
|
|
||||||
fullArgs := append([]string{}, cl.DefaultSubmitArgs...)
|
binary := cl.SparkSubmitExecuteBin
|
||||||
fullArgs = append(fullArgs, translateArgs(userArgs, d.UploadStore, d.Logger)...)
|
|
||||||
|
// Build argv inline, mirroring the Python reference's output order exactly.
|
||||||
|
// 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.
|
||||||
|
cmd := []string{binary}
|
||||||
|
cmd = append(cmd, "--master", master)
|
||||||
|
cmd = append(cmd, "--queue", queue)
|
||||||
|
cmd = append(cmd, "--executor-memory", executorMemory)
|
||||||
|
cmd = append(cmd, "--executor-cores", strconv.Itoa(executorCores))
|
||||||
|
cmd = append(cmd, "--num-executors", strconv.Itoa(numExecutors))
|
||||||
|
|
||||||
|
for _, k := range sortedStringKeys(sparkConf) {
|
||||||
|
cmd = append(cmd, "--conf", k+"="+sparkConf[k])
|
||||||
|
}
|
||||||
|
for _, k := range sortedStringKeys(extraArgs) {
|
||||||
|
cmd = append(cmd, "--"+k, extraArgs[k])
|
||||||
|
}
|
||||||
|
cmd = append(cmd, scriptPath)
|
||||||
|
|
||||||
|
if d.Logger != nil {
|
||||||
|
d.Logger.Debug("spark_submit.built", "cluster_id", clusterID, "argv", cmd)
|
||||||
|
}
|
||||||
|
|
||||||
result, err := executor.Run(ctx, executor.SparkSubmitOpts{
|
result, err := executor.Run(ctx, executor.SparkSubmitOpts{
|
||||||
Binary: cl.SparkSubmitExecuteBin,
|
Binary: binary,
|
||||||
Args: fullArgs,
|
Args: cmd,
|
||||||
Timeout: d.SparkSubmitTimeout,
|
Timeout: d.SparkSubmitTimeout,
|
||||||
})
|
})
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -103,52 +194,57 @@ 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
|
// requireNonNegativeInt extracts an integer parameter from the request and
|
||||||
// cluster-local paths untouched. It is a code-level guardrail, not a
|
// rejects negative or non-integer values.
|
||||||
// validator: non-minted absolute paths are intentionally accepted.
|
func requireNonNegativeInt(args map[string]any, key string) (int, error) {
|
||||||
func translateArgs(userArgs []string, store *uploads.Store, logger *slog.Logger) []string {
|
v, ok := args[key]
|
||||||
if store == nil {
|
if !ok || v == nil {
|
||||||
return userArgs
|
return 0, fmt.Errorf("%s is required", key)
|
||||||
}
|
}
|
||||||
out := make([]string, len(userArgs))
|
switch n := v.(type) {
|
||||||
for i, arg := range userArgs {
|
case int:
|
||||||
pathPart, isPath := pathLikeArg(arg)
|
if n < 0 {
|
||||||
if isPath {
|
return 0, fmt.Errorf("%s must be a non-negative integer", key)
|
||||||
if fileID, err := store.Validate(pathPart); err == nil {
|
|
||||||
if canonical, err := store.AbsPath(fileID); err == nil {
|
|
||||||
out[i] = canonical
|
|
||||||
continue
|
|
||||||
}
|
}
|
||||||
} else if logger != nil && filepath.IsAbs(pathPart) {
|
return n, nil
|
||||||
// Absolute path that this server did not mint. Still allowed,
|
case float64:
|
||||||
// but warn so operators can spot LLM-invented remote paths.
|
ni := int(n)
|
||||||
logger.Warn("spark_submit.unminted_path",
|
if n < 0 || float64(ni) != n {
|
||||||
"path", pathPart,
|
return 0, fmt.Errorf("%s must be a non-negative integer", key)
|
||||||
"note", "not produced by upload_file on this server; cluster-local paths are valid, but if this came from a remote agent the path is unreachable")
|
|
||||||
}
|
}
|
||||||
|
return ni, nil
|
||||||
|
default:
|
||||||
|
return 0, fmt.Errorf("%s must be a number", key)
|
||||||
}
|
}
|
||||||
out[i] = arg
|
|
||||||
}
|
|
||||||
return out
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// pathLikeArg extracts an absolute path from a spark-submit argv element.
|
// parseStringMap extracts an optional object whose values are strings.
|
||||||
// It handles plain absolute paths ("/opt/spark/pi.py") and flag-style values
|
func parseStringMap(args map[string]any, key string) (map[string]string, error) {
|
||||||
// ("--conf=/x/y"). Non-path args like "--class" or "Main" return ok=false.
|
out := make(map[string]string)
|
||||||
func pathLikeArg(arg string) (string, bool) {
|
v, ok := args[key]
|
||||||
if arg == "" {
|
if !ok || v == nil {
|
||||||
return "", false
|
return out, nil
|
||||||
}
|
}
|
||||||
if strings.HasPrefix(arg, "/") {
|
m, ok := v.(map[string]any)
|
||||||
return arg, true
|
if !ok {
|
||||||
|
return nil, fmt.Errorf("%s must be an object", key)
|
||||||
}
|
}
|
||||||
if strings.HasPrefix(arg, "-") {
|
for k, val := range m {
|
||||||
if eq := strings.Index(arg, "="); eq > 0 && eq < len(arg)-1 {
|
s, ok := val.(string)
|
||||||
val := arg[eq+1:]
|
if !ok {
|
||||||
if strings.HasPrefix(val, "/") {
|
return nil, fmt.Errorf("%s[%q] must be a string", key, k)
|
||||||
return val, true
|
|
||||||
}
|
}
|
||||||
|
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)
|
||||||
}
|
}
|
||||||
return "", false
|
sort.Strings(keys)
|
||||||
|
return keys
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,20 +1,19 @@
|
|||||||
package tools
|
package tools
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"bytes"
|
|
||||||
"context"
|
"context"
|
||||||
"encoding/json"
|
|
||||||
"log/slog"
|
|
||||||
"os"
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
"strings"
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"github.com/mark3labs/mcp-go/mcp"
|
||||||
|
|
||||||
"spark-mcp-go/internal/cluster"
|
"spark-mcp-go/internal/cluster"
|
||||||
)
|
)
|
||||||
|
|
||||||
func TestSparkSubmit_PathPassthrough(t *testing.T) {
|
func TestSparkSubmit_StructuredCommand(t *testing.T) {
|
||||||
deps, repo := testDepsWithDataDir(t)
|
deps, repo := testDepsWithDataDir(t)
|
||||||
deps.SparkSubmitTimeout = 5 * time.Second
|
deps.SparkSubmitTimeout = 5 * time.Second
|
||||||
store := deps.UploadStore
|
store := deps.UploadStore
|
||||||
@@ -42,7 +41,15 @@ func TestSparkSubmit_PathPassthrough(t *testing.T) {
|
|||||||
|
|
||||||
req := newToolRequest(SparkSubmitName, map[string]any{
|
req := newToolRequest(SparkSubmitName, map[string]any{
|
||||||
"cluster_id": "cluster-echo",
|
"cluster_id": "cluster-echo",
|
||||||
"args": []any{absPath, "--class", "Main"},
|
"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)
|
res, err := deps.SparkSubmitHandler(context.Background(), req)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -57,9 +64,21 @@ func TestSparkSubmit_PathPassthrough(t *testing.T) {
|
|||||||
t.Fatalf("read echoed args: %v", err)
|
t.Fatalf("read echoed args: %v", err)
|
||||||
}
|
}
|
||||||
lines := strings.Split(strings.TrimSpace(string(got)), "\n")
|
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) {
|
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 {
|
for i, l := range lines {
|
||||||
if l != want[i] {
|
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) {
|
if !strings.HasSuffix(absPath, "/"+fileID) {
|
||||||
t.Errorf("absPath=%q does not end with fileID %q", 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)
|
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{
|
req := newToolRequest(SparkSubmitName, map[string]any{
|
||||||
"cluster_id": "cluster-echo",
|
"cluster_id": "cluster-echo",
|
||||||
"args": []any{clusterLocalPath, "--class", "Main"},
|
"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)
|
res, err := deps.SparkSubmitHandler(context.Background(), req)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatalf("handler error: %v", err)
|
t.Fatalf("handler error: %v", err)
|
||||||
}
|
}
|
||||||
if res.IsError {
|
if !res.IsError {
|
||||||
t.Fatalf("unexpected error result: %v", res.Content)
|
t.Fatalf("expected error result, got: %v", res.Content)
|
||||||
}
|
}
|
||||||
|
text, ok := mcp.AsTextContent(res.Content[0])
|
||||||
got, err := os.ReadFile(echoed)
|
if !ok {
|
||||||
if err != nil {
|
t.Fatalf("content is not text: %T", res.Content[0])
|
||||||
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])
|
|
||||||
}
|
}
|
||||||
|
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