tao.chen/spark-mcp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 2 Wiki Activity

rm: list_applications 加 limit/queue + 日志 fallback 链补 NM 直连

Browse Source 
P0 (Python parity 缺口):
  - rm.Client.ListApps 新增 queue + limit 查询参数
  - list_applications MCP tool schema 暴露 queue (string) 和 limit
    (number, default 100), handler 透传给 RM client
  - 限流原因: 生产 RM 上无 limit 会拉回 N MB JSON 撞 MaxResponseBytes
  - limit <= 0 不发送 query 参数, 跟 RM 默认行为一致; 老的调用方
    (TestListApplications_EndToEnd 等) 不传参时行为不变
  - 4 个旧 ListApps 测试调用点跟着更新, 加 3 个新测试:
    TestListApps_WithLimit, TestListApps_WithQueue,
    TestListApps_LimitZeroNoParam

P1.3 (Python fallback 行为补全):
  - rm.Client.GetLogs 加第 4 步 fallback: 当前 3 步
    (amContainerLogs/aggregated-logs/logs) 全失败时, 调 GetApp
    解析 app.amContainerLogs 字段, GET 该 URL 直连 NodeManager
  - 新 source 名 'amContainerLogs-direct' 区分 RM-level endpoint
  - 提取 doBytesWithHosts(extraHosts...) 支持 per-call host 白名单,
    NM host 动态加进 allowedHosts, SSRF 保护不破 (URL 是 RM 响应里
    回来的, 不是 LLM 任意填的)
  - 新增 TestGetLogs_FallbackToAMDirect (成功路径) 和
    TestGetLogs_FallbackFailsWhenAMFieldMissing (amContainerLogs
    字段为空时正常返回 error)

P1.4 (零代码改动 + 文档化):
  - get_application_logs / get_application_status 工具 description
    更新, 提示 LLM raw RM JSON 里包含 amContainerLogs 字段, 可在
    aggregated-logs 全部失败时自己用 fetch_url 直连 NM
  - get_application_status 函数体不变 (本就是透传 raw JSON)

验证:
  - go build / go vet / go test 全部通过
  - 5 个新测试全 PASS, 老的 TestListApplications_EndToEnd 仍 PASS
  - 单 ListApps 调用点 (list_applications.go:81) 编译过

