wenshu/k8s-orchestrator/internal/api/server.go

package api

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"strings"

	"github.com/go-chi/chi/v5"
	apierrors "k8s.io/apimachinery/pkg/api/errors"

	"github.com/cyrus/spark-orchestrator/internal/service"
)

type Server struct {
	svc *service.OrchestratorService
}

func NewServer(svc *service.OrchestratorService) *Server {
	return &Server{svc: svc}
}

func (s *Server) Router() http.Handler {
	r := chi.NewRouter()

	r.Get("/healthz", func(w http.ResponseWriter, _ *http.Request) {
		writeJSON(w, http.StatusOK, map[string]string{"status": "ok"})
	})

	r.Route("/v1", func(r chi.Router) {
		r.Post("/nodes", s.handleAddNode)
		r.Delete("/nodes/{nodeId}", s.handleDeleteNode)
		r.Get("/nodes", s.handleListNodes)

		r.Post("/spark-clusters", s.handleCreateCluster)
		r.Get("/spark-clusters/{clusterId}", s.handleGetCluster)
		r.Post("/spark-clusters/{clusterId}/release", s.handleReleaseCluster)
		r.Post("/spark-clusters/{clusterId}/scale-out", s.handleScaleOut)
		r.Post("/spark-clusters/{clusterId}/scale-in", s.handleScaleIn)

			r.Post("/starrocks-clusters", s.handleCreateStarRocksCluster)
			r.Get("/starrocks-clusters/{clusterId}", s.handleGetStarRocksCluster)
			r.Get("/starrocks-clusters/{clusterId}/direct-endpoint", s.handleGetStarRocksDirectEndpoint)
			r.Post("/starrocks-clusters/{clusterId}/release", s.handleReleaseStarRocksCluster)
			r.Post("/starrocks-clusters/{clusterId}/scale-out", s.handleScaleOutStarRocksCluster)
			r.Post("/starrocks-clusters/{clusterId}/scale-in", s.handleScaleInStarRocksCluster)
		r.Post("/starrocks-jobs", s.handleSubmitStarRocksJob)
		r.Get("/starrocks-jobs/operations/{operationId}", s.handleGetStarRocksJobStatus)

		r.Post("/spark-jobs", s.handleSubmitSparkJob)
		r.Get("/spark-jobs/operations/{operationId}", s.handleGetSparkJobStatus)
		r.Delete("/spark-jobs/operations/{operationId}", s.handleDeleteSparkJob)
		r.Get("/spark-jobs/operations/{operationId}/result", s.handleGetSparkJobResult)
		r.Get("/operations/{operationId}", s.handleGetOperation)
	})

	return r
}

func (s *Server) handleAddNode(w http.ResponseWriter, r *http.Request) {
	idKey, err := requiredIdempotencyKey(r)
	if err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}
	var req service.AddNodeRequest
	if err := decodeJSON(r, &req); err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}
	resp, err := s.svc.AddNode(r.Context(), req, idKey)
	if err != nil {
		writeServiceError(w, err)
		return
	}
	writeJSON(w, http.StatusAccepted, resp)
}

func (s *Server) handleDeleteNode(w http.ResponseWriter, r *http.Request) {
	idKey, err := requiredIdempotencyKey(r)
	if err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}
	nodeID := chi.URLParam(r, "nodeId")
	resp, err := s.svc.DeleteNode(r.Context(), nodeID, idKey)
	if err != nil {
		writeServiceError(w, err)
		return
	}
	writeJSON(w, http.StatusAccepted, resp)
}

func (s *Server) handleListNodes(w http.ResponseWriter, r *http.Request) {
	nodes, err := s.svc.ListNodes(r.Context())
	if err != nil {
		writeServiceError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, map[string]interface{}{"items": nodes})
}

