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:
@@ -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()
|
||||
}
|
||||
Reference in New Issue
Block a user