spark-mcp/internal/rm/client.go

package rm

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"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=&queue=&limit=].
func (r *Client) ListApps(ctx context.Context, state, user, queue string, limit int) (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)
	}
	if queue != "" {
		q.Set("queue", queue)
	}
	if limit > 0 {
		q.Set("limit", strconv.Itoa(limit))
	}
	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 endpoint, often 307 to a NodeManager)
//  2. GET /ws/v1/cluster/apps/{id}/aggregated-logs?container=...
//  3. GET /ws/v1/cluster/apps/{id}/logs
//  4. GET the NodeManager URL from the app's amContainerLogs field.
//
// 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.doBytesWithHosts(ctx, http.MethodGet, r.cluster.RMURL+ep.path, nil)
		if e == nil {
			return b, ep.name, nil
		}
		lastErr = e
	}
	// Try fetching the NodeManager's stdout URL directly.
	app, err := r.GetApp(ctx, appID)
	if err == nil {
		var resp struct {
			App struct {
				AMContainerLogs string `json:"amContainerLogs"`
			} `json:"app"`
		}
		if err := json.Unmarshal(app, &resp); err == nil && resp.App.AMContainerLogs != "" {
			nmHost := extractHost(resp.App.AMContainerLogs)
			b, e := r.doBytesWithHosts(ctx, http.MethodGet, resp.App.AMContainerLogs, nil, nmHost)
			if e == nil {
				return b, "amContainerLogs-direct", nil
			}
			lastErr = e
		}
	} else {
		lastErr = err
	}
	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) doBytesWithHosts(ctx context.Context, method, rawURL string, body io.Reader, extraHosts ...string) ([]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, append(r.allowedHosts(), extraHosts...)...)
	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) doBytes(ctx context.Context, method, rawURL string, body io.Reader) ([]byte, error) {
	return r.doBytesWithHosts(ctx, method, rawURL, 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()
}