Sfoglia il codice sorgente

feat: add "/api/status/test"

[email protected] 1 anno fa
parent
commit
fe7f42fc2e
3 ha cambiato i file con 49 aggiunte e 0 eliminazioni
  1. 16 0
      controller/misc.go
  2. 32 0
      model/main.go
  3. 1 0
      router/api-router.go

+ 16 - 0
controller/misc.go

@@ -11,6 +11,22 @@ import (
 	"github.com/gin-gonic/gin"
 )
 
+func TestStatus(c *gin.Context) {
+	err := model.PingDB()
+	if err != nil {
+		c.JSON(http.StatusServiceUnavailable, gin.H{
+			"success": false,
+			"message": "数据库连接失败",
+		})
+		return
+	}
+	c.JSON(http.StatusOK, gin.H{
+		"success": true,
+		"message": "Server is running",
+	})
+	return
+}
+
 func GetStatus(c *gin.Context) {
 	c.JSON(http.StatusOK, gin.H{
 		"success": true,

+ 32 - 0
model/main.go

@@ -5,9 +5,11 @@ import (
 	"gorm.io/driver/postgres"
 	"gorm.io/driver/sqlite"
 	"gorm.io/gorm"
+	"log"
 	"one-api/common"
 	"os"
 	"strings"
+	"sync"
 	"time"
 )
 
@@ -148,3 +150,33 @@ func CloseDB() error {
 	err = sqlDB.Close()
 	return err
 }
+
+var (
+	lastPingTime time.Time
+	pingMutex    sync.Mutex
+)
+
+func PingDB() error {
+	pingMutex.Lock()
+	defer pingMutex.Unlock()
+
+	if time.Since(lastPingTime) < time.Second*10 {
+		return nil
+	}
+
+	sqlDB, err := DB.DB()
+	if err != nil {
+		log.Printf("Error getting sql.DB from GORM: %v", err)
+		return err
+	}
+
+	err = sqlDB.Ping()
+	if err != nil {
+		log.Printf("Error pinging DB: %v", err)
+		return err
+	}
+
+	lastPingTime = time.Now()
+	common.SysLog("Database pinged successfully")
+	return nil
+}

+ 1 - 0
router/api-router.go

@@ -14,6 +14,7 @@ func SetApiRouter(router *gin.Engine) {
 	apiRouter.Use(middleware.GlobalAPIRateLimit())
 	{
 		apiRouter.GET("/status", controller.GetStatus)
+		apiRouter.GET("/status/test", middleware.AdminAuth(), controller.TestStatus)
 		apiRouter.GET("/notice", controller.GetNotice)
 		apiRouter.GET("/about", controller.GetAbout)
 		apiRouter.GET("/midjourney", controller.GetMidjourney)