mirror of
https://github.com/EZ-Api/ez-api.git
synced 2026-01-13 17:47:51 +00:00
2c5ccd56ee
Introduce StatsService integration to admin and master handlers, exposing realtime metrics (requests, tokens, QPS, rate limit status) via new endpoints: - GET /admin/masters/:id/realtime - GET /v1/realtime Also embed realtime stats in the existing GET /admin/masters/:id response and change GlobalQPS default to 0 with validation to reject negative values.
360 lines
13 KiB
Go
360 lines
13 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"expvar"
|
|
"log/slog"
|
|
"net/http"
|
|
"os"
|
|
"os/signal"
|
|
"strings"
|
|
"syscall"
|
|
"time"
|
|
|
|
_ "github.com/ez-api/ez-api/docs"
|
|
"github.com/ez-api/ez-api/internal/api"
|
|
"github.com/ez-api/ez-api/internal/config"
|
|
"github.com/ez-api/ez-api/internal/cron"
|
|
"github.com/ez-api/ez-api/internal/middleware"
|
|
"github.com/ez-api/ez-api/internal/model"
|
|
"github.com/ez-api/ez-api/internal/service"
|
|
"github.com/ez-api/foundation/logging"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/redis/go-redis/v9"
|
|
swaggerFiles "github.com/swaggo/files"
|
|
ginSwagger "github.com/swaggo/gin-swagger"
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
// @title EZ-API Control Plane
|
|
// @version 0.0.1
|
|
// @description Management API for EZ-API Gateway system.
|
|
// @termsOfService http://swagger.io/terms/
|
|
|
|
// @contact.name API Support
|
|
// @contact.url http://www.swagger.io/support
|
|
// @contact.email support@swagger.io
|
|
|
|
// @license.name Apache 2.0
|
|
// @license.url http://www.apache.org/licenses/LICENSE-2.0.html
|
|
|
|
// @host localhost:8080
|
|
// @BasePath /
|
|
|
|
// @securityDefinitions.apikey AdminAuth
|
|
// @in header
|
|
// @name Authorization
|
|
|
|
// @securityDefinitions.apikey MasterAuth
|
|
// @in header
|
|
// @name Authorization
|
|
|
|
func fatal(logger *slog.Logger, msg string, args ...any) {
|
|
logger.Error(msg, args...)
|
|
os.Exit(1)
|
|
}
|
|
|
|
func isOriginAllowed(allowed []string, origin string) bool {
|
|
if len(allowed) == 0 {
|
|
return false
|
|
}
|
|
for _, item := range allowed {
|
|
if item == "*" {
|
|
return true
|
|
}
|
|
if strings.EqualFold(strings.TrimSpace(item), strings.TrimSpace(origin)) {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func main() {
|
|
logger, _ := logging.New(logging.Options{Service: "ez-api"})
|
|
|
|
// 1. Load Configuration
|
|
cfg, err := config.Load()
|
|
if err != nil {
|
|
fatal(logger, "failed to load config", "err", err)
|
|
}
|
|
|
|
// 2. Initialize Redis Client
|
|
rdb := redis.NewClient(&redis.Options{
|
|
Addr: cfg.Redis.Addr,
|
|
Password: cfg.Redis.Password,
|
|
DB: cfg.Redis.DB,
|
|
})
|
|
|
|
// Verify Redis connection
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
if err := rdb.Ping(ctx).Err(); err != nil {
|
|
fatal(logger, "failed to connect to redis", "err", err)
|
|
}
|
|
logger.Info("connected to redis successfully")
|
|
|
|
// 3. Initialize GORM (PostgreSQL)
|
|
db, err := gorm.Open(postgres.Open(cfg.Postgres.DSN), &gorm.Config{})
|
|
if err != nil {
|
|
fatal(logger, "failed to connect to postgresql", "err", err)
|
|
}
|
|
sqlDB, err := db.DB()
|
|
if err != nil {
|
|
fatal(logger, "failed to get generic database object", "err", err)
|
|
}
|
|
// Verify DB connection
|
|
if err := sqlDB.Ping(); err != nil {
|
|
fatal(logger, "failed to ping postgresql", "err", err)
|
|
}
|
|
logger.Info("connected to postgresql successfully")
|
|
|
|
logDB := db
|
|
if cfg.Log.DSN != "" {
|
|
logDB, err = gorm.Open(postgres.Open(cfg.Log.DSN), &gorm.Config{})
|
|
if err != nil {
|
|
fatal(logger, "failed to connect to log database", "err", err)
|
|
}
|
|
sqlLogDB, err := logDB.DB()
|
|
if err != nil {
|
|
fatal(logger, "failed to get log database object", "err", err)
|
|
}
|
|
if err := sqlLogDB.Ping(); err != nil {
|
|
fatal(logger, "failed to ping log database", "err", err)
|
|
}
|
|
logger.Info("connected to log database successfully")
|
|
}
|
|
|
|
// Auto Migrate
|
|
if logDB != db {
|
|
if err := db.AutoMigrate(&model.Master{}, &model.Key{}, &model.Provider{}, &model.Model{}, &model.Binding{}, &model.Namespace{}, &model.OperationLog{}); err != nil {
|
|
fatal(logger, "failed to auto migrate", "err", err)
|
|
}
|
|
if err := logDB.AutoMigrate(&model.LogRecord{}); err != nil {
|
|
fatal(logger, "failed to auto migrate log tables", "err", err)
|
|
}
|
|
if err := service.EnsureLogIndexes(logDB); err != nil {
|
|
fatal(logger, "failed to ensure log indexes", "err", err)
|
|
}
|
|
} else {
|
|
if err := db.AutoMigrate(&model.Master{}, &model.Key{}, &model.Provider{}, &model.Model{}, &model.Binding{}, &model.Namespace{}, &model.OperationLog{}, &model.LogRecord{}); err != nil {
|
|
fatal(logger, "failed to auto migrate", "err", err)
|
|
}
|
|
if err := service.EnsureLogIndexes(db); err != nil {
|
|
fatal(logger, "failed to ensure log indexes", "err", err)
|
|
}
|
|
}
|
|
|
|
// 4. Setup Services and Handlers
|
|
syncService := service.NewSyncService(rdb)
|
|
logPartitioner := service.NewLogPartitioner(logDB, cfg.Log.Partitioning)
|
|
if logPartitioner.Enabled() {
|
|
if _, err := logPartitioner.EnsurePartitionFor(time.Now().UTC()); err != nil {
|
|
fatal(logger, "failed to ensure log partition", "err", err)
|
|
}
|
|
}
|
|
logWriter := service.NewLogWriter(logDB, cfg.Log.QueueCapacity, cfg.Log.BatchSize, cfg.Log.FlushInterval, logPartitioner)
|
|
logCtx, cancelLogs := context.WithCancel(context.Background())
|
|
defer cancelLogs()
|
|
logWriter.Start(logCtx)
|
|
quotaResetter := cron.NewQuotaResetter(db, syncService, time.Duration(cfg.Quota.ResetIntervalSeconds)*time.Second)
|
|
quotaCtx, cancelQuota := context.WithCancel(context.Background())
|
|
defer cancelQuota()
|
|
go quotaResetter.Start(quotaCtx)
|
|
logCleaner := cron.NewLogCleaner(logDB, rdb, cfg.Log.RetentionDays, int64(cfg.Log.MaxRecords), time.Hour, logPartitioner)
|
|
cleanerCtx, cancelCleaner := context.WithCancel(context.Background())
|
|
defer cancelCleaner()
|
|
go logCleaner.Start(cleanerCtx)
|
|
|
|
adminService, err := service.NewAdminService()
|
|
if err != nil {
|
|
fatal(logger, "failed to create admin service", "err", err)
|
|
}
|
|
masterService := service.NewMasterService(db)
|
|
statsService := service.NewStatsService(rdb)
|
|
healthService := service.NewHealthCheckService(db, rdb)
|
|
healthHandler := api.NewHealthHandler(healthService)
|
|
|
|
handler := api.NewHandler(db, logDB, syncService, logWriter, rdb, logPartitioner)
|
|
adminHandler := api.NewAdminHandler(db, logDB, masterService, syncService, statsService, logPartitioner)
|
|
masterHandler := api.NewMasterHandler(db, logDB, masterService, syncService, statsService, logPartitioner)
|
|
internalHandler := api.NewInternalHandler(db)
|
|
featureHandler := api.NewFeatureHandler(rdb)
|
|
modelRegistryService := service.NewModelRegistryService(db, rdb, service.ModelRegistryConfig{
|
|
Enabled: cfg.ModelRegistry.Enabled,
|
|
RefreshEvery: time.Duration(cfg.ModelRegistry.RefreshSeconds) * time.Second,
|
|
ModelsDevBaseURL: cfg.ModelRegistry.ModelsDevBaseURL,
|
|
ModelsDevAPIBaseURL: cfg.ModelRegistry.ModelsDevAPIBaseURL,
|
|
ModelsDevRef: cfg.ModelRegistry.ModelsDevRef,
|
|
CacheDir: cfg.ModelRegistry.CacheDir,
|
|
Timeout: time.Duration(cfg.ModelRegistry.TimeoutSeconds) * time.Second,
|
|
})
|
|
modelRegistryHandler := api.NewModelRegistryHandler(modelRegistryService)
|
|
|
|
// 4.1 Prime Redis snapshots so DP can start with data
|
|
if err := syncService.SyncAll(db); err != nil {
|
|
logger.Warn("initial sync warning", "err", err)
|
|
}
|
|
modelRegistryService.Start(context.Background())
|
|
|
|
// 5. Setup Gin Router
|
|
r := gin.Default()
|
|
r.Use(middleware.RequestID())
|
|
|
|
allowedOrigins := cfg.CORS.AllowOrigins
|
|
allowAllOrigins := isOriginAllowed(allowedOrigins, "*")
|
|
|
|
// CORS Middleware
|
|
r.Use(func(c *gin.Context) {
|
|
origin := c.Request.Header.Get("Origin")
|
|
if allowAllOrigins {
|
|
c.Writer.Header().Set("Access-Control-Allow-Origin", "*")
|
|
} else if origin != "" && isOriginAllowed(allowedOrigins, origin) {
|
|
c.Writer.Header().Set("Access-Control-Allow-Origin", origin)
|
|
c.Writer.Header().Add("Vary", "Origin")
|
|
}
|
|
c.Writer.Header().Set("Access-Control-Allow-Credentials", "true")
|
|
c.Writer.Header().Set("Access-Control-Allow-Headers", "Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization, accept, origin, Cache-Control, X-Requested-With")
|
|
c.Writer.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS, GET, PUT, DELETE")
|
|
|
|
if c.Request.Method == "OPTIONS" {
|
|
c.AbortWithStatus(204)
|
|
return
|
|
}
|
|
|
|
c.Next()
|
|
})
|
|
|
|
// Health Check Endpoint
|
|
r.GET("/health", func(c *gin.Context) {
|
|
status := healthService.Check(c.Request.Context())
|
|
httpStatus := http.StatusOK
|
|
if status.Status == "down" {
|
|
httpStatus = http.StatusServiceUnavailable
|
|
}
|
|
c.JSON(httpStatus, status)
|
|
})
|
|
r.GET("/api/status/test", healthHandler.TestDeps)
|
|
|
|
// Swagger Documentation
|
|
r.GET("/swagger/*any", ginSwagger.WrapHandler(swaggerFiles.Handler))
|
|
|
|
// API Routes
|
|
// Internal Routes
|
|
internalGroup := r.Group("/internal")
|
|
internalGroup.Use(middleware.InternalAuthMiddleware(cfg.Internal.StatsToken))
|
|
{
|
|
internalGroup.POST("/stats/flush", internalHandler.FlushStats)
|
|
internalGroup.GET("/metrics", gin.WrapH(expvar.Handler()))
|
|
}
|
|
|
|
// Admin Routes
|
|
adminGroup := r.Group("/admin")
|
|
adminGroup.Use(middleware.AdminAuthMiddleware(adminService))
|
|
adminGroup.Use(middleware.OperationLogMiddleware(db))
|
|
{
|
|
adminGroup.POST("/masters", adminHandler.CreateMaster)
|
|
adminGroup.GET("/masters", adminHandler.ListMasters)
|
|
adminGroup.GET("/masters/:id", adminHandler.GetMaster)
|
|
adminGroup.GET("/masters/:id/realtime", adminHandler.GetMasterRealtime)
|
|
adminGroup.PUT("/masters/:id", adminHandler.UpdateMaster)
|
|
adminGroup.DELETE("/masters/:id", adminHandler.DeleteMaster)
|
|
adminGroup.POST("/masters/batch", adminHandler.BatchMasters)
|
|
adminGroup.POST("/masters/:id/manage", adminHandler.ManageMaster)
|
|
adminGroup.POST("/masters/:id/keys", adminHandler.IssueChildKeyForMaster)
|
|
adminGroup.GET("/masters/:id/access", handler.GetMasterAccess)
|
|
adminGroup.PUT("/masters/:id/access", handler.UpdateMasterAccess)
|
|
adminGroup.GET("/keys/:id/access", handler.GetKeyAccess)
|
|
adminGroup.PUT("/keys/:id/access", handler.UpdateKeyAccess)
|
|
adminGroup.GET("/operation-logs", adminHandler.ListOperationLogs)
|
|
adminGroup.POST("/namespaces", handler.CreateNamespace)
|
|
adminGroup.GET("/namespaces", handler.ListNamespaces)
|
|
adminGroup.GET("/namespaces/:id", handler.GetNamespace)
|
|
adminGroup.PUT("/namespaces/:id", handler.UpdateNamespace)
|
|
adminGroup.DELETE("/namespaces/:id", handler.DeleteNamespace)
|
|
adminGroup.GET("/features", featureHandler.ListFeatures)
|
|
adminGroup.PUT("/features", featureHandler.UpdateFeatures)
|
|
adminGroup.GET("/model-registry/status", modelRegistryHandler.GetStatus)
|
|
adminGroup.POST("/model-registry/check", modelRegistryHandler.Check)
|
|
adminGroup.POST("/model-registry/refresh", modelRegistryHandler.Refresh)
|
|
adminGroup.POST("/model-registry/rollback", modelRegistryHandler.Rollback)
|
|
// Other admin routes for managing providers, models, etc.
|
|
adminGroup.POST("/providers", handler.CreateProvider)
|
|
adminGroup.GET("/providers", handler.ListProviders)
|
|
adminGroup.GET("/providers/:id", handler.GetProvider)
|
|
adminGroup.POST("/providers/preset", handler.CreateProviderPreset)
|
|
adminGroup.POST("/providers/custom", handler.CreateProviderCustom)
|
|
adminGroup.POST("/providers/google", handler.CreateProviderGoogle)
|
|
adminGroup.PUT("/providers/:id", handler.UpdateProvider)
|
|
adminGroup.DELETE("/providers/:id", handler.DeleteProvider)
|
|
adminGroup.POST("/providers/batch", handler.BatchProviders)
|
|
adminGroup.POST("/providers/:id/test", handler.TestProvider)
|
|
adminGroup.POST("/providers/:id/fetch-models", handler.FetchProviderModels)
|
|
adminGroup.POST("/models", handler.CreateModel)
|
|
adminGroup.GET("/models", handler.ListModels)
|
|
adminGroup.PUT("/models/:id", handler.UpdateModel)
|
|
adminGroup.DELETE("/models/:id", handler.DeleteModel)
|
|
adminGroup.POST("/models/batch", handler.BatchModels)
|
|
adminGroup.GET("/logs", handler.ListLogs)
|
|
adminGroup.DELETE("/logs", handler.DeleteLogs)
|
|
adminGroup.GET("/logs/stats", handler.LogStats)
|
|
adminGroup.GET("/logs/webhook", handler.GetLogWebhookConfig)
|
|
adminGroup.PUT("/logs/webhook", handler.UpdateLogWebhookConfig)
|
|
adminGroup.GET("/stats", adminHandler.GetAdminStats)
|
|
adminGroup.POST("/bindings", handler.CreateBinding)
|
|
adminGroup.GET("/bindings", handler.ListBindings)
|
|
adminGroup.GET("/bindings/:id", handler.GetBinding)
|
|
adminGroup.PUT("/bindings/:id", handler.UpdateBinding)
|
|
adminGroup.DELETE("/bindings/:id", handler.DeleteBinding)
|
|
adminGroup.POST("/bindings/batch", handler.BatchBindings)
|
|
adminGroup.POST("/sync/snapshot", handler.SyncSnapshot)
|
|
}
|
|
|
|
// Master Routes
|
|
masterGroup := r.Group("/v1")
|
|
masterGroup.Use(middleware.MasterAuthMiddleware(masterService))
|
|
{
|
|
masterGroup.GET("/self", masterHandler.GetSelf)
|
|
masterGroup.POST("/tokens", masterHandler.IssueChildKey)
|
|
masterGroup.GET("/tokens", masterHandler.ListTokens)
|
|
masterGroup.GET("/tokens/:id", masterHandler.GetToken)
|
|
masterGroup.PUT("/tokens/:id", masterHandler.UpdateToken)
|
|
masterGroup.DELETE("/tokens/:id", masterHandler.DeleteToken)
|
|
masterGroup.GET("/logs", masterHandler.ListSelfLogs)
|
|
masterGroup.GET("/logs/stats", masterHandler.GetSelfLogStats)
|
|
masterGroup.GET("/realtime", masterHandler.GetSelfRealtime)
|
|
masterGroup.GET("/stats", masterHandler.GetSelfStats)
|
|
}
|
|
|
|
// Public/General Routes (if any)
|
|
r.POST("/logs", handler.IngestLog)
|
|
|
|
srv := &http.Server{
|
|
Addr: ":" + cfg.Server.Port,
|
|
Handler: r,
|
|
}
|
|
|
|
// 6. Start Server with Graceful Shutdown
|
|
go func() {
|
|
logger.Info("starting ez-api", "port", cfg.Server.Port)
|
|
if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
|
fatal(logger, "server failed", "err", err)
|
|
}
|
|
}()
|
|
|
|
// Wait for interrupt signal
|
|
quit := make(chan os.Signal, 1)
|
|
signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
|
|
<-quit
|
|
logger.Info("shutting down server...")
|
|
|
|
// Shutdown with timeout
|
|
ctx, cancel = context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
if err := srv.Shutdown(ctx); err != nil {
|
|
fatal(logger, "server forced to shutdown", "err", err)
|
|
}
|
|
|
|
logger.Info("server exited properly")
|
|
}
|