From acd1f1894c713c67c92921979bea1d284ef0f7fc Mon Sep 17 00:00:00 2001
From: zijiren233 <pyh1670605849@gmail.com>
Date: Sun, 12 Jan 2025 18:42:37 +0800
Subject: [PATCH] fix: db

---
 internal/db/update.go | 42 ++++++++++++++++++++++++++++++------------
 1 file changed, 30 insertions(+), 12 deletions(-)

diff --git a/internal/db/update.go b/internal/db/update.go
index 3d57231..3e6a62b 100644
--- a/internal/db/update.go
+++ b/internal/db/update.go
@@ -135,19 +135,37 @@ 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())
-		// 		}
-		// 	}()
-		// }
+		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:
+	case conf.DatabaseTypeSqlite3:
+		if err := db.Exec("PRAGMA foreign_keys = OFF").Error; err != nil {
+			return err
+		}
+		defer func() {
+			err := db.Exec("PRAGMA foreign_keys = ON").Error
+			if err != nil {
+				log.Fatalf("failed to set foreign key checks: %s", err.Error())
+			}
+		}()
+		return db.AutoMigrate(dst...)
+	case conf.DatabaseTypePostgres:
+		if err := db.Exec("SET CONSTRAINTS ALL DEFERRED").Error; err != nil {
+			return err
+		}
+		defer func() {
+			err := db.Exec("SET CONSTRAINTS ALL IMMEDIATE").Error
+			if err != nil {
+				log.Fatalf("failed to set foreign key checks: %s", err.Error())
+			}
+		}()
 		return db.AutoMigrate(dst...)
 	default:
 		return fmt.Errorf("unknown database type: %s", conf.Conf.Database.Type)