Files
ez-api/cmd/server/main.go
2025-12-21 00:53:52 +08:00

272 lines
9.6 KiB
Go

package main
import (
"context"
"log/slog"
"net/http"
"os"
"os/signal"
"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 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")
// Auto Migrate
if err := db.AutoMigrate(&model.Master{}, &model.Key{}, &model.Provider{}, &model.Model{}, &model.Binding{}, &model.LogRecord{}); err != nil {
fatal(logger, "failed to auto migrate", "err", err)
}
// 4. Setup Services and Handlers
syncService := service.NewSyncService(rdb)
logWriter := service.NewLogWriter(db, cfg.Log.QueueCapacity, cfg.Log.BatchSize, cfg.Log.FlushInterval)
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)
adminService, err := service.NewAdminService()
if err != nil {
fatal(logger, "failed to create admin service", "err", err)
}
masterService := service.NewMasterService(db)
healthService := service.NewHealthCheckService(db, rdb)
handler := api.NewHandler(db, syncService, logWriter)
adminHandler := api.NewAdminHandler(db, masterService, syncService)
masterHandler := api.NewMasterHandler(db, masterService, syncService)
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())
// CORS Middleware
r.Use(func(c *gin.Context) {
c.Writer.Header().Set("Access-Control-Allow-Origin", "*") // TODO: Restrict this in production
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)
})
// 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)
}
// Admin Routes
adminGroup := r.Group("/admin")
adminGroup.Use(middleware.AdminAuthMiddleware(adminService))
{
adminGroup.POST("/masters", adminHandler.CreateMaster)
adminGroup.GET("/masters", adminHandler.ListMasters)
adminGroup.GET("/masters/:id", adminHandler.GetMaster)
adminGroup.PUT("/masters/:id", adminHandler.UpdateMaster)
adminGroup.DELETE("/masters/:id", adminHandler.DeleteMaster)
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("/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/:id/test", handler.TestProvider)
adminGroup.POST("/models", handler.CreateModel)
adminGroup.GET("/models", handler.ListModels)
adminGroup.PUT("/models/:id", handler.UpdateModel)
adminGroup.GET("/logs", handler.ListLogs)
adminGroup.DELETE("/logs", handler.DeleteLogs)
adminGroup.GET("/logs/stats", handler.LogStats)
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("/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("/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")
}