func (s *Server) handleCreateCluster(w http.ResponseWriter, r *http.Request) {
	idKey, err := requiredIdempotencyKey(r)
	if err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}
	var req service.CreateClusterRequest
	if err := decodeJSON(r, &req); err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}
	resp, err := s.svc.CreateCluster(r.Context(), req, idKey)
	if err != nil {
		writeServiceError(w, err)
		return
	}
	writeJSON(w, http.StatusAccepted, resp)
}

func (s *Server) handleGetCluster(w http.ResponseWriter, r *http.Request) {
	cluster, err := s.svc.GetCluster(r.Context(), chi.URLParam(r, "clusterId"))
	if err != nil {
		writeServiceError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, cluster)
}

func (s *Server) handleReleaseCluster(w http.ResponseWriter, r *http.Request) {
	idKey, err := requiredIdempotencyKey(r)
	if err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}
	clusterID := chi.URLParam(r, "clusterId")
	resp, err := s.svc.ReleaseCluster(r.Context(), clusterID, idKey)
	if err != nil {
		writeServiceError(w, err)
		return
	}
	writeJSON(w, http.StatusAccepted, resp)
}

func (s *Server) handleScaleOut(w http.ResponseWriter, r *http.Request) {
	s.handleResize(w, r, true)
}

func (s *Server) handleScaleIn(w http.ResponseWriter, r *http.Request) {
	s.handleResize(w, r, false)
}

func (s *Server) handleResize(w http.ResponseWriter, r *http.Request, scaleOut bool) {
	idKey, err := requiredIdempotencyKey(r)
	if err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}
	var req service.ResizeClusterRequest
	if err := decodeJSON(r, &req); err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}
	clusterID := chi.URLParam(r, "clusterId")
	var resp *service.OperationResponse
	if scaleOut {
		resp, err = s.svc.ScaleOut(r.Context(), clusterID, req, idKey)
	} else {
		resp, err = s.svc.ScaleIn(r.Context(), clusterID, req, idKey)
	}
	if err != nil {
		writeServiceError(w, err)
		return
	}
	writeJSON(w, http.StatusAccepted, resp)
}

func (s *Server) handleSubmitSparkJob(w http.ResponseWriter, r *http.Request) {
	idKey, err := requiredIdempotencyKey(r)
	if err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}
	var req service.SubmitSparkJobRequest
	if err := decodeJSON(r, &req); err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}
	resp, err := s.svc.SubmitSparkJob(r.Context(), req, idKey)
	if err != nil {
		writeServiceError(w, err)
		return
	}
	writeJSON(w, http.StatusAccepted, resp)
}

func (s *Server) handleGetSparkJobStatus(w http.ResponseWriter, r *http.Request) {
	status, err := s.svc.GetSparkJobStatusByOperation(r.Context(), chi.URLParam(r, "operationId"))
	if err != nil {
		writeServiceError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, status)
}

func (s *Server) handleDeleteSparkJob(w http.ResponseWriter, r *http.Request) {
	resp, err := s.svc.DeleteSparkJobByOperation(r.Context(), chi.URLParam(r, "operationId"))
	if err != nil {
		writeServiceError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, resp)
}

func (s *Server) handleGetSparkJobResult(w http.ResponseWriter, r *http.Request) {
	result, err := s.svc.GetSparkJobResultByOperation(r.Context(), chi.URLParam(r, "operationId"))
	if err != nil {
		writeServiceError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, result)
}

func (s *Server) handleCreateStarRocksCluster(w http.ResponseWriter, r *http.Request) {
	idKey, err := requiredIdempotencyKey(r)
	if err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}
	var req service.CreateStarRocksClusterRequest
	if err := decodeJSON(r, &req); err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}
	resp, err := s.svc.CreateStarRocksCluster(r.Context(), req, idKey)
	if err != nil {
		writeServiceError(w, err)
		return
	}
	writeJSON(w, http.StatusAccepted, resp)
}

func (s *Server) handleGetStarRocksCluster(w http.ResponseWriter, r *http.Request) {
	cluster, err := s.svc.GetStarRocksCluster(r.Context(), chi.URLParam(r, "clusterId"))
	if err != nil {
		writeServiceError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, cluster)
}

