Phase 5 Batch 2: 4 RM Tool + rm 客户端 + 日志降级链
- 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 <noreply@anthropic.com>
This commit is contained in:
tao.chen
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user