Feat: execute database script to calibrate user storage

bootstrap/script.go

@@ -0,0 +1,18 @@
+package bootstrap
+
+import (
+	"context"
+	"github.com/cloudreve/Cloudreve/v3/models/scripts"
+	"github.com/cloudreve/Cloudreve/v3/pkg/util"
+)
+
+func RunScript(name string) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	if err := scripts.RunDBScript(name, ctx); err != nil {
+		util.Log().Error("数据库脚本执行失败: %s", err)
+		return
+	}
+
+	util.Log().Info("数据库脚本 [%s] 执行完毕", name)
+}

main.go

@@ -10,13 +10,15 @@ import (
 )
 
 var (
-	isEject  bool
-	confPath string
+	isEject    bool
+	confPath   string
+	scriptName string
 )
 
 func init() {
 	flag.StringVar(&confPath, "c", util.RelativePath("conf.ini"), "配置文件路径")
 	flag.BoolVar(&isEject, "eject", false, "导出内置静态资源")
+	flag.StringVar(&scriptName, "database-script", "", "运行内置数据库助手脚本")
 	flag.Parse()
 	bootstrap.Init(confPath)
 }
@@ -28,6 +30,12 @@ func main() {
 		return
 	}
 
+	if scriptName != "" {
+		// 开始运行助手数据库脚本
+		bootstrap.RunScript(scriptName)
+		return
+	}
+
 	api := routers.InitRouter()
 
 	// 如果启用了SSL

models/scripts/invoker.go

@@ -0,0 +1,25 @@
+package scripts
+
+import (
+	"context"
+	"fmt"
+)
+
+type DBScript interface {
+	Run(ctx context.Context)
+}
+
+var availableScripts = make(map[string]DBScript)
+
+func RunDBScript(name string, ctx context.Context) error {
+	if script, ok := availableScripts[name]; ok {
+		script.Run(ctx)
+		return nil
+	}
+
+	return fmt.Errorf("数据库脚本 [%s] 不存在", name)
+}
+
+func register(name string, script DBScript) {
+	availableScripts[name] = script
+}

models/scripts/storage.go

@@ -0,0 +1,37 @@
+package scripts
+
+import (
+	"context"
+	model "github.com/cloudreve/Cloudreve/v3/models"
+	"github.com/cloudreve/Cloudreve/v3/pkg/util"
+)
+
+type UserStorageCalibration int
+
+func init() {
+	register("CalibrateUserStorage", UserStorageCalibration(0))
+}
+
+type storageResult struct {
+	Total uint64
+}
+
+// Run 运行脚本校准所有用户容量
+func (script UserStorageCalibration) Run(ctx context.Context) {
+	// 列出所有用户
+	var res []model.User
+	model.DB.Model(&model.User{}).Find(&res)
+
+	// 逐个检查容量
+	for _, user := range res {
+		// 计算正确的容量
+		var total storageResult
+		model.DB.Model(&model.File{}).Where("user_id = ?", user.ID).Select("sum(size) as total").Scan(&total)
+		// 更新用户的容量
+		if user.Storage != total.Total {
+			util.Log().Info("将用户 [%s] 的容量由 %d 校准为 %d", user.Email,
+				user.Storage, total.Total)
+			model.DB.Model(&user).Update("storage", total.Total)
+		}
+	}
+}