func (s *Server) handleGetStarRocksDirectEndpoint(w http.ResponseWriter, r *http.Request) {
	endpoint, err := s.svc.GetStarRocksDirectEndpoint(r.Context(), chi.URLParam(r, "clusterId"))
	if err != nil {
		writeServiceError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, endpoint)
}

func (s *Server) handleReleaseStarRocksCluster(w http.ResponseWriter, r *http.Request) {
	idKey, err := requiredIdempotencyKey(r)
	if err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}
	clusterID := chi.URLParam(r, "clusterId")
	resp, err := s.svc.ReleaseStarRocksCluster(r.Context(), clusterID, idKey)
	if err != nil {
		writeServiceError(w, err)
		return
	}
	writeJSON(w, http.StatusAccepted, resp)
}

func (s *Server) handleScaleOutStarRocksCluster(w http.ResponseWriter, r *http.Request) {
	s.handleResizeStarRocksCluster(w, r, true)
}

func (s *Server) handleScaleInStarRocksCluster(w http.ResponseWriter, r *http.Request) {
	s.handleResizeStarRocksCluster(w, r, false)
}

func (s *Server) handleResizeStarRocksCluster(w http.ResponseWriter, r *http.Request, scaleOut bool) {
	idKey, err := requiredIdempotencyKey(r)
	if err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}
	var req service.ResizeStarRocksClusterRequest
	if err := decodeJSON(r, &req); err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}
	clusterID := chi.URLParam(r, "clusterId")
	var resp *service.OperationResponse
	if scaleOut {
		resp, err = s.svc.ScaleOutStarRocksCluster(r.Context(), clusterID, req, idKey)
	} else {
		resp, err = s.svc.ScaleInStarRocksCluster(r.Context(), clusterID, req, idKey)
	}
	if err != nil {
		writeServiceError(w, err)
		return
	}
	writeJSON(w, http.StatusAccepted, resp)
}

func (s *Server) handleSubmitStarRocksJob(w http.ResponseWriter, r *http.Request) {
	idKey, err := requiredIdempotencyKey(r)
	if err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}
	var req service.SubmitStarRocksJobRequest
	if err := decodeJSON(r, &req); err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}
	resp, err := s.svc.SubmitStarRocksJob(r.Context(), req, idKey)
	if err != nil {
		writeServiceError(w, err)
		return
	}
	writeJSON(w, http.StatusAccepted, resp)
}

func (s *Server) handleGetStarRocksJobStatus(w http.ResponseWriter, r *http.Request) {
	status, err := s.svc.GetStarRocksJobStatusByOperation(r.Context(), chi.URLParam(r, "operationId"))
	if err != nil {
		writeServiceError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, status)
}

func (s *Server) handleGetOperation(w http.ResponseWriter, r *http.Request) {
	op, err := s.svc.GetOperation(r.Context(), chi.URLParam(r, "operationId"))
	if err != nil {
		writeServiceError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, op)
}

func requiredIdempotencyKey(r *http.Request) (string, error) {
	v := strings.TrimSpace(r.Header.Get("Idempotency-Key"))
	if v == "" {
		return "", errors.New("Idempotency-Key header is required")
	}
	return v, nil
}

func decodeJSON(r *http.Request, out interface{}) error {
	defer r.Body.Close()
	dec := json.NewDecoder(r.Body)
	dec.DisallowUnknownFields()
	if err := dec.Decode(out); err != nil {
		return fmt.Errorf("invalid json payload: %w", err)
	}
	return nil
}

func writeServiceError(w http.ResponseWriter, err error) {
	if apierrors.IsNotFound(err) {
		writeError(w, http.StatusNotFound, err)
		return
	}
	writeError(w, http.StatusBadRequest, err)
}

func writeError(w http.ResponseWriter, status int, err error) {
	writeJSON(w, status, map[string]string{"error": err.Error()})
}

func writeJSON(w http.ResponseWriter, status int, payload interface{}) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(payload)
}