未提交: spark-mcp-linux-amd64 (本地 build 产物, 当前 .gitignore
没拦, 需另行决定)

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
tao.chen
2026-07-13 15:15:21 +08:00
co-authored by Claude
parent bf2640d8fd
commit 2537976722
5 changed files with 176 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/rm/client_test.go
+116 -4
View File
@@ -3,6 +3,7 @@ package rm
import (
"context"
"encoding/base64"
"fmt"
"io"
"net/http"
"net/http/httptest"
@@ -30,7 +31,7 @@ func TestListApps(t *testing.T) {
defer srv.Close()
cl := &cluster.Cluster{RMURL: srv.URL}
raw, err := testClient(srv, cl).ListApps(context.Background(), "", "")
raw, err := testClient(srv, cl).ListApps(context.Background(), "", "", "", 0)
if err != nil {
t.Fatalf("ListApps: %v", err)
}
@@ -54,7 +55,7 @@ func TestListApps_WithQuery(t *testing.T) {
defer srv.Close()
cl := &cluster.Cluster{RMURL: srv.URL}
raw, err := testClient(srv, cl).ListApps(context.Background(), "RUNNING", "yarn")
raw, err := testClient(srv, cl).ListApps(context.Background(), "RUNNING", "yarn", "", 0)
if err != nil {
t.Fatalf("ListApps: %v", err)
}
@@ -212,7 +213,7 @@ func TestApplyAuth_SimpleUserName(t *testing.T) {
defer srv.Close()
cl := &cluster.Cluster{RMURL: srv.URL, AuthType: cluster.AuthSimple, AuthUsername: "alice"}
_, err := testClient(srv, cl).ListApps(context.Background(), "", "")
_, err := testClient(srv, cl).ListApps(context.Background(), "", "", "", 0)
if err != nil {
t.Fatalf("ListApps: %v", err)
}
@@ -227,7 +228,7 @@ func TestApplyAuth_BasicHeader(t *testing.T) {
defer srv.Close()
cl := &cluster.Cluster{RMURL: srv.URL, AuthType: cluster.AuthBasic, AuthUsername: "u", AuthPassword: "p"}
_, err := testClient(srv, cl).ListApps(context.Background(), "", "")
_, err := testClient(srv, cl).ListApps(context.Background(), "", "", "", 0)
if err != nil {
t.Fatalf("ListApps: %v", err)
}
@@ -236,3 +237,114 @@ func TestApplyAuth_BasicHeader(t *testing.T) {
t.Errorf("Authorization=%q, want %q", captured, want)
}
}
func TestListApps_WithLimit(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Query().Get("limit") != "50" {
t.Errorf("limit=%q, want 50", r.URL.Query().Get("limit"))
}
w.Write([]byte(`{"apps":{"app":[]}}`))
}))
defer srv.Close()
cl := &cluster.Cluster{RMURL: srv.URL}
_, err := testClient(srv, cl).ListApps(context.Background(), "", "", "", 50)
if err != nil {
t.Fatalf("ListApps: %v", err)
}
}
func TestListApps_WithQueue(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Query().Get("queue") != "root.users.alice" {
t.Errorf("queue=%q, want root.users.alice", r.URL.Query().Get("queue"))
}
w.Write([]byte(`{"apps":{"app":[]}}`))
}))
defer srv.Close()
cl := &cluster.Cluster{RMURL: srv.URL}
_, err := testClient(srv, cl).ListApps(context.Background(), "", "", "root.users.alice", 0)
if err != nil {
t.Fatalf("ListApps: %v", err)
}
}
func TestListApps_LimitZeroNoParam(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Query().Get("limit") != "" {
t.Errorf("limit should not be sent when 0; got %q", r.URL.Query().Get("limit"))
}
w.Write([]byte(`{"apps":{"app":[]}}`))
}))
defer srv.Close()
cl := &cluster.Cluster{RMURL: srv.URL}
_, err := testClient(srv, cl).ListApps(context.Background(), "", "", "", 0)
if err != nil {
t.Fatalf("ListApps: %v", err)
}
}
func TestGetLogs_FallbackToAMDirect(t *testing.T) {
// Mock NM that we want GetLogs to hit after standard endpoints fail.
nmSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Write([]byte("nm direct logs"))
}))
defer nmSrv.Close()
nmHost := nmSrv.Listener.Addr().String()
rmSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch {
case r.URL.Path == "/ws/v1/cluster/apps/app_123/amContainerLogs":
w.WriteHeader(http.StatusInternalServerError)
case r.URL.Path == "/ws/v1/cluster/apps/app_123/aggregated-logs":
w.WriteHeader(http.StatusNotImplemented)
case r.URL.Path == "/ws/v1/cluster/apps/app_123/logs":
w.WriteHeader(http.StatusNotFound)
case r.URL.Path == "/ws/v1/cluster/apps/app_123" && r.Method == http.MethodGet:
// GetApp returns amContainerLogs pointing to the mock NM.
w.Write([]byte(fmt.Sprintf(`{"app":{"id":"app_123","amContainerLogs":"%s/node/containerlogs/container_1/root"}}`, nmSrv.URL)))
default:
t.Errorf("unexpected RM path %q", r.URL.Path)
}
}))
defer rmSrv.Close()
cl := &cluster.Cluster{
RMURL: rmSrv.URL,
URLAllowlist: []string{nmHost}, // allow NM host for SSRF
}
body, source, err := testClient(rmSrv, cl).GetLogs(context.Background(), "app_123", "container_1")
if err != nil {
t.Fatalf("GetLogs: %v", err)
}
if source != "amContainerLogs-direct" {
t.Errorf("source=%q, want amContainerLogs-direct", source)
}
if string(body) != "nm direct logs" {
t.Errorf("body=%q", string(body))
}
}
func TestGetLogs_FallbackFailsWhenAMFieldMissing(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch {
case r.URL.Path == "/ws/v1/cluster/apps/app_123/amContainerLogs":
w.WriteHeader(http.StatusInternalServerError)
case r.URL.Path == "/ws/v1/cluster/apps/app_123/aggregated-logs":
w.WriteHeader(http.StatusNotImplemented)
case r.URL.Path == "/ws/v1/cluster/apps/app_123/logs":
w.WriteHeader(http.StatusNotFound)
case r.URL.Path == "/ws/v1/cluster/apps/app_123" && r.Method == http.MethodGet:
// amContainerLogs is empty
w.Write([]byte(`{"app":{"id":"app_123"}}`))
default:
t.Errorf("unexpected path %q", r.URL.Path)
}
}))
defer srv.Close()
cl := &cluster.Cluster{RMURL: srv.URL}
_, _, err := testClient(srv, cl).GetLogs(context.Background(), "app_123", "container_1")
if err == nil {
t.Fatal("expected error when all endpoints fail and amContainerLogs is empty")
}
}