292 lines
8.1 KiB
Go
292 lines
8.1 KiB
Go
package httpapi
|
|
|
|
import (
|
|
"encoding/json"
|
|
"errors"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
|
|
"ai-service/internal/config"
|
|
"ai-service/internal/model"
|
|
"ai-service/internal/store"
|
|
|
|
"github.com/google/uuid"
|
|
)
|
|
|
|
type Server struct {
|
|
store *store.Store
|
|
cfg config.Config
|
|
}
|
|
|
|
func NewServer(store *store.Store, cfg config.Config) http.Handler {
|
|
return &Server{store: store, cfg: cfg}
|
|
}
|
|
|
|
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
|
path := strings.TrimSuffix(r.URL.Path, "/")
|
|
if path == "" {
|
|
path = "/"
|
|
}
|
|
switch {
|
|
case r.Method == http.MethodGet && path == "/healthz":
|
|
writeJSON(w, http.StatusOK, map[string]string{"status": "ok"})
|
|
case r.Method == http.MethodGet && path == "/readyz":
|
|
s.handleReady(w, r)
|
|
case r.Method == http.MethodGet && path == "/":
|
|
writeJSON(w, http.StatusOK, map[string]string{"service": "ai-service"})
|
|
case r.Method == http.MethodPost && path == "/api/v1/jobs":
|
|
s.handleCreateJob(w, r)
|
|
case r.Method == http.MethodPost && path == "/api/v1/jobs/batch":
|
|
s.handleCreateBatch(w, r)
|
|
case r.Method == http.MethodPost && path == "/api/v1/jobs/claim":
|
|
s.handleClaimJobs(w, r)
|
|
case r.Method == http.MethodGet && strings.HasPrefix(path, "/api/v1/jobs/"):
|
|
s.handleGetJob(w, r, path)
|
|
case r.Method == http.MethodPost && strings.HasPrefix(path, "/api/v1/jobs/") && strings.HasSuffix(path, "/retry"):
|
|
s.handleRetryJob(w, r, path)
|
|
case r.Method == http.MethodPost && strings.HasPrefix(path, "/api/v1/jobs/") && strings.HasSuffix(path, "/complete"):
|
|
s.handleCompleteJob(w, r, path)
|
|
case r.Method == http.MethodPost && strings.HasPrefix(path, "/api/v1/jobs/") && strings.HasSuffix(path, "/fail"):
|
|
s.handleFailJob(w, r, path)
|
|
case r.Method == http.MethodGet && path == "/api/v1/stats":
|
|
s.handleStats(w, r)
|
|
default:
|
|
writeError(w, http.StatusNotFound, "not found")
|
|
}
|
|
}
|
|
|
|
func (s *Server) handleReady(w http.ResponseWriter, r *http.Request) {
|
|
ctx, cancel := contextWithTimeout(r, 3*time.Second)
|
|
defer cancel()
|
|
if err := s.store.Ping(ctx); err != nil {
|
|
writeError(w, http.StatusServiceUnavailable, err.Error())
|
|
return
|
|
}
|
|
writeJSON(w, http.StatusOK, map[string]string{"status": "ready"})
|
|
}
|
|
|
|
func (s *Server) handleCreateJob(w http.ResponseWriter, r *http.Request) {
|
|
var req model.CreateJob
|
|
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
|
writeError(w, http.StatusBadRequest, "bad json")
|
|
return
|
|
}
|
|
ctx, cancel := contextWithTimeout(r, 8*time.Second)
|
|
defer cancel()
|
|
job, err := s.store.CreateJob(ctx, req)
|
|
if err != nil {
|
|
status := http.StatusInternalServerError
|
|
if isValidationError(err) {
|
|
status = http.StatusBadRequest
|
|
}
|
|
writeError(w, status, err.Error())
|
|
return
|
|
}
|
|
writeJSON(w, http.StatusCreated, job)
|
|
}
|
|
|
|
type createBatchRequest struct {
|
|
OwnerService string `json:"owner_service"`
|
|
TaskType string `json:"task_type"`
|
|
ModelProfile string `json:"model_profile"`
|
|
Priority int `json:"priority"`
|
|
MaxAttempts int `json:"max_attempts"`
|
|
Jobs []model.CreateJob `json:"jobs"`
|
|
}
|
|
|
|
type createBatchResponse struct {
|
|
Jobs []*model.Job `json:"jobs"`
|
|
}
|
|
|
|
func (s *Server) handleCreateBatch(w http.ResponseWriter, r *http.Request) {
|
|
var req createBatchRequest
|
|
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
|
writeError(w, http.StatusBadRequest, "bad json")
|
|
return
|
|
}
|
|
if len(req.Jobs) == 0 {
|
|
writeError(w, http.StatusBadRequest, "jobs is required")
|
|
return
|
|
}
|
|
ctx, cancel := contextWithTimeout(r, 20*time.Second)
|
|
defer cancel()
|
|
out := createBatchResponse{Jobs: make([]*model.Job, 0, len(req.Jobs))}
|
|
for _, item := range req.Jobs {
|
|
if item.OwnerService == "" {
|
|
item.OwnerService = req.OwnerService
|
|
}
|
|
if item.TaskType == "" {
|
|
item.TaskType = req.TaskType
|
|
}
|
|
if item.ModelProfile == "" {
|
|
item.ModelProfile = req.ModelProfile
|
|
}
|
|
if item.Priority == 0 {
|
|
item.Priority = req.Priority
|
|
}
|
|
if item.MaxAttempts == 0 {
|
|
item.MaxAttempts = req.MaxAttempts
|
|
}
|
|
job, err := s.store.CreateJob(ctx, item)
|
|
if err != nil {
|
|
status := http.StatusInternalServerError
|
|
if isValidationError(err) {
|
|
status = http.StatusBadRequest
|
|
}
|
|
writeError(w, status, err.Error())
|
|
return
|
|
}
|
|
out.Jobs = append(out.Jobs, job)
|
|
}
|
|
writeJSON(w, http.StatusCreated, out)
|
|
}
|
|
|
|
type claimJobsResponse struct {
|
|
Jobs []*model.Job `json:"jobs"`
|
|
}
|
|
|
|
func (s *Server) handleClaimJobs(w http.ResponseWriter, r *http.Request) {
|
|
var req model.ClaimJobs
|
|
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
|
writeError(w, http.StatusBadRequest, "bad json")
|
|
return
|
|
}
|
|
ctx, cancel := contextWithTimeout(r, 8*time.Second)
|
|
defer cancel()
|
|
jobs, err := s.store.ClaimJobs(ctx, req)
|
|
if err != nil {
|
|
writeError(w, http.StatusInternalServerError, err.Error())
|
|
return
|
|
}
|
|
writeJSON(w, http.StatusOK, claimJobsResponse{Jobs: jobs})
|
|
}
|
|
|
|
func (s *Server) handleGetJob(w http.ResponseWriter, r *http.Request, path string) {
|
|
id, err := jobIDFromPath(path, false)
|
|
if err != nil {
|
|
writeError(w, http.StatusBadRequest, err.Error())
|
|
return
|
|
}
|
|
ctx, cancel := contextWithTimeout(r, 8*time.Second)
|
|
defer cancel()
|
|
job, err := s.store.GetJob(ctx, id)
|
|
if err != nil {
|
|
writeError(w, http.StatusInternalServerError, err.Error())
|
|
return
|
|
}
|
|
if job == nil {
|
|
writeError(w, http.StatusNotFound, "job not found")
|
|
return
|
|
}
|
|
writeJSON(w, http.StatusOK, job)
|
|
}
|
|
|
|
func (s *Server) handleRetryJob(w http.ResponseWriter, r *http.Request, path string) {
|
|
id, err := jobIDFromPath(path, true)
|
|
if err != nil {
|
|
writeError(w, http.StatusBadRequest, err.Error())
|
|
return
|
|
}
|
|
ctx, cancel := contextWithTimeout(r, 8*time.Second)
|
|
defer cancel()
|
|
job, err := s.store.RetryJob(ctx, id)
|
|
if err != nil {
|
|
writeError(w, http.StatusInternalServerError, err.Error())
|
|
return
|
|
}
|
|
if job == nil {
|
|
writeError(w, http.StatusNotFound, "retryable job not found")
|
|
return
|
|
}
|
|
writeJSON(w, http.StatusOK, job)
|
|
}
|
|
|
|
func (s *Server) handleCompleteJob(w http.ResponseWriter, r *http.Request, path string) {
|
|
id, err := jobIDFromActionPath(path, "complete")
|
|
if err != nil {
|
|
writeError(w, http.StatusBadRequest, err.Error())
|
|
return
|
|
}
|
|
var req model.CompleteJob
|
|
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
|
writeError(w, http.StatusBadRequest, "bad json")
|
|
return
|
|
}
|
|
ctx, cancel := contextWithTimeout(r, 8*time.Second)
|
|
defer cancel()
|
|
job, err := s.store.CompleteJob(ctx, id, req)
|
|
if err != nil {
|
|
writeError(w, http.StatusInternalServerError, err.Error())
|
|
return
|
|
}
|
|
if job == nil {
|
|
writeError(w, http.StatusNotFound, "running job not found")
|
|
return
|
|
}
|
|
writeJSON(w, http.StatusOK, job)
|
|
}
|
|
|
|
func (s *Server) handleFailJob(w http.ResponseWriter, r *http.Request, path string) {
|
|
id, err := jobIDFromActionPath(path, "fail")
|
|
if err != nil {
|
|
writeError(w, http.StatusBadRequest, err.Error())
|
|
return
|
|
}
|
|
var req model.FailJob
|
|
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
|
writeError(w, http.StatusBadRequest, "bad json")
|
|
return
|
|
}
|
|
ctx, cancel := contextWithTimeout(r, 8*time.Second)
|
|
defer cancel()
|
|
job, err := s.store.FailJob(ctx, id, req)
|
|
if err != nil {
|
|
writeError(w, http.StatusInternalServerError, err.Error())
|
|
return
|
|
}
|
|
if job == nil {
|
|
writeError(w, http.StatusNotFound, "running job not found")
|
|
return
|
|
}
|
|
writeJSON(w, http.StatusOK, job)
|
|
}
|
|
|
|
func (s *Server) handleStats(w http.ResponseWriter, r *http.Request) {
|
|
ctx, cancel := contextWithTimeout(r, 8*time.Second)
|
|
defer cancel()
|
|
stats, err := s.store.Stats(ctx)
|
|
if err != nil {
|
|
writeError(w, http.StatusInternalServerError, err.Error())
|
|
return
|
|
}
|
|
writeJSON(w, http.StatusOK, stats)
|
|
}
|
|
|
|
func jobIDFromPath(path string, retry bool) (uuid.UUID, error) {
|
|
raw := strings.TrimPrefix(path, "/api/v1/jobs/")
|
|
if retry {
|
|
raw = strings.TrimSuffix(raw, "/retry")
|
|
}
|
|
id, err := uuid.Parse(strings.Trim(raw, "/"))
|
|
if err != nil {
|
|
return uuid.Nil, errors.New("bad job id")
|
|
}
|
|
return id, nil
|
|
}
|
|
|
|
func jobIDFromActionPath(path string, action string) (uuid.UUID, error) {
|
|
raw := strings.TrimPrefix(path, "/api/v1/jobs/")
|
|
raw = strings.TrimSuffix(raw, "/"+action)
|
|
id, err := uuid.Parse(strings.Trim(raw, "/"))
|
|
if err != nil {
|
|
return uuid.Nil, errors.New("bad job id")
|
|
}
|
|
return id, nil
|
|
}
|
|
|
|
func isValidationError(err error) bool {
|
|
msg := err.Error()
|
|
return strings.Contains(msg, " is required")
|
|
}
|