From 6a24e8953fe165d37a854f435f18631aa3090580 Mon Sep 17 00:00:00 2001
From: CaIon <1808837298@qq.com>
Date: Wed, 10 Jan 2024 13:23:43 +0800
Subject: [PATCH] feat: able to fix channels

---
 controller/channel.go               | 16 +++++++++++
 model/ability.go                    | 43 +++++++++++++++++++++++++++++
 router/api-router.go                |  1 +
 web/src/components/ChannelsTable.js | 20 ++++++++++++++
 4 files changed, 80 insertions(+)

diff --git a/controller/channel.go b/controller/channel.go
index 4725bf0..dd71259 100644
--- a/controller/channel.go
+++ b/controller/channel.go
@@ -35,6 +35,22 @@ func GetAllChannels(c *gin.Context) {
 	return
 }
 
+func FixChannelsAbilities(c *gin.Context) {
+	count, err := model.FixAbility()
+	if err != nil {
+		c.JSON(http.StatusOK, gin.H{
+			"success": false,
+			"message": err.Error(),
+		})
+		return
+	}
+	c.JSON(http.StatusOK, gin.H{
+		"success": true,
+		"message": "",
+		"data":    count,
+	})
+}
+
 func SearchChannels(c *gin.Context) {
 	keyword := c.Query("keyword")
 	group := c.Query("group")
diff --git a/model/ability.go b/model/ability.go
index 0e5eb2c..7a2cc39 100644
--- a/model/ability.go
+++ b/model/ability.go
@@ -1,6 +1,7 @@
 package model
 
 import (
+	"fmt"
 	"one-api/common"
 	"strings"
 )
@@ -118,3 +119,45 @@ func (channel *Channel) UpdateAbilities() error {
 func UpdateAbilityStatus(channelId int, status bool) error {
 	return DB.Model(&Ability{}).Where("channel_id = ?", channelId).Select("enabled").Update("enabled", status).Error
 }
+
+func FixAbility() (int, error) {
+	var channelIds []int
+	count := 0
+	// Find all channel ids from channel table
+	err := DB.Model(&Channel{}).Pluck("id", &channelIds).Error
+	if err != nil {
+		common.SysError(fmt.Sprintf("Get channel ids from channel table failed: %s", err.Error()))
+		return 0, err
+	}
+	// Delete abilities of channels that are not in channel table
+	err = DB.Where("channel_id NOT IN (?)", channelIds).Delete(&Ability{}).Error
+	if err != nil {
+		common.SysError(fmt.Sprintf("Delete abilities of channels that are not in channel table failed: %s", err.Error()))
+		return 0, err
+	}
+	common.SysLog(fmt.Sprintf("Delete abilities of channels that are not in channel table successfully, ids: %v", channelIds))
+	count += len(channelIds)
+
+	// Use channelIds to find channel not in abilities table
+	var abilityChannelIds []int
+	err = DB.Model(&Ability{}).Pluck("channel_id", &abilityChannelIds).Error
+	if err != nil {
+		common.SysError(fmt.Sprintf("Get channel ids from abilities table failed: %s", err.Error()))
+		return 0, err
+	}
+	var channels []Channel
+	err = DB.Where("id NOT IN (?)", abilityChannelIds).Find(&channels).Error
+	if err != nil {
+		return 0, err
+	}
+	for _, channel := range channels {
+		err := channel.UpdateAbilities()
+		if err != nil {
+			common.SysError(fmt.Sprintf("Update abilities of channel %d failed: %s", channel.Id, err.Error()))
+		} else {
+			common.SysLog(fmt.Sprintf("Update abilities of channel %d successfully", channel.Id))
+			count++
+		}
+	}
+	return count, nil
+}
diff --git a/router/api-router.go b/router/api-router.go
index 3c3b999..04d3490 100644
--- a/router/api-router.go
+++ b/router/api-router.go
@@ -84,6 +84,7 @@ func SetApiRouter(router *gin.Engine) {
 			channelRoute.DELETE("/disabled", controller.DeleteDisabledChannel)
 			channelRoute.DELETE("/:id", controller.DeleteChannel)
 			channelRoute.POST("/batch", controller.DeleteChannelBatch)
+			channelRoute.POST("/fix", controller.FixChannelsAbilities)
 		}
 		tokenRoute := apiRouter.Group("/token")
 		tokenRoute.Use(middleware.UserAuth())
diff --git a/web/src/components/ChannelsTable.js b/web/src/components/ChannelsTable.js
index 0b07420..f194660 100644
--- a/web/src/components/ChannelsTable.js
+++ b/web/src/components/ChannelsTable.js
@@ -531,6 +531,17 @@ const ChannelsTable = () => {
         setLoading(false);
     }
 
+    const fixChannelsAbilities = async () => {
+        const res = await API.post(`/api/channel/fix`);
+        const {success, message, data} = res.data;
+        if (success) {
+            showSuccess(`已修复 ${data} 个通道！`);
+            await refresh();
+        } else {
+            showError(message);
+        }
+    }
+
     const sortChannel = (key) => {
         if (channels.length === 0) return;
         setLoading(true);
@@ -722,6 +733,15 @@ const ChannelsTable = () => {
                     >
                         <Button disabled={!enableBatchDelete} theme='light' type='danger' style={{marginRight: 8}}>删除所选通道</Button>
                     </Popconfirm>
+                    <Popconfirm
+                        title="确定是否要修复数据库一致性？"
+                        content="进行该操作时，可能导致渠道访问错误，请仅在数据库出现问题时使用"
+                        okType={'warning'}
+                        onConfirm={fixChannelsAbilities}
+                        position={'top'}
+                    >
+                        <Button theme='light' type='secondary' style={{marginRight: 8}}>修复数据库一致性</Button>
+                    </Popconfirm>
                 </Space>
             </div>
         </>