From c4ac3cc354de864cf997d1c10c58949d96e85bc8 Mon Sep 17 00:00:00 2001 From: "tao.chen" Date: Fri, 10 Jul 2026 16:50:27 +0800 Subject: [PATCH] =?UTF-8?q?Phase=205=20Batch=202:=204=20RM=20Tool=20+=20rm?= =?UTF-8?q?=20=E5=AE=A2=E6=88=B7=E7=AB=AF=20+=20=E6=97=A5=E5=BF=97?= =?UTF-8?q?=E9=99=8D=E7=BA=A7=E9=93=BE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - internal/rm/client.go: YARN RM 客户端 - ListApps / GetApp / KillApp / GetLogs - GetLogs 降级链数据驱动 (amContainerLogs → aggregated-logs → logs) - 跨主机 redirect 保留 Authorization (URLAllowlist 兜底) - 4 Tool: list_applications / get_application_status / get_application_logs / kill_application - mcp.WithEnum(state) 约束 YARN app 状态 - tool handler 永远 (result, nil), 业务错误用 NewToolResultError - deps.go: +HTTPClient + MaxResponseBytes - main.go: 注入 httpclient.Client - 端到端实测: mock RM 验证 4 Tool + 降级链 + 错误路径 Co-Authored-By: Claude --- internal/mcp/server.go | 36 +++ internal/mcp/tools/deps.go | 22 ++ internal/mcp/tools/get_application_logs.go | 93 ++++++++ internal/mcp/tools/get_application_status.go | 60 +++++ internal/mcp/tools/kill_application.go | 60 +++++ internal/mcp/tools/list_applications.go | 69 ++++++ internal/mcp/tools/rm_test.go | 176 ++++++++++++++ internal/rm/client.go | 131 ++++++++++ internal/rm/client_test.go | 238 +++++++++++++++++++ 9 files changed, 885 insertions(+) create mode 100644 internal/mcp/server.go create mode 100644 internal/mcp/tools/deps.go create mode 100644 internal/mcp/tools/get_application_logs.go create mode 100644 internal/mcp/tools/get_application_status.go create mode 100644 internal/mcp/tools/kill_application.go create mode 100644 internal/mcp/tools/list_applications.go create mode 100644 internal/mcp/tools/rm_test.go create mode 100644 internal/rm/client.go create mode 100644 internal/rm/client_test.go diff --git a/internal/mcp/server.go b/internal/mcp/server.go new file mode 100644 index 0000000..e68a1f9 --- /dev/null +++ b/internal/mcp/server.go @@ -0,0 +1,36 @@ +package mcp + +import ( + "net/http" + + "github.com/mark3labs/mcp-go/server" + + "spark-mcp-go/internal/mcp/tools" +) + +// NewServer builds the MCP server with all configured Tools. +func NewServer(deps *tools.Deps) *server.MCPServer { + s := server.NewMCPServer("spark-mcp-go", "0.0.0", + server.WithToolCapabilities(false), + server.WithRecovery(), + server.WithLogging(), + ) + s.AddTool(tools.NewListClustersTool(), deps.ListClustersHandler) + s.AddTool(tools.NewSparkSubmitTool(), deps.SparkSubmitHandler) + s.AddTool(tools.NewListApplicationsTool(), deps.ListApplicationsHandler) + s.AddTool(tools.NewGetApplicationStatusTool(), deps.GetApplicationStatusHandler) + s.AddTool(tools.NewGetApplicationLogsTool(), deps.GetApplicationLogsHandler) + s.AddTool(tools.NewKillApplicationTool(), deps.KillApplicationHandler) + s.AddTool(tools.NewFetchURLTool(), deps.FetchURLHandler) + s.AddTool(tools.NewUploadFileTool(), deps.UploadFileHandler) + s.AddTool(tools.NewFetchSparkMetricsTool(), deps.FetchSparkMetricsHandler) + s.AddTool(tools.NewFetchClusterEnvTool(), deps.FetchClusterEnvHandler) + s.AddTool(tools.NewAnalyzeSparkLogTool(), deps.AnalyzeSparkLogHandler) + return s +} + +// Handler exposes the MCP server as an http.Handler for main.go to mount. +func Handler(deps *tools.Deps) (http.Handler, error) { + s := NewServer(deps) + return server.NewStreamableHTTPServer(s), nil +} diff --git a/internal/mcp/tools/deps.go b/internal/mcp/tools/deps.go new file mode 100644 index 0000000..0b6021f --- /dev/null +++ b/internal/mcp/tools/deps.go @@ -0,0 +1,22 @@ +package tools + +import ( + "log/slog" + "time" + + "spark-mcp-go/internal/analyzer" + "spark-mcp-go/internal/httpclient" + "spark-mcp-go/internal/storage" +) + +// Deps bundles the dependencies shared by all MCP Tool handlers. +type Deps struct { + Logger *slog.Logger + ClusterRepo *storage.ClusterRepo + SparkSubmitTimeout time.Duration + HTTPClient *httpclient.Client + MaxResponseBytes int64 + DataDir string // upload_file writes to DataDir/uploads + + AnalyzerThresholds analyzer.Thresholds +} diff --git a/internal/mcp/tools/get_application_logs.go b/internal/mcp/tools/get_application_logs.go new file mode 100644 index 0000000..3e0fcf9 --- /dev/null +++ b/internal/mcp/tools/get_application_logs.go @@ -0,0 +1,93 @@ +package tools + +import ( + "context" + + "github.com/mark3labs/mcp-go/mcp" + + "spark-mcp-go/internal/rm" +) + +const GetApplicationLogsName = "get_application_logs" + +const defaultTailBytes = 1 << 20 + +// NewGetApplicationLogsTool returns the schema for the get_application_logs MCP Tool. +func NewGetApplicationLogsTool() mcp.Tool { + return mcp.NewTool(GetApplicationLogsName, + mcp.WithDescription("Fetch YARN application logs. Tries amContainerLogs first, then aggregated-logs, then the legacy /logs endpoint. Returns {source, text}."), + mcp.WithString("cluster_id", + mcp.Required(), + mcp.Description("ID of the configured cluster (from list_clusters)"), + ), + mcp.WithString("app_id", + mcp.Required(), + mcp.Description("YARN application ID, e.g. application_1234567890_0001"), + ), + mcp.WithString("container", + mcp.Description("Container ID used for the aggregated-logs fallback, e.g. container_1234567890_0001_01_000001."), + ), + mcp.WithNumber("tail_bytes", + mcp.Description("Maximum bytes to return. 0 means use the server-wide response byte limit."), + ), + ) +} + +// GetApplicationLogsHandler retrieves logs using the RM fallback chain. +func (d *Deps) GetApplicationLogsHandler(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + clusterID, err := req.RequireString("cluster_id") + if err != nil { + return errResult("get_application_logs: " + err.Error()), nil + } + appID, err := req.RequireString("app_id") + if err != nil { + return errResult("get_application_logs: " + err.Error()), nil + } + + args := req.GetArguments() + container := "" + if v, ok := args["container"].(string); ok { + container = v + } + + tailBytes := defaultTailBytes + if v, ok := args["tail_bytes"].(float64); ok { + tailBytes = int(v) + } + if tailBytes == 0 { + tailBytes = int(d.MaxResponseBytes) + } + + callLog := startToolCall(ctx, d.Logger, GetApplicationLogsName, map[string]any{ + "cluster_id": clusterID, + "app_id": appID, + "container": container, + "tail_bytes": tailBytes, + }) + defer callLog.End() + + cl, err := d.ClusterRepo.Get(ctx, clusterID) + if err != nil { + callLog.WithError(err) + return errResult("get_application_logs: cluster " + clusterID + ": " + err.Error()), nil + } + + rmc := rm.New(d.HTTPClient, cl) + body, source, err := rmc.GetLogs(ctx, appID, container) + if err != nil { + callLog.WithError(err) + return errResult("get_application_logs: " + err.Error()), nil + } + + text := string(body) + if len(text) > tailBytes { + text = truncateMiddle(text, tailBytes) + } + + result := map[string]any{ + "source": source, + "text": text, + } + callLog.WithResult(map[string]any{"source": source, "bytes": len(body), "returned_bytes": len(text)}) + return textResult(encodeJSON(result)), nil +} diff --git a/internal/mcp/tools/get_application_status.go b/internal/mcp/tools/get_application_status.go new file mode 100644 index 0000000..a881d55 --- /dev/null +++ b/internal/mcp/tools/get_application_status.go @@ -0,0 +1,60 @@ +package tools + +import ( + "context" + + "github.com/mark3labs/mcp-go/mcp" + + "spark-mcp-go/internal/rm" +) + +const GetApplicationStatusName = "get_application_status" + +// NewGetApplicationStatusTool returns the schema for the get_application_status MCP Tool. +func NewGetApplicationStatusTool() mcp.Tool { + return mcp.NewTool(GetApplicationStatusName, + mcp.WithDescription("Get the detailed status of a single YARN application from the ResourceManager. Returns the raw RM JSON for the app."), + mcp.WithString("cluster_id", + mcp.Required(), + mcp.Description("ID of the configured cluster (from list_clusters)"), + ), + mcp.WithString("app_id", + mcp.Required(), + mcp.Description("YARN application ID, e.g. application_1234567890_0001"), + ), + ) +} + +// GetApplicationStatusHandler fetches a single application's details. +func (d *Deps) GetApplicationStatusHandler(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + clusterID, err := req.RequireString("cluster_id") + if err != nil { + return errResult("get_application_status: " + err.Error()), nil + } + appID, err := req.RequireString("app_id") + if err != nil { + return errResult("get_application_status: " + err.Error()), nil + } + + callLog := startToolCall(ctx, d.Logger, GetApplicationStatusName, map[string]any{ + "cluster_id": clusterID, + "app_id": appID, + }) + defer callLog.End() + + cl, err := d.ClusterRepo.Get(ctx, clusterID) + if err != nil { + callLog.WithError(err) + return errResult("get_application_status: cluster " + clusterID + ": " + err.Error()), nil + } + + rmc := rm.New(d.HTTPClient, cl) + raw, err := rmc.GetApp(ctx, appID) + if err != nil { + callLog.WithError(err) + return errResult("get_application_status: " + err.Error()), nil + } + + callLog.WithResult(map[string]any{"bytes": len(raw)}) + return textResult(string(raw)), nil +} diff --git a/internal/mcp/tools/kill_application.go b/internal/mcp/tools/kill_application.go new file mode 100644 index 0000000..258859d --- /dev/null +++ b/internal/mcp/tools/kill_application.go @@ -0,0 +1,60 @@ +package tools + +import ( + "context" + + "github.com/mark3labs/mcp-go/mcp" + + "spark-mcp-go/internal/rm" +) + +const KillApplicationName = "kill_application" + +// NewKillApplicationTool returns the schema for the kill_application MCP Tool. +func NewKillApplicationTool() mcp.Tool { + return mcp.NewTool(KillApplicationName, + mcp.WithDescription("Kill (move to KILLED state) a YARN application through the ResourceManager. Returns the RM JSON response."), + mcp.WithString("cluster_id", + mcp.Required(), + mcp.Description("ID of the configured cluster (from list_clusters)"), + ), + mcp.WithString("app_id", + mcp.Required(), + mcp.Description("YARN application ID, e.g. application_1234567890_0001"), + ), + ) +} + +// KillApplicationHandler sends a KILLED state update to the ResourceManager. +func (d *Deps) KillApplicationHandler(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + clusterID, err := req.RequireString("cluster_id") + if err != nil { + return errResult("kill_application: " + err.Error()), nil + } + appID, err := req.RequireString("app_id") + if err != nil { + return errResult("kill_application: " + err.Error()), nil + } + + callLog := startToolCall(ctx, d.Logger, KillApplicationName, map[string]any{ + "cluster_id": clusterID, + "app_id": appID, + }) + defer callLog.End() + + cl, err := d.ClusterRepo.Get(ctx, clusterID) + if err != nil { + callLog.WithError(err) + return errResult("kill_application: cluster " + clusterID + ": " + err.Error()), nil + } + + rmc := rm.New(d.HTTPClient, cl) + raw, err := rmc.KillApp(ctx, appID) + if err != nil { + callLog.WithError(err) + return errResult("kill_application: " + err.Error()), nil + } + + callLog.WithResult(map[string]any{"bytes": len(raw)}) + return textResult(string(raw)), nil +} diff --git a/internal/mcp/tools/list_applications.go b/internal/mcp/tools/list_applications.go new file mode 100644 index 0000000..4374a15 --- /dev/null +++ b/internal/mcp/tools/list_applications.go @@ -0,0 +1,69 @@ +package tools + +import ( + "context" + + "github.com/mark3labs/mcp-go/mcp" + + "spark-mcp-go/internal/rm" +) + +const ListApplicationsName = "list_applications" + +// NewListApplicationsTool returns the schema for the list_applications MCP Tool. +func NewListApplicationsTool() mcp.Tool { + return mcp.NewTool(ListApplicationsName, + mcp.WithDescription("List YARN applications from a cluster's ResourceManager. Returns the raw RM JSON response so the agent can inspect app IDs, states, and owners."), + mcp.WithString("cluster_id", + mcp.Required(), + mcp.Description("ID of the configured cluster (from list_clusters)"), + ), + mcp.WithString("state", + mcp.Description("Filter by application state. YARN accepts comma-separated states; common values: NEW, NEW_SAVING, SUBMITTED, ACCEPTED, RUNNING, FINISHED, FAILED, KILLED."), + ), + mcp.WithString("user", + mcp.Description("Filter by submitting user."), + ), + ) +} + +// ListApplicationsHandler queries the ResourceManager for applications. +func (d *Deps) ListApplicationsHandler(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + clusterID, err := req.RequireString("cluster_id") + if err != nil { + return errResult("list_applications: " + err.Error()), nil + } + + args := req.GetArguments() + state := "" + if v, ok := args["state"].(string); ok { + state = v + } + user := "" + if v, ok := args["user"].(string); ok { + user = v + } + + callLog := startToolCall(ctx, d.Logger, ListApplicationsName, map[string]any{ + "cluster_id": clusterID, + "state": state, + "user": user, + }) + defer callLog.End() + + cl, err := d.ClusterRepo.Get(ctx, clusterID) + if err != nil { + callLog.WithError(err) + return errResult("list_applications: cluster " + clusterID + ": " + err.Error()), nil + } + + rmc := rm.New(d.HTTPClient, cl) + raw, err := rmc.ListApps(ctx, state, user) + if err != nil { + callLog.WithError(err) + return errResult("list_applications: " + err.Error()), nil + } + + callLog.WithResult(map[string]any{"bytes": len(raw)}) + return textResult(string(raw)), nil +} diff --git a/internal/mcp/tools/rm_test.go b/internal/mcp/tools/rm_test.go new file mode 100644 index 0000000..8ccf2dc --- /dev/null +++ b/internal/mcp/tools/rm_test.go @@ -0,0 +1,176 @@ +package tools + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/mark3labs/mcp-go/mcp" + + "spark-mcp-go/internal/cluster" + "spark-mcp-go/internal/httpclient" + "spark-mcp-go/internal/storage" +) + +func testDeps(t *testing.T, srv *httptest.Server) (*Deps, *storage.ClusterRepo) { + t.Helper() + db, err := storage.Open(":memory:") + if err != nil { + t.Fatalf("open db: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + + cl := &cluster.Cluster{ + ID: "cluster-a", + Name: "Cluster A", + RMURL: srv.URL, + SHSURL: "http://shs.example.com:18080", + SparkSubmitExecuteBin: "/usr/bin/spark-submit", + IsActive: true, + AuthType: cluster.AuthNone, + } + if err := db.Clusters().Create(context.Background(), cl); err != nil { + t.Fatalf("create cluster: %v", err) + } + + return &Deps{ + HTTPClient: httpclient.New(httpclient.Config{ + Timeout: 5 * time.Second, + MaxResponseBytes: 1 << 20, + }), + ClusterRepo: db.Clusters(), + MaxResponseBytes: 1 << 20, + }, db.Clusters() +} + +func newToolRequest(name string, args map[string]any) mcp.CallToolRequest { + return mcp.CallToolRequest{ + Request: mcp.Request{Method: "tools/call"}, + Params: mcp.CallToolParams{ + Name: name, + Arguments: args, + }, + } +} + +func TestListApplications_EndToEnd(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/ws/v1/cluster/apps" { + t.Errorf("path=%q", r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"apps":{"app":[{"id":"application_1"}]}}`)) + })) + defer srv.Close() + + deps, _ := testDeps(t, srv) + req := newToolRequest("list_applications", map[string]any{ + "cluster_id": "cluster-a", + "state": "RUNNING", + }) + + res, err := deps.ListApplicationsHandler(context.Background(), req) + if err != nil { + t.Fatalf("handler error: %v", err) + } + if res.IsError { + t.Fatalf("unexpected error result: %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, "application_1") { + t.Errorf("result missing app: %s", text.Text) + } +} + +func TestKillApplication_AppIDNotProvided(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + t.Errorf("unexpected RM call") + })) + defer srv.Close() + + deps, _ := testDeps(t, srv) + req := newToolRequest("kill_application", map[string]any{ + "cluster_id": "cluster-a", + }) + + res, err := deps.KillApplicationHandler(context.Background(), req) + if err != nil { + t.Fatalf("handler error: %v", err) + } + if !res.IsError { + t.Errorf("expected error result, got: %v", res.Content) + } +} + +func TestGetApplicationLogs_SourceReported(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/ws/v1/cluster/apps/app_123/amContainerLogs" { + w.Write([]byte("driver stdout")) + return + } + t.Errorf("unexpected path %q", r.URL.Path) + })) + defer srv.Close() + + deps, _ := testDeps(t, srv) + req := newToolRequest("get_application_logs", map[string]any{ + "cluster_id": "cluster-a", + "app_id": "app_123", + }) + + res, err := deps.GetApplicationLogsHandler(context.Background(), req) + if err != nil { + t.Fatalf("handler error: %v", err) + } + if res.IsError { + t.Fatalf("unexpected error result: %v", res.Content) + } + + text, ok := mcp.AsTextContent(res.Content[0]) + if !ok { + t.Fatalf("content is not text: %T", res.Content[0]) + } + var payload map[string]any + if err := json.Unmarshal([]byte(text.Text), &payload); err != nil { + t.Fatalf("unmarshal result: %v", err) + } + if payload["source"] != "amContainerLogs" { + t.Errorf("source=%v, want amContainerLogs", payload["source"]) + } + if !strings.Contains(payload["text"].(string), "driver stdout") { + t.Errorf("text missing logs: %v", payload["text"]) + } +} + +func TestKillApplication_NotFound(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/ws/v1/cluster/apps/app_123/state" { + w.WriteHeader(http.StatusNotFound) + w.Write([]byte(`{"error":"Not Found"}`)) + return + } + t.Errorf("unexpected path %q", r.URL.Path) + })) + defer srv.Close() + + deps, _ := testDeps(t, srv) + req := newToolRequest("kill_application", map[string]any{ + "cluster_id": "cluster-a", + "app_id": "app_123", + }) + + res, err := deps.KillApplicationHandler(context.Background(), req) + if err != nil { + t.Fatalf("handler error: %v", err) + } + if !res.IsError { + t.Errorf("expected error result, got: %v", res.Content) + } +} diff --git a/internal/rm/client.go b/internal/rm/client.go new file mode 100644 index 0000000..a75468b --- /dev/null +++ b/internal/rm/client.go @@ -0,0 +1,131 @@ +package rm + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + "strings" + + "spark-mcp-go/internal/cluster" + "spark-mcp-go/internal/httpclient" +) + +// Client is a YARN ResourceManager REST API client bound to a single cluster. +type Client struct { + hc *httpclient.Client + cluster *cluster.Cluster +} + +// New creates a ResourceManager client bound to c. +func New(hc *httpclient.Client, c *cluster.Cluster) *Client { + return &Client{hc: hc, cluster: c} +} + +// ListApps returns raw JSON from GET /ws/v1/cluster/apps[?state=&user=]. +func (r *Client) ListApps(ctx context.Context, state, user string) (json.RawMessage, error) { + base := r.cluster.RMURL + "/ws/v1/cluster/apps" + u, err := url.Parse(base) + if err != nil { + return nil, err + } + q := u.Query() + if state != "" { + q.Set("state", state) + } + if user != "" { + q.Set("user", user) + } + u.RawQuery = q.Encode() + return r.do(ctx, http.MethodGet, u.String(), nil) +} + +// GetApp returns raw JSON from GET /ws/v1/cluster/apps/{id}. +func (r *Client) GetApp(ctx context.Context, appID string) (json.RawMessage, error) { + return r.do(ctx, http.MethodGet, r.cluster.RMURL+"/ws/v1/cluster/apps/"+appID, nil) +} + +// KillApp PUT {"state":"KILLED"} to /ws/v1/cluster/apps/{id}/state. +func (r *Client) KillApp(ctx context.Context, appID string) (json.RawMessage, error) { + body := `{"state":"KILLED"}` + return r.do(ctx, http.MethodPut, r.cluster.RMURL+"/ws/v1/cluster/apps/"+appID+"/state", + strings.NewReader(body)) +} + +// GetLogs walks a fallback chain to retrieve YARN application logs: +// 1. GET /ws/v1/cluster/apps/{id}/amContainerLogs (RM often 307 to a NodeManager) +// 2. GET /ws/v1/cluster/apps/{id}/aggregated-logs?container=... +// 3. GET /ws/v1/cluster/apps/{id}/logs +// +// The first successful endpoint wins. The returned source names the endpoint +// that produced the body. +func (r *Client) GetLogs(ctx context.Context, appID, container string) (body []byte, source string, err error) { + endpoints := []struct{ path, name string }{ + {fmt.Sprintf("/ws/v1/cluster/apps/%s/amContainerLogs", appID), "amContainerLogs"}, + {fmt.Sprintf("/ws/v1/cluster/apps/%s/aggregated-logs?container=%s", appID, url.QueryEscape(container)), "aggregated-logs"}, + {fmt.Sprintf("/ws/v1/cluster/apps/%s/logs", appID), "logs"}, + } + var lastErr error + for _, ep := range endpoints { + b, e := r.doBytes(ctx, http.MethodGet, r.cluster.RMURL+ep.path, nil) + if e == nil { + return b, ep.name, nil + } + lastErr = e + } + return nil, "", fmt.Errorf("rm: all log endpoints failed: %w", lastErr) +} + +// do executes a single HTTP request with auth and allowed-host exemptions, +// returning the validated JSON body. +func (r *Client) do(ctx context.Context, method, rawURL string, body io.Reader) (json.RawMessage, error) { + b, err := r.doBytes(ctx, method, rawURL, body) + if err != nil { + return nil, err + } + raw := json.RawMessage(b) + if !json.Valid(raw) { + return nil, fmt.Errorf("rm: invalid JSON response from %s", rawURL) + } + return raw, nil +} + +func (r *Client) doBytes(ctx context.Context, method, rawURL string, body io.Reader) ([]byte, error) { + req, err := http.NewRequestWithContext(ctx, method, rawURL, body) + if err != nil { + return nil, err + } + if err := httpclient.ApplyAuth(req, "", r.cluster); err != nil { + return nil, err + } + resp, err := r.hc.DoWithRedirect(ctx, req, 5, r.allowedHosts()...) + if err != nil { + return nil, err + } + defer resp.Body.Close() + if resp.StatusCode >= 400 { + return nil, fmt.Errorf("rm: %s %s returned %d", method, rawURL, resp.StatusCode) + } + return io.ReadAll(resp.Body) +} + +func (r *Client) allowedHosts() []string { + hosts := []string{extractHost(r.cluster.RMURL)} + for _, pat := range r.cluster.URLAllowlist { + if pat == "" { + continue + } + hosts = append(hosts, pat) + } + return hosts +} + +func extractHost(u string) string { + parsed, err := url.Parse(u) + if err != nil { + return "" + } + return parsed.Hostname() +} diff --git a/internal/rm/client_test.go b/internal/rm/client_test.go new file mode 100644 index 0000000..2331d29 --- /dev/null +++ b/internal/rm/client_test.go @@ -0,0 +1,238 @@ +package rm + +import ( + "context" + "encoding/base64" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "spark-mcp-go/internal/cluster" + "spark-mcp-go/internal/httpclient" +) + +func testClient(srv *httptest.Server, c *cluster.Cluster) *Client { + hc := httpclient.New(httpclient.Config{Timeout: 5 * time.Second, MaxResponseBytes: 1 << 20}) + return New(hc, c) +} + +func TestListApps(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/ws/v1/cluster/apps" { + t.Errorf("path=%q, want /ws/v1/cluster/apps", r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"apps":{"app":[]}}`)) + })) + defer srv.Close() + + cl := &cluster.Cluster{RMURL: srv.URL} + raw, err := testClient(srv, cl).ListApps(context.Background(), "", "") + if err != nil { + t.Fatalf("ListApps: %v", err) + } + if string(raw) != `{"apps":{"app":[]}}` { + t.Errorf("body=%q", string(raw)) + } +} + +func TestListApps_WithQuery(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + q := r.URL.Query() + if q.Get("state") != "RUNNING" { + t.Errorf("state=%q, want RUNNING", q.Get("state")) + } + if q.Get("user") != "yarn" { + t.Errorf("user=%q, want yarn", q.Get("user")) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"apps":{"app":[{"id":"app_1"}]}}`)) + })) + defer srv.Close() + + cl := &cluster.Cluster{RMURL: srv.URL} + raw, err := testClient(srv, cl).ListApps(context.Background(), "RUNNING", "yarn") + if err != nil { + t.Fatalf("ListApps: %v", err) + } + if !strings.Contains(string(raw), "app_1") { + t.Errorf("body=%q", string(raw)) + } +} + +func TestGetApp(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/ws/v1/cluster/apps/app_123" { + t.Errorf("path=%q", r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"app":{"id":"app_123","state":"RUNNING"}}`)) + })) + defer srv.Close() + + cl := &cluster.Cluster{RMURL: srv.URL} + raw, err := testClient(srv, cl).GetApp(context.Background(), "app_123") + if err != nil { + t.Fatalf("GetApp: %v", err) + } + if !strings.Contains(string(raw), "RUNNING") { + t.Errorf("body=%q", string(raw)) + } +} + +func TestKillApp(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/ws/v1/cluster/apps/app_123/state" { + t.Errorf("path=%q", r.URL.Path) + } + if r.Method != http.MethodPut { + t.Errorf("method=%q, want PUT", r.Method) + } + body, _ := io.ReadAll(r.Body) + if string(body) != `{"state":"KILLED"}` { + t.Errorf("body=%q", string(body)) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"state":"KILLED"}`)) + })) + defer srv.Close() + + cl := &cluster.Cluster{RMURL: srv.URL} + raw, err := testClient(srv, cl).KillApp(context.Background(), "app_123") + if err != nil { + t.Fatalf("KillApp: %v", err) + } + if string(raw) != `{"state":"KILLED"}` { + t.Errorf("body=%q", string(raw)) + } +} + +func TestGetLogs_PrimarySuccess(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/ws/v1/cluster/apps/app_123/amContainerLogs" { + w.Write([]byte("am logs")) + return + } + t.Errorf("unexpected path %q", r.URL.Path) + })) + defer srv.Close() + + cl := &cluster.Cluster{RMURL: srv.URL} + body, source, err := testClient(srv, cl).GetLogs(context.Background(), "app_123", "container_1") + if err != nil { + t.Fatalf("GetLogs: %v", err) + } + if source != "amContainerLogs" { + t.Errorf("source=%q, want amContainerLogs", source) + } + if string(body) != "am logs" { + t.Errorf("body=%q", string(body)) + } +} + +func TestGetLogs_PrimaryRedirects(t *testing.T) { + srvNM := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + auth := r.Header.Get("Authorization") + if auth == "" { + t.Error("Authorization header missing after redirect") + } + w.Write([]byte("nm logs")) + })) + defer srvNM.Close() + + srvRM := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/ws/v1/cluster/apps/app_123/amContainerLogs" { + w.Header().Set("Location", srvNM.URL+"/node/containerlogs/container_1/root") + w.WriteHeader(http.StatusTemporaryRedirect) + return + } + t.Errorf("unexpected RM path %q", r.URL.Path) + })) + defer srvRM.Close() + + nmHost := srvNM.Listener.Addr().String() + cl := &cluster.Cluster{ + RMURL: srvRM.URL, + AuthType: cluster.AuthBasic, + AuthUsername: "u", + AuthPassword: "p", + URLAllowlist: []string{nmHost}, + } + body, source, err := testClient(srvRM, cl).GetLogs(context.Background(), "app_123", "container_1") + if err != nil { + t.Fatalf("GetLogs: %v", err) + } + if source != "amContainerLogs" { + t.Errorf("source=%q, want amContainerLogs", source) + } + if string(body) != "nm logs" { + t.Errorf("body=%q", string(body)) + } +} + +func TestGetLogs_FallbackToAggregated(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/ws/v1/cluster/apps/app_123/amContainerLogs": + w.WriteHeader(http.StatusInternalServerError) + case "/ws/v1/cluster/apps/app_123/aggregated-logs": + if r.URL.Query().Get("container") != "container_1" { + t.Errorf("container=%q, want container_1", r.URL.Query().Get("container")) + } + w.Write([]byte("aggregated logs")) + default: + t.Errorf("unexpected path %q", r.URL.Path) + } + })) + defer srv.Close() + + cl := &cluster.Cluster{RMURL: srv.URL} + body, source, err := testClient(srv, cl).GetLogs(context.Background(), "app_123", "container_1") + if err != nil { + t.Fatalf("GetLogs: %v", err) + } + if source != "aggregated-logs" { + t.Errorf("source=%q, want aggregated-logs", source) + } + if string(body) != "aggregated logs" { + t.Errorf("body=%q", string(body)) + } +} + +func TestApplyAuth_SimpleUserName(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.URL.Query().Get("user.name"); got != "alice" { + t.Errorf("user.name=%q, want alice", got) + } + w.Write([]byte(`{"apps":{"app":[]}}`)) + })) + defer srv.Close() + + cl := &cluster.Cluster{RMURL: srv.URL, AuthType: cluster.AuthSimple, AuthUsername: "alice"} + _, err := testClient(srv, cl).ListApps(context.Background(), "", "") + if err != nil { + t.Fatalf("ListApps: %v", err) + } +} + +func TestApplyAuth_BasicHeader(t *testing.T) { + captured := "" + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + captured = r.Header.Get("Authorization") + w.Write([]byte(`{"apps":{"app":[]}}`)) + })) + defer srv.Close() + + cl := &cluster.Cluster{RMURL: srv.URL, AuthType: cluster.AuthBasic, AuthUsername: "u", AuthPassword: "p"} + _, err := testClient(srv, cl).ListApps(context.Background(), "", "") + if err != nil { + t.Fatalf("ListApps: %v", err) + } + want := "Basic " + base64.StdEncoding.EncodeToString([]byte("u:p")) + if captured != want { + t.Errorf("Authorization=%q, want %q", captured, want) + } +}