Initial project scaffold
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Cursor
@@ -0,0 +1,222 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
)
|
||||
|
||||
type Permission string
|
||||
|
||||
const (
|
||||
PermissionPublic Permission = "public"
|
||||
PermissionBasic Permission = "basic"
|
||||
PermissionCreat Permission = "creat"
|
||||
PermissionPower Permission = "power"
|
||||
PermissionManager Permission = "manager"
|
||||
)
|
||||
|
||||
type User struct {
|
||||
ID string `json:"sub"`
|
||||
Username string `json:"username"`
|
||||
Roles []string `json:"role,omitempty"`
|
||||
TenantID string `json:"tenantId,omitempty"`
|
||||
SSOID string `json:"sso_id,omitempty"`
|
||||
APIKeyID string `json:"apiKeyId,omitempty"`
|
||||
APIKeySecret string `json:"apiKeySecret,omitempty"`
|
||||
APIKeyName string `json:"apiKeyName,omitempty"`
|
||||
}
|
||||
|
||||
type contextKey string
|
||||
|
||||
const userContextKey contextKey = "easyai-auth-user"
|
||||
|
||||
var ErrUnauthorized = errors.New("unauthorized")
|
||||
|
||||
type Authenticator struct {
|
||||
JWTSecret string
|
||||
ServerMainBaseURL string
|
||||
ServerMainInternalToken string
|
||||
HTTPClient *http.Client
|
||||
}
|
||||
|
||||
func New(jwtSecret string, serverMainBaseURL string, internalToken string) *Authenticator {
|
||||
return &Authenticator{
|
||||
JWTSecret: jwtSecret,
|
||||
ServerMainBaseURL: strings.TrimRight(serverMainBaseURL, "/"),
|
||||
ServerMainInternalToken: internalToken,
|
||||
HTTPClient: &http.Client{
|
||||
Timeout: 10 * time.Second,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func UserFromContext(ctx context.Context) (*User, bool) {
|
||||
user, ok := ctx.Value(userContextKey).(*User)
|
||||
return user, ok
|
||||
}
|
||||
|
||||
func (a *Authenticator) Require(permission Permission, next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
user, err := a.Authenticate(r)
|
||||
if err != nil {
|
||||
if permission == PermissionPublic {
|
||||
next.ServeHTTP(w, r)
|
||||
return
|
||||
}
|
||||
http.Error(w, "unauthorized", http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
if !hasPermission(user.Roles, permission) {
|
||||
http.Error(w, "forbidden", http.StatusForbidden)
|
||||
return
|
||||
}
|
||||
next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), userContextKey, user)))
|
||||
})
|
||||
}
|
||||
|
||||
func (a *Authenticator) Authenticate(r *http.Request) (*User, error) {
|
||||
token := extractBearer(r.Header.Get("Authorization"))
|
||||
if token == "" {
|
||||
token = strings.TrimSpace(r.Header.Get("x-comfy-api-key"))
|
||||
}
|
||||
if token == "" {
|
||||
return nil, ErrUnauthorized
|
||||
}
|
||||
if strings.HasPrefix(token, "sk-") {
|
||||
return a.verifyAPIKey(r.Context(), token)
|
||||
}
|
||||
return a.verifyJWT(token)
|
||||
}
|
||||
|
||||
func (a *Authenticator) verifyJWT(tokenString string) (*User, error) {
|
||||
token, err := jwt.ParseWithClaims(tokenString, jwt.MapClaims{}, func(token *jwt.Token) (any, error) {
|
||||
if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
|
||||
return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
|
||||
}
|
||||
return []byte(a.JWTSecret), nil
|
||||
})
|
||||
if err != nil || !token.Valid {
|
||||
return nil, ErrUnauthorized
|
||||
}
|
||||
claims, ok := token.Claims.(jwt.MapClaims)
|
||||
if !ok {
|
||||
return nil, ErrUnauthorized
|
||||
}
|
||||
|
||||
user := &User{
|
||||
ID: stringClaim(claims, "sub"),
|
||||
Username: stringClaim(claims, "username"),
|
||||
Roles: stringSliceClaim(claims, "role"),
|
||||
TenantID: stringClaim(claims, "tenantId"),
|
||||
SSOID: stringClaim(claims, "sso_id"),
|
||||
APIKeyID: stringClaim(claims, "apiKeyId"),
|
||||
APIKeySecret: stringClaim(claims, "apiKeySecret"),
|
||||
APIKeyName: stringClaim(claims, "apiKeyName"),
|
||||
}
|
||||
if user.ID == "" {
|
||||
return nil, ErrUnauthorized
|
||||
}
|
||||
return user, nil
|
||||
}
|
||||
|
||||
func (a *Authenticator) verifyAPIKey(ctx context.Context, apiKey string) (*User, error) {
|
||||
if a.ServerMainBaseURL == "" || a.ServerMainInternalToken == "" {
|
||||
return nil, ErrUnauthorized
|
||||
}
|
||||
body, _ := json.Marshal(map[string]string{"apiKey": apiKey})
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, a.ServerMainBaseURL+"/internal/platform/auth/verify-api-key", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+a.ServerMainInternalToken)
|
||||
|
||||
resp, err := a.HTTPClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, ErrUnauthorized
|
||||
}
|
||||
var user User
|
||||
if err := json.NewDecoder(resp.Body).Decode(&user); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if user.ID == "" {
|
||||
return nil, ErrUnauthorized
|
||||
}
|
||||
return &user, nil
|
||||
}
|
||||
|
||||
func extractBearer(value string) string {
|
||||
fields := strings.Fields(value)
|
||||
if len(fields) == 2 && strings.EqualFold(fields[0], "bearer") {
|
||||
return fields[1]
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func hasPermission(roles []string, required Permission) bool {
|
||||
if required == PermissionPublic {
|
||||
return true
|
||||
}
|
||||
granted := map[Permission]bool{PermissionPublic: true}
|
||||
for _, role := range roles {
|
||||
for _, permission := range permissionsForRole(role) {
|
||||
granted[permission] = true
|
||||
}
|
||||
}
|
||||
return granted[required]
|
||||
}
|
||||
|
||||
func permissionsForRole(role string) []Permission {
|
||||
switch role {
|
||||
case "admin", "manager":
|
||||
return []Permission{PermissionPublic, PermissionBasic, PermissionCreat, PermissionPower, PermissionManager}
|
||||
case "operator":
|
||||
return []Permission{PermissionPublic, PermissionBasic, PermissionCreat, PermissionPower}
|
||||
case "creator":
|
||||
return []Permission{PermissionPublic, PermissionBasic, PermissionCreat}
|
||||
case "user":
|
||||
return []Permission{PermissionPublic, PermissionBasic}
|
||||
default:
|
||||
return []Permission{PermissionPublic}
|
||||
}
|
||||
}
|
||||
|
||||
func stringClaim(claims jwt.MapClaims, key string) string {
|
||||
value, _ := claims[key].(string)
|
||||
return value
|
||||
}
|
||||
|
||||
func stringSliceClaim(claims jwt.MapClaims, key string) []string {
|
||||
value := claims[key]
|
||||
switch typed := value.(type) {
|
||||
case []string:
|
||||
return typed
|
||||
case []any:
|
||||
out := make([]string, 0, len(typed))
|
||||
for _, item := range typed {
|
||||
if s, ok := item.(string); ok && s != "" {
|
||||
out = append(out, s)
|
||||
}
|
||||
}
|
||||
return out
|
||||
case string:
|
||||
if typed == "" {
|
||||
return nil
|
||||
}
|
||||
return []string{typed}
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type Config struct {
|
||||
AppEnv string
|
||||
HTTPAddr string
|
||||
DatabaseURL string
|
||||
JWTSecret string
|
||||
ServerMainBaseURL string
|
||||
ServerMainInternalToken string
|
||||
CORSAllowedOrigin string
|
||||
LogLevel slog.Level
|
||||
}
|
||||
|
||||
func Load() Config {
|
||||
return Config{
|
||||
AppEnv: env("APP_ENV", "development"),
|
||||
HTTPAddr: env("HTTP_ADDR", ":8088"),
|
||||
DatabaseURL: gatewayDatabaseURL(),
|
||||
JWTSecret: env("CONFIG_JWT_SECRET", "this is a very secret secret"),
|
||||
ServerMainBaseURL: strings.TrimRight(
|
||||
env("SERVER_MAIN_BASE_URL", "http://localhost:3000"),
|
||||
"/",
|
||||
),
|
||||
ServerMainInternalToken: env("SERVER_MAIN_INTERNAL_TOKEN", ""),
|
||||
CORSAllowedOrigin: env("CORS_ALLOWED_ORIGIN", "http://localhost:5178"),
|
||||
LogLevel: logLevel(env("LOG_LEVEL", "info")),
|
||||
}
|
||||
}
|
||||
|
||||
func gatewayDatabaseURL() string {
|
||||
if value := envValue("AI_GATEWAY_DATABASE_URL"); value != "" {
|
||||
return normalizePostgresURL(value)
|
||||
}
|
||||
if value := envValue("DATABASE_URL"); value != "" {
|
||||
return normalizePostgresURL(value)
|
||||
}
|
||||
if memoryURL := envValue("MEMORY_DATABASE_URL"); memoryURL != "" {
|
||||
return normalizePostgresURL(withDatabase(memoryURL, env("AI_GATEWAY_DATABASE_NAME", "easyai_ai_gateway")))
|
||||
}
|
||||
return normalizePostgresURL("postgresql://easyai:easyai2025@localhost:5432/easyai_ai_gateway?sslmode=disable")
|
||||
}
|
||||
|
||||
func normalizePostgresURL(raw string) string {
|
||||
parsed, err := url.Parse(raw)
|
||||
if err != nil {
|
||||
return raw
|
||||
}
|
||||
values := parsed.Query()
|
||||
schema := values.Get("schema")
|
||||
if schema == "" {
|
||||
return raw
|
||||
}
|
||||
values.Del("schema")
|
||||
if values.Get("search_path") == "" {
|
||||
values.Set("search_path", schema)
|
||||
}
|
||||
parsed.RawQuery = values.Encode()
|
||||
return parsed.String()
|
||||
}
|
||||
|
||||
func withDatabase(raw string, databaseName string) string {
|
||||
parsed, err := url.Parse(raw)
|
||||
if err != nil || databaseName == "" {
|
||||
return raw
|
||||
}
|
||||
parsed.Path = "/" + databaseName
|
||||
return parsed.String()
|
||||
}
|
||||
|
||||
func envValue(key string) string {
|
||||
return strings.TrimSpace(os.Getenv(key))
|
||||
}
|
||||
|
||||
func env(key string, fallback string) string {
|
||||
if value := envValue(key); value != "" {
|
||||
return value
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
func logLevel(value string) slog.Level {
|
||||
switch strings.ToLower(value) {
|
||||
case "debug":
|
||||
return slog.LevelDebug
|
||||
case "warn", "warning":
|
||||
return slog.LevelWarn
|
||||
case "error":
|
||||
return slog.LevelError
|
||||
default:
|
||||
return slog.LevelInfo
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,218 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func (s *Server) health(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, http.StatusOK, map[string]any{
|
||||
"ok": true,
|
||||
"service": "easyai-ai-gateway",
|
||||
"env": s.cfg.AppEnv,
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) ready(w http.ResponseWriter, r *http.Request) {
|
||||
if err := s.store.Ping(r.Context()); err != nil {
|
||||
writeError(w, http.StatusServiceUnavailable, "postgres unavailable")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{"ok": true})
|
||||
}
|
||||
|
||||
func (s *Server) me(w http.ResponseWriter, r *http.Request) {
|
||||
user, _ := auth.UserFromContext(r.Context())
|
||||
writeJSON(w, http.StatusOK, user)
|
||||
}
|
||||
|
||||
func (s *Server) listPlatforms(w http.ResponseWriter, r *http.Request) {
|
||||
platforms, err := s.store.ListPlatforms(r.Context())
|
||||
if err != nil {
|
||||
s.logger.Error("list platforms failed", "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "list platforms failed")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{"items": platforms})
|
||||
}
|
||||
|
||||
func (s *Server) createPlatform(w http.ResponseWriter, r *http.Request) {
|
||||
var input store.CreatePlatformInput
|
||||
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid json body")
|
||||
return
|
||||
}
|
||||
if input.Provider == "" || input.Name == "" {
|
||||
writeError(w, http.StatusBadRequest, "provider and name are required")
|
||||
return
|
||||
}
|
||||
if input.AuthType == "" {
|
||||
input.AuthType = "bearer"
|
||||
}
|
||||
platform, err := s.store.CreatePlatform(r.Context(), input)
|
||||
if err != nil {
|
||||
s.logger.Error("create platform failed", "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "create platform failed")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusCreated, platform)
|
||||
}
|
||||
|
||||
func (s *Server) listModels(w http.ResponseWriter, r *http.Request) {
|
||||
models, err := s.store.ListModels(r.Context())
|
||||
if err != nil {
|
||||
s.logger.Error("list models failed", "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "list models failed")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{"items": models})
|
||||
}
|
||||
|
||||
func (s *Server) listCatalogProviders(w http.ResponseWriter, r *http.Request) {
|
||||
items, err := s.store.ListCatalogProviders(r.Context())
|
||||
if err != nil {
|
||||
s.logger.Error("list catalog providers failed", "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "list catalog providers failed")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{"items": items})
|
||||
}
|
||||
|
||||
func (s *Server) listBaseModels(w http.ResponseWriter, r *http.Request) {
|
||||
items, err := s.store.ListBaseModels(r.Context())
|
||||
if err != nil {
|
||||
s.logger.Error("list base models failed", "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "list base models failed")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{"items": items})
|
||||
}
|
||||
|
||||
func (s *Server) listPricingRules(w http.ResponseWriter, r *http.Request) {
|
||||
items, err := s.store.ListPricingRules(r.Context())
|
||||
if err != nil {
|
||||
s.logger.Error("list pricing rules failed", "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "list pricing rules failed")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{"items": items})
|
||||
}
|
||||
|
||||
func (s *Server) estimatePricing(w http.ResponseWriter, r *http.Request) {
|
||||
var body map[string]any
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid json body")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{
|
||||
"items": []any{},
|
||||
"resolver": "effective-pricing-placeholder",
|
||||
"request": body,
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) listRateLimitWindows(w http.ResponseWriter, r *http.Request) {
|
||||
items, err := s.store.ListRateLimitWindows(r.Context())
|
||||
if err != nil {
|
||||
s.logger.Error("list rate limit windows failed", "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "list rate limit windows failed")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{"items": items})
|
||||
}
|
||||
|
||||
func (s *Server) createTask(kind string) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
user, ok := auth.UserFromContext(r.Context())
|
||||
if !ok {
|
||||
writeError(w, http.StatusUnauthorized, "unauthorized")
|
||||
return
|
||||
}
|
||||
|
||||
var body map[string]any
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid json body")
|
||||
return
|
||||
}
|
||||
model, _ := body["model"].(string)
|
||||
if model == "" {
|
||||
writeError(w, http.StatusBadRequest, "model is required")
|
||||
return
|
||||
}
|
||||
|
||||
task, err := s.store.CreateTask(r.Context(), store.CreateTaskInput{
|
||||
Kind: kind,
|
||||
Model: model,
|
||||
Request: body,
|
||||
}, user)
|
||||
if err != nil {
|
||||
s.logger.Error("create task failed", "kind", kind, "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "create task failed")
|
||||
return
|
||||
}
|
||||
|
||||
writeJSON(w, http.StatusAccepted, map[string]any{
|
||||
"task": task,
|
||||
"next": map[string]string{
|
||||
"events": fmt.Sprintf("/api/v1/tasks/%s/events", task.ID),
|
||||
"detail": fmt.Sprintf("/api/v1/tasks/%s", task.ID),
|
||||
},
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) getTask(w http.ResponseWriter, r *http.Request) {
|
||||
task, err := s.store.GetTask(r.Context(), r.PathValue("taskID"))
|
||||
if err == nil {
|
||||
writeJSON(w, http.StatusOK, task)
|
||||
return
|
||||
}
|
||||
if store.IsNotFound(err) {
|
||||
writeError(w, http.StatusNotFound, "task not found")
|
||||
return
|
||||
}
|
||||
s.logger.Error("get task failed", "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "get task failed")
|
||||
}
|
||||
|
||||
func (s *Server) taskEvents(w http.ResponseWriter, r *http.Request) {
|
||||
task, err := s.store.GetTask(r.Context(), r.PathValue("taskID"))
|
||||
if err != nil {
|
||||
if store.IsNotFound(err) {
|
||||
writeError(w, http.StatusNotFound, "task not found")
|
||||
return
|
||||
}
|
||||
writeError(w, http.StatusInternalServerError, "get task failed")
|
||||
return
|
||||
}
|
||||
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Cache-Control", "no-cache")
|
||||
w.Header().Set("Connection", "keep-alive")
|
||||
|
||||
sendSSE(w, "task.accepted", map[string]any{
|
||||
"taskId": task.ID,
|
||||
"status": task.Status,
|
||||
})
|
||||
if flusher, ok := w.(http.Flusher); ok {
|
||||
flusher.Flush()
|
||||
}
|
||||
|
||||
timer := time.NewTimer(250 * time.Millisecond)
|
||||
defer timer.Stop()
|
||||
|
||||
select {
|
||||
case <-r.Context().Done():
|
||||
return
|
||||
case <-timer.C:
|
||||
sendSSE(w, "task.placeholder", map[string]any{
|
||||
"taskId": task.ID,
|
||||
"message": "runtime worker is not wired yet",
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
)
|
||||
|
||||
func writeJSON(w http.ResponseWriter, status int, value any) {
|
||||
w.Header().Set("Content-Type", "application/json; charset=utf-8")
|
||||
w.WriteHeader(status)
|
||||
_ = json.NewEncoder(w).Encode(value)
|
||||
}
|
||||
|
||||
func writeError(w http.ResponseWriter, status int, message string) {
|
||||
writeJSON(w, status, map[string]any{
|
||||
"error": map[string]any{
|
||||
"message": message,
|
||||
"status": status,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func sendSSE(w http.ResponseWriter, event string, payload any) {
|
||||
bytes, _ := json.Marshal(payload)
|
||||
_, _ = fmt.Fprintf(w, "event: %s\n", event)
|
||||
_, _ = fmt.Fprintf(w, "data: %s\n\n", bytes)
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
type Server struct {
|
||||
cfg config.Config
|
||||
store *store.Store
|
||||
auth *auth.Authenticator
|
||||
logger *slog.Logger
|
||||
}
|
||||
|
||||
func NewServer(cfg config.Config, db *store.Store, logger *slog.Logger) http.Handler {
|
||||
server := &Server{
|
||||
cfg: cfg,
|
||||
store: db,
|
||||
auth: auth.New(cfg.JWTSecret, cfg.ServerMainBaseURL, cfg.ServerMainInternalToken),
|
||||
logger: logger,
|
||||
}
|
||||
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("GET /healthz", server.health)
|
||||
mux.HandleFunc("GET /readyz", server.ready)
|
||||
|
||||
mux.Handle("GET /api/v1/me", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.me)))
|
||||
mux.Handle("GET /api/v1/catalog/providers", server.auth.Require(auth.PermissionPower, http.HandlerFunc(server.listCatalogProviders)))
|
||||
mux.Handle("GET /api/v1/catalog/base-models", server.auth.Require(auth.PermissionPower, http.HandlerFunc(server.listBaseModels)))
|
||||
mux.Handle("GET /api/v1/pricing/rules", server.auth.Require(auth.PermissionPower, http.HandlerFunc(server.listPricingRules)))
|
||||
mux.Handle("POST /api/v1/pricing/estimate", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.estimatePricing)))
|
||||
mux.Handle("GET /api/v1/platforms", server.auth.Require(auth.PermissionPower, http.HandlerFunc(server.listPlatforms)))
|
||||
mux.Handle("POST /api/v1/platforms", server.auth.Require(auth.PermissionPower, http.HandlerFunc(server.createPlatform)))
|
||||
mux.Handle("GET /api/v1/models", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.listModels)))
|
||||
mux.Handle("GET /api/v1/runtime/rate-limit-windows", server.auth.Require(auth.PermissionPower, http.HandlerFunc(server.listRateLimitWindows)))
|
||||
mux.Handle("POST /api/v1/chat/completions", server.auth.Require(auth.PermissionBasic, server.createTask("chat.completions")))
|
||||
mux.Handle("POST /api/v1/images/generations", server.auth.Require(auth.PermissionBasic, server.createTask("images.generations")))
|
||||
mux.Handle("POST /api/v1/videos/generations", server.auth.Require(auth.PermissionBasic, server.createTask("videos.generations")))
|
||||
mux.Handle("GET /api/v1/tasks/{taskID}", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.getTask)))
|
||||
mux.Handle("GET /api/v1/tasks/{taskID}/events", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.taskEvents)))
|
||||
|
||||
return server.recover(server.cors(mux))
|
||||
}
|
||||
|
||||
func (s *Server) cors(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
origin := r.Header.Get("Origin")
|
||||
if origin != "" && (s.cfg.CORSAllowedOrigin == "*" || strings.EqualFold(origin, s.cfg.CORSAllowedOrigin)) {
|
||||
w.Header().Set("Access-Control-Allow-Origin", origin)
|
||||
w.Header().Set("Vary", "Origin")
|
||||
w.Header().Set("Access-Control-Allow-Credentials", "true")
|
||||
w.Header().Set("Access-Control-Allow-Headers", "Authorization, Content-Type, X-Comfy-Api-Key")
|
||||
w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PATCH, DELETE, OPTIONS")
|
||||
}
|
||||
if r.Method == http.MethodOptions {
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
return
|
||||
}
|
||||
next.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) recover(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
defer func() {
|
||||
if err := recover(); err != nil {
|
||||
s.logger.Error("panic recovered", "error", err, "path", r.URL.Path)
|
||||
writeError(w, http.StatusInternalServerError, "internal server error")
|
||||
}
|
||||
}()
|
||||
next.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,508 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
type Store struct {
|
||||
pool *pgxpool.Pool
|
||||
}
|
||||
|
||||
func Connect(ctx context.Context, databaseURL string) (*Store, error) {
|
||||
pool, err := pgxpool.New(ctx, databaseURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := pool.Ping(ctx); err != nil {
|
||||
pool.Close()
|
||||
return nil, err
|
||||
}
|
||||
return &Store{pool: pool}, nil
|
||||
}
|
||||
|
||||
func (s *Store) Close() {
|
||||
s.pool.Close()
|
||||
}
|
||||
|
||||
func (s *Store) Ping(ctx context.Context) error {
|
||||
return s.pool.Ping(ctx)
|
||||
}
|
||||
|
||||
type Platform struct {
|
||||
ID string `json:"id"`
|
||||
Provider string `json:"provider"`
|
||||
PlatformKey string `json:"platformKey"`
|
||||
Name string `json:"name"`
|
||||
BaseURL string `json:"baseUrl,omitempty"`
|
||||
AuthType string `json:"authType"`
|
||||
Status string `json:"status"`
|
||||
Priority int `json:"priority"`
|
||||
DefaultPricingMode string `json:"defaultPricingMode"`
|
||||
DefaultDiscountFactor float64 `json:"defaultDiscountFactor"`
|
||||
Config map[string]any `json:"config,omitempty"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
|
||||
type CreatePlatformInput struct {
|
||||
Provider string `json:"provider"`
|
||||
PlatformKey string `json:"platformKey"`
|
||||
Name string `json:"name"`
|
||||
BaseURL string `json:"baseUrl"`
|
||||
AuthType string `json:"authType"`
|
||||
Credentials map[string]any `json:"credentials"`
|
||||
Config map[string]any `json:"config"`
|
||||
DefaultPricingMode string `json:"defaultPricingMode"`
|
||||
DefaultDiscountFactor float64 `json:"defaultDiscountFactor"`
|
||||
Priority int `json:"priority"`
|
||||
}
|
||||
|
||||
type PlatformModel struct {
|
||||
ID string `json:"id"`
|
||||
PlatformID string `json:"platformId"`
|
||||
BaseModelID string `json:"baseModelId,omitempty"`
|
||||
Provider string `json:"provider,omitempty"`
|
||||
PlatformName string `json:"platformName,omitempty"`
|
||||
ModelName string `json:"modelName"`
|
||||
ModelAlias string `json:"modelAlias,omitempty"`
|
||||
ModelType string `json:"modelType"`
|
||||
DisplayName string `json:"displayName"`
|
||||
CapabilityOverride map[string]any `json:"capabilityOverride,omitempty"`
|
||||
Capabilities map[string]any `json:"capabilities,omitempty"`
|
||||
PricingMode string `json:"pricingMode"`
|
||||
DiscountFactor float64 `json:"discountFactor,omitempty"`
|
||||
BillingConfigOverride map[string]any `json:"billingConfigOverride,omitempty"`
|
||||
BillingConfig map[string]any `json:"billingConfig,omitempty"`
|
||||
Enabled bool `json:"enabled"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
|
||||
type CatalogProvider struct {
|
||||
ID string `json:"id"`
|
||||
ProviderKey string `json:"providerKey"`
|
||||
DisplayName string `json:"displayName"`
|
||||
ProviderType string `json:"providerType"`
|
||||
CapabilitySchema map[string]any `json:"capabilitySchema,omitempty"`
|
||||
DefaultRateLimitPolicy map[string]any `json:"defaultRateLimitPolicy,omitempty"`
|
||||
Status string `json:"status"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
|
||||
type BaseModel struct {
|
||||
ID string `json:"id"`
|
||||
ProviderKey string `json:"providerKey"`
|
||||
CanonicalModelKey string `json:"canonicalModelKey"`
|
||||
ProviderModelName string `json:"providerModelName"`
|
||||
ModelType string `json:"modelType"`
|
||||
DisplayName string `json:"displayName"`
|
||||
Capabilities map[string]any `json:"capabilities,omitempty"`
|
||||
BaseBillingConfig map[string]any `json:"baseBillingConfig,omitempty"`
|
||||
DefaultRateLimitPolicy map[string]any `json:"defaultRateLimitPolicy,omitempty"`
|
||||
PricingVersion int `json:"pricingVersion"`
|
||||
Status string `json:"status"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
|
||||
type PricingRule struct {
|
||||
ID string `json:"id"`
|
||||
ScopeType string `json:"scopeType"`
|
||||
ScopeID string `json:"scopeId,omitempty"`
|
||||
ResourceType string `json:"resourceType"`
|
||||
Unit string `json:"unit"`
|
||||
BasePrice float64 `json:"basePrice"`
|
||||
Currency string `json:"currency"`
|
||||
BaseWeight map[string]any `json:"baseWeight,omitempty"`
|
||||
DynamicWeight map[string]any `json:"dynamicWeight,omitempty"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
|
||||
type RateLimitWindow struct {
|
||||
ScopeType string `json:"scopeType"`
|
||||
ScopeKey string `json:"scopeKey"`
|
||||
Metric string `json:"metric"`
|
||||
WindowStart time.Time `json:"windowStart"`
|
||||
LimitValue float64 `json:"limitValue"`
|
||||
UsedValue float64 `json:"usedValue"`
|
||||
ReservedValue float64 `json:"reservedValue"`
|
||||
ResetAt time.Time `json:"resetAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
|
||||
type CreateTaskInput struct {
|
||||
Kind string `json:"kind"`
|
||||
Model string `json:"model"`
|
||||
Request map[string]any `json:"request"`
|
||||
}
|
||||
|
||||
type GatewayTask struct {
|
||||
ID string `json:"id"`
|
||||
Kind string `json:"kind"`
|
||||
UserID string `json:"userId"`
|
||||
TenantID string `json:"tenantId,omitempty"`
|
||||
Model string `json:"model"`
|
||||
Request map[string]any `json:"request,omitempty"`
|
||||
Status string `json:"status"`
|
||||
Result map[string]any `json:"result,omitempty"`
|
||||
Billings []any `json:"billings,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
|
||||
func (s *Store) ListPlatforms(ctx context.Context) ([]Platform, error) {
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT id::text, provider, platform_key, name, COALESCE(base_url, ''), auth_type, status, priority,
|
||||
default_pricing_mode, default_discount_factor::float8, config, created_at, updated_at
|
||||
FROM integration_platforms
|
||||
ORDER BY priority ASC, created_at DESC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
platforms := make([]Platform, 0)
|
||||
for rows.Next() {
|
||||
var platform Platform
|
||||
var configBytes []byte
|
||||
if err := rows.Scan(
|
||||
&platform.ID,
|
||||
&platform.Provider,
|
||||
&platform.PlatformKey,
|
||||
&platform.Name,
|
||||
&platform.BaseURL,
|
||||
&platform.AuthType,
|
||||
&platform.Status,
|
||||
&platform.Priority,
|
||||
&platform.DefaultPricingMode,
|
||||
&platform.DefaultDiscountFactor,
|
||||
&configBytes,
|
||||
&platform.CreatedAt,
|
||||
&platform.UpdatedAt,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
platform.Config = decodeObject(configBytes)
|
||||
platforms = append(platforms, platform)
|
||||
}
|
||||
return platforms, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Store) CreatePlatform(ctx context.Context, input CreatePlatformInput) (Platform, error) {
|
||||
credentials, _ := json.Marshal(input.Credentials)
|
||||
config, _ := json.Marshal(input.Config)
|
||||
if input.DefaultPricingMode == "" {
|
||||
input.DefaultPricingMode = "inherit_discount"
|
||||
}
|
||||
if input.DefaultDiscountFactor == 0 {
|
||||
input.DefaultDiscountFactor = 1
|
||||
}
|
||||
if input.Priority == 0 {
|
||||
input.Priority = 100
|
||||
}
|
||||
var platform Platform
|
||||
var configBytes []byte
|
||||
err := s.pool.QueryRow(ctx, `
|
||||
INSERT INTO integration_platforms (provider, platform_key, name, base_url, auth_type, credentials, config, default_pricing_mode, default_discount_factor, priority)
|
||||
VALUES ($1, COALESCE(NULLIF($2, ''), 'platform_' || replace(gen_random_uuid()::text, '-', '')), $3, $4, $5, $6, $7, $8, $9, $10)
|
||||
RETURNING id::text, provider, platform_key, name, COALESCE(base_url, ''), auth_type, status, priority,
|
||||
default_pricing_mode, default_discount_factor::float8, config, created_at, updated_at`,
|
||||
input.Provider, input.PlatformKey, input.Name, input.BaseURL, input.AuthType, credentials, config, input.DefaultPricingMode, input.DefaultDiscountFactor, input.Priority,
|
||||
).Scan(
|
||||
&platform.ID,
|
||||
&platform.Provider,
|
||||
&platform.PlatformKey,
|
||||
&platform.Name,
|
||||
&platform.BaseURL,
|
||||
&platform.AuthType,
|
||||
&platform.Status,
|
||||
&platform.Priority,
|
||||
&platform.DefaultPricingMode,
|
||||
&platform.DefaultDiscountFactor,
|
||||
&configBytes,
|
||||
&platform.CreatedAt,
|
||||
&platform.UpdatedAt,
|
||||
)
|
||||
if err != nil {
|
||||
return Platform{}, err
|
||||
}
|
||||
platform.Config = decodeObject(configBytes)
|
||||
return platform, nil
|
||||
}
|
||||
|
||||
func (s *Store) ListModels(ctx context.Context) ([]PlatformModel, error) {
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT m.id::text, m.platform_id::text, COALESCE(m.base_model_id::text, ''), p.provider, p.name,
|
||||
m.model_name, COALESCE(m.model_alias, ''), m.model_type, m.display_name,
|
||||
m.capability_override, m.capabilities, m.pricing_mode, COALESCE(m.discount_factor, 0)::float8,
|
||||
m.billing_config_override, m.billing_config, m.enabled, m.created_at, m.updated_at
|
||||
FROM platform_models m
|
||||
JOIN integration_platforms p ON p.id = m.platform_id
|
||||
ORDER BY m.model_type ASC, m.model_name ASC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
models := make([]PlatformModel, 0)
|
||||
for rows.Next() {
|
||||
var model PlatformModel
|
||||
var capabilityOverride []byte
|
||||
var capabilities []byte
|
||||
var billingConfigOverride []byte
|
||||
var billingConfig []byte
|
||||
if err := rows.Scan(
|
||||
&model.ID,
|
||||
&model.PlatformID,
|
||||
&model.BaseModelID,
|
||||
&model.Provider,
|
||||
&model.PlatformName,
|
||||
&model.ModelName,
|
||||
&model.ModelAlias,
|
||||
&model.ModelType,
|
||||
&model.DisplayName,
|
||||
&capabilityOverride,
|
||||
&capabilities,
|
||||
&model.PricingMode,
|
||||
&model.DiscountFactor,
|
||||
&billingConfigOverride,
|
||||
&billingConfig,
|
||||
&model.Enabled,
|
||||
&model.CreatedAt,
|
||||
&model.UpdatedAt,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
model.CapabilityOverride = decodeObject(capabilityOverride)
|
||||
model.Capabilities = decodeObject(capabilities)
|
||||
model.BillingConfigOverride = decodeObject(billingConfigOverride)
|
||||
model.BillingConfig = decodeObject(billingConfig)
|
||||
models = append(models, model)
|
||||
}
|
||||
return models, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Store) ListCatalogProviders(ctx context.Context) ([]CatalogProvider, error) {
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT id::text, provider_key, display_name, provider_type, capability_schema,
|
||||
default_rate_limit_policy, status, created_at, updated_at
|
||||
FROM model_catalog_providers
|
||||
ORDER BY provider_key ASC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
items := make([]CatalogProvider, 0)
|
||||
for rows.Next() {
|
||||
var item CatalogProvider
|
||||
var capabilitySchema []byte
|
||||
var rateLimitPolicy []byte
|
||||
if err := rows.Scan(
|
||||
&item.ID,
|
||||
&item.ProviderKey,
|
||||
&item.DisplayName,
|
||||
&item.ProviderType,
|
||||
&capabilitySchema,
|
||||
&rateLimitPolicy,
|
||||
&item.Status,
|
||||
&item.CreatedAt,
|
||||
&item.UpdatedAt,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
item.CapabilitySchema = decodeObject(capabilitySchema)
|
||||
item.DefaultRateLimitPolicy = decodeObject(rateLimitPolicy)
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Store) ListBaseModels(ctx context.Context) ([]BaseModel, error) {
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT id::text, provider_key, canonical_model_key, provider_model_name, model_type, display_name,
|
||||
capabilities, base_billing_config, default_rate_limit_policy, pricing_version,
|
||||
status, created_at, updated_at
|
||||
FROM base_model_catalog
|
||||
ORDER BY provider_key ASC, model_type ASC, canonical_model_key ASC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
items := make([]BaseModel, 0)
|
||||
for rows.Next() {
|
||||
var item BaseModel
|
||||
var capabilities []byte
|
||||
var billingConfig []byte
|
||||
var rateLimitPolicy []byte
|
||||
if err := rows.Scan(
|
||||
&item.ID,
|
||||
&item.ProviderKey,
|
||||
&item.CanonicalModelKey,
|
||||
&item.ProviderModelName,
|
||||
&item.ModelType,
|
||||
&item.DisplayName,
|
||||
&capabilities,
|
||||
&billingConfig,
|
||||
&rateLimitPolicy,
|
||||
&item.PricingVersion,
|
||||
&item.Status,
|
||||
&item.CreatedAt,
|
||||
&item.UpdatedAt,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
item.Capabilities = decodeObject(capabilities)
|
||||
item.BaseBillingConfig = decodeObject(billingConfig)
|
||||
item.DefaultRateLimitPolicy = decodeObject(rateLimitPolicy)
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Store) ListPricingRules(ctx context.Context) ([]PricingRule, error) {
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT id::text, scope_type, COALESCE(scope_id::text, ''), resource_type, unit,
|
||||
base_price::float8, currency, base_weight, dynamic_weight, created_at, updated_at
|
||||
FROM model_pricing_rules
|
||||
ORDER BY scope_type ASC, resource_type ASC, created_at DESC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
items := make([]PricingRule, 0)
|
||||
for rows.Next() {
|
||||
var item PricingRule
|
||||
var baseWeight []byte
|
||||
var dynamicWeight []byte
|
||||
if err := rows.Scan(
|
||||
&item.ID,
|
||||
&item.ScopeType,
|
||||
&item.ScopeID,
|
||||
&item.ResourceType,
|
||||
&item.Unit,
|
||||
&item.BasePrice,
|
||||
&item.Currency,
|
||||
&baseWeight,
|
||||
&dynamicWeight,
|
||||
&item.CreatedAt,
|
||||
&item.UpdatedAt,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
item.BaseWeight = decodeObject(baseWeight)
|
||||
item.DynamicWeight = decodeObject(dynamicWeight)
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Store) ListRateLimitWindows(ctx context.Context) ([]RateLimitWindow, error) {
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT scope_type, scope_key, metric, window_start, limit_value::float8, used_value::float8,
|
||||
reserved_value::float8, reset_at, updated_at
|
||||
FROM gateway_rate_limit_counters
|
||||
WHERE reset_at >= now() - interval '5 minutes'
|
||||
ORDER BY window_start DESC, scope_type ASC, scope_key ASC, metric ASC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
items := make([]RateLimitWindow, 0)
|
||||
for rows.Next() {
|
||||
var item RateLimitWindow
|
||||
if err := rows.Scan(
|
||||
&item.ScopeType,
|
||||
&item.ScopeKey,
|
||||
&item.Metric,
|
||||
&item.WindowStart,
|
||||
&item.LimitValue,
|
||||
&item.UsedValue,
|
||||
&item.ReservedValue,
|
||||
&item.ResetAt,
|
||||
&item.UpdatedAt,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Store) CreateTask(ctx context.Context, input CreateTaskInput, user *auth.User) (GatewayTask, error) {
|
||||
requestBody, _ := json.Marshal(input.Request)
|
||||
var task GatewayTask
|
||||
var requestBytes []byte
|
||||
var resultBytes []byte
|
||||
var billingsBytes []byte
|
||||
err := s.pool.QueryRow(ctx, `
|
||||
INSERT INTO gateway_tasks (kind, user_id, tenant_id, model, request, status)
|
||||
VALUES ($1, $2, NULLIF($3, ''), $4, $5, 'queued')
|
||||
RETURNING id::text, kind, user_id, COALESCE(tenant_id, ''), model, request, status, result, billings, COALESCE(error, ''), created_at, updated_at`,
|
||||
input.Kind, user.ID, user.TenantID, input.Model, requestBody,
|
||||
).Scan(&task.ID, &task.Kind, &task.UserID, &task.TenantID, &task.Model, &requestBytes, &task.Status, &resultBytes, &billingsBytes, &task.Error, &task.CreatedAt, &task.UpdatedAt)
|
||||
if err != nil {
|
||||
return GatewayTask{}, err
|
||||
}
|
||||
task.Request = decodeObject(requestBytes)
|
||||
task.Result = decodeObject(resultBytes)
|
||||
task.Billings = decodeArray(billingsBytes)
|
||||
return task, nil
|
||||
}
|
||||
|
||||
func (s *Store) GetTask(ctx context.Context, taskID string) (GatewayTask, error) {
|
||||
var task GatewayTask
|
||||
var requestBytes []byte
|
||||
var resultBytes []byte
|
||||
var billingsBytes []byte
|
||||
err := s.pool.QueryRow(ctx, `
|
||||
SELECT id::text, kind, user_id, COALESCE(tenant_id, ''), model, request, status, result, billings, COALESCE(error, ''), created_at, updated_at
|
||||
FROM gateway_tasks
|
||||
WHERE id=$1`, taskID,
|
||||
).Scan(&task.ID, &task.Kind, &task.UserID, &task.TenantID, &task.Model, &requestBytes, &task.Status, &resultBytes, &billingsBytes, &task.Error, &task.CreatedAt, &task.UpdatedAt)
|
||||
if err != nil {
|
||||
return GatewayTask{}, err
|
||||
}
|
||||
task.Request = decodeObject(requestBytes)
|
||||
task.Result = decodeObject(resultBytes)
|
||||
task.Billings = decodeArray(billingsBytes)
|
||||
return task, nil
|
||||
}
|
||||
|
||||
func IsNotFound(err error) bool {
|
||||
return err == pgx.ErrNoRows
|
||||
}
|
||||
|
||||
func decodeObject(bytes []byte) map[string]any {
|
||||
if len(bytes) == 0 {
|
||||
return nil
|
||||
}
|
||||
var out map[string]any
|
||||
if err := json.Unmarshal(bytes, &out); err != nil {
|
||||
return nil
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func decodeArray(bytes []byte) []any {
|
||||
if len(bytes) == 0 {
|
||||
return nil
|
||||
}
|
||||
var out []any
|
||||
if err := json.Unmarshal(bytes, &out); err != nil {
|
||||
return nil
|
||||
}
|
||||
return out
|
||||
}
|
||||
Reference in New Issue
Block a user