From f7c80391169001e75737f773ea82d3eab344d341 Mon Sep 17 00:00:00 2001
From: HFO4 <912394456@qq.com>
Date: Sun, 6 Dec 2020 16:49:49 +0800
Subject: [PATCH] Feat: execute database script to calibrate user storage

---
 assets                    |  2 +-
 bootstrap/script.go       | 18 ++++++++++++++++++
 main.go                   | 12 ++++++++++--
 models/scripts/invoker.go | 25 +++++++++++++++++++++++++
 models/scripts/storage.go | 37 +++++++++++++++++++++++++++++++++++++
 5 files changed, 91 insertions(+), 3 deletions(-)
 create mode 100644 bootstrap/script.go
 create mode 100644 models/scripts/invoker.go
 create mode 100644 models/scripts/storage.go

diff --git a/assets b/assets
index 1bd0933..b473ad0 160000
--- a/assets
+++ b/assets
@@ -1 +1 @@
-Subproject commit 1bd093315526ac8585a8c7a3499e50f2a13d82fa
+Subproject commit b473ad0e45d4dcf4b090fbc155871009c084a92a
diff --git a/bootstrap/script.go b/bootstrap/script.go
new file mode 100644
index 0000000..9168dfa
--- /dev/null
+++ b/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)
+}
diff --git a/main.go b/main.go
index d1b4e53..2183de7 100644
--- a/main.go
+++ b/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
diff --git a/models/scripts/invoker.go b/models/scripts/invoker.go
new file mode 100644
index 0000000..af0155b
--- /dev/null
+++ b/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
+}
diff --git a/models/scripts/storage.go b/models/scripts/storage.go
new file mode 100644
index 0000000..9e152d5
--- /dev/null
+++ b/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)
+		}
+	}
+}