From c5871b9175bb04197b1d5b36b40f830665f3cb12 Mon Sep 17 00:00:00 2001
From: zijiren233 <pyh1670605849@gmail.com>
Date: Thu, 28 Dec 2023 16:06:52 +0800
Subject: [PATCH] Opt: db upgrade func

---
 internal/db/update.go | 54 ++++++++++++++++++-------------------------
 1 file changed, 23 insertions(+), 31 deletions(-)

diff --git a/internal/db/update.go b/internal/db/update.go
index 48e22dc..a52a972 100644
--- a/internal/db/update.go
+++ b/internal/db/update.go
@@ -38,37 +38,22 @@ var dbVersions = map[string]dbVersion{
 	"0.0.2": {
 		NextVersion: "0.0.3",
 		Upgrade: func(db *gorm.DB) error {
+			// alist and emby movies path are changed, so we need to delete them
+			_ = db.Exec("DELETE FROM movies WHERE base_vendor_info_vendor IN ('alist', 'emby')").Error
+			_ = db.Migrator().DropTable("alist_vendors", "emby_vendors")
 			// delete all vendors, because we are going to change the more vendor table, e.g. bilibili_vendors
 			return db.Migrator().DropTable("streaming_vendor_infos")
 		},
 	},
 	"0.0.3": {
 		NextVersion: "",
-		Upgrade: func(db *gorm.DB) error {
-			// alist and emby movies path are changed, so we need to delete them
-			_ = db.Exec("DELETE FROM movies WHERE base_vendor_info_vendor IN ('alist', 'emby')").Error
-			return db.Migrator().DropTable("alist_vendors", "emby_vendors")
-		},
 	},
 }
 
 func UpgradeDatabase() error {
-	if conf.Conf.Database.Type == conf.DatabaseTypeMysql {
-		if err := db.Exec("SET FOREIGN_KEY_CHECKS = 0").Error; err != nil {
-			return err
-		}
-		defer func() {
-			err := db.Exec("SET FOREIGN_KEY_CHECKS = 1").Error
-			if err != nil {
-				log.Fatalf("failed to set foreign key checks: %s", err.Error())
-			}
-		}()
-	}
-
 	if !db.Migrator().HasTable(&model.Setting{}) {
 		return autoMigrate(models...)
 	}
-
 	setting := model.Setting{
 		Name:  "database_version",
 		Type:  model.SettingTypeString,
@@ -88,28 +73,24 @@ func UpgradeDatabase() error {
 			}
 		}()
 	}
-
-	version, ok := dbVersions[currentVersion]
-	if !ok {
-		return nil
-	}
-	currentVersion = version.NextVersion
 	for currentVersion != "" {
 		version, ok := dbVersions[currentVersion]
 		if !ok {
 			break
 		}
-		log.Infof("Upgrading database to version %s", currentVersion)
-		if version.Upgrade != nil {
-			err := version.Upgrade(db)
+		if version.NextVersion != "" {
+			log.Infof("Upgrading database to version %s", version.NextVersion)
+			if version.Upgrade != nil {
+				err := version.Upgrade(db)
+				if err != nil {
+					return err
+				}
+			}
+			err := UpdateSettingItemValue("database_version", version.NextVersion)
 			if err != nil {
 				return err
 			}
 		}
-		err := UpdateSettingItemValue("database_version", currentVersion)
-		if err != nil {
-			return err
-		}
 		currentVersion = version.NextVersion
 	}
 	return nil
@@ -119,6 +100,17 @@ func autoMigrate(dst ...any) error {
 	log.Info("migrating database...")
 	switch conf.Conf.Database.Type {
 	case conf.DatabaseTypeMysql:
+		if conf.Conf.Database.Type == conf.DatabaseTypeMysql {
+			if err := db.Exec("SET FOREIGN_KEY_CHECKS = 0").Error; err != nil {
+				return err
+			}
+			defer func() {
+				err := db.Exec("SET FOREIGN_KEY_CHECKS = 1").Error
+				if err != nil {
+					log.Fatalf("failed to set foreign key checks: %s", err.Error())
+				}
+			}()
+		}
 		return db.Set("gorm:table_options", "ENGINE=InnoDB CHARSET=utf8mb4").AutoMigrate(dst...)
 	case conf.DatabaseTypeSqlite3, conf.DatabaseTypePostgres:
 		return db.AutoMigrate(dst...)