fix(db/migrator): remove foreign key constraints and better root_id…

… generation (#835)

Add function to delete all foreign key constraints when migration.

Leave relationship maintenance to the program and reduce the difficulty of database management. Because there are many different DBs and the implementation of foreign keys may be different, and the DB may not support foreign keys, so don't rely on the foreign key function of the DB system.

internal/dao/migrate.go

@@ -7,44 +7,85 @@ import (
 
 func (dao *Dao) MigrateModels() {
 	// Upgrade the database
-	dao.migrateRootID()
+	dao.MigrateRootID()
 
 	// Migrate the schema
 	dao.DB().AutoMigrate(&entity.Site{}, &entity.Page{}, &entity.User{},
-		&entity.Comment{}, &entity.Notify{}, &entity.Vote{}) // 注意表的创建顺序，因为有关联字段
+		&entity.Comment{}, &entity.Notify{}, &entity.Vote{})
+
+	// Delete all foreign key constraints
+	// Leave relationship maintenance to the program and reduce the difficulty of database management.
+	// because there are many different DBs and the implementation of foreign keys may be different,
+	// and the DB may not support foreign keys, so don't rely on the foreign key function of the DB system.
+	dao.DropConstraintsIfExist()
 }
 
-func (dao *Dao) migrateRootID() {
+// Remove all constraints
+func (dao *Dao) DropConstraintsIfExist() {
+	if dao.DB().Dialector.Name() == "sqlite" {
+		return // sqlite dose not support constraints by default
+	}
+
+	TAG := "[DB Migrator] "
+
+	list := []struct {
+		model      any
+		constraint string
+	}{
+		{&entity.Comment{}, "fk_comments_page"},
+		{&entity.Comment{}, "fk_comments_user"},
+		{&entity.Page{}, "fk_pages_site"},
+	}
+
+	for _, item := range list {
+		if dao.DB().Migrator().HasConstraint(item.model, item.constraint) {
+			log.Info(TAG, "Dropping constraint: ", item.constraint)
+			err := dao.DB().Migrator().DropConstraint(item.model, item.constraint)
+			if err != nil {
+				log.Fatal(TAG, "Failed to drop constraint: ", item.constraint)
+			}
+		}
+	}
+}
+
+func (dao *Dao) MigrateRootID() {
 	const TAG = "[DB Migrator] "
 
+	// if comments table does not exist which means the first time to setup Artalk
 	if !dao.DB().Migrator().HasTable(&entity.Comment{}) {
-		return
+		return // so no need to upgrade
 	}
+
+	// if root_id column already exists which means the migration has been done
 	if dao.DB().Migrator().HasColumn(&entity.Comment{}, "root_id") {
-		return
+		return // so no need to migrate again
 	}
+
+	log.Info(TAG, "Generating root IDs...")
+
 	dao.DB().Migrator().AddColumn(&entity.Comment{}, "root_id")
 
-	batchSize := 1000
-	var offset uint = 0
-	for {
-		var comments []entity.Comment
-		dao.DB().Limit(batchSize).Offset(int(offset)).Find(&comments)
+	err := dao.DB().Raw(`WITH RECURSIVE CommentHierarchy AS (
+		SELECT id, id AS root_id, rid
+		FROM comments
+		WHERE rid = 0
 
-		if len(comments) == 0 {
-			break
-		}
+		UNION ALL
 
-		for i := range comments {
-			if comments[i].Rid != 0 {
-				rootID := dao.FindCommentRootID(comments[i].Rid)
-				comments[i].RootID = rootID
-			}
-			dao.DB().Save(&comments[i])
-		}
+		SELECT c.id, ch.root_id, c.rid
+		FROM comments c
+		INNER JOIN CommentHierarchy ch ON c.rid = ch.id
+	)
+	UPDATE comments SET root_id = (
+		SELECT root_id
+		FROM CommentHierarchy
+		WHERE comments.id = CommentHierarchy.id
+	);
+	`).Scan(&struct{}{}).Error
 
-		offset += uint(batchSize)
-		log.Debug(TAG, "Processed ", offset, " comments")
+	if err != nil {
+		dao.DB().Migrator().DropColumn(&entity.Comment{}, "root_id")
+		log.Fatal(TAG, "Failed to generate root IDs, please feedback this issue to the Artalk team.")
 	}
 
 	log.Info(TAG, "Root IDs generated successfully.")

internal/db/db.go

@@ -15,6 +15,7 @@ func NewDB(conf config.DBConf) (*gorm.DB, error) {
 		NamingStrategy: schema.NamingStrategy{
 			TablePrefix: conf.TablePrefix,
 		},
+		DisableForeignKeyConstraintWhenMigrating: true,
 	}
 
 	// Enable Prepared Statement by default
@@ -46,7 +47,7 @@ func NewDB(conf config.DBConf) (*gorm.DB, error) {
 }
 
 func NewTestDB() (*gorm.DB, error) {
-	return OpenSQLite("file::memory:?cache=shared", &gorm.Config{})
+	return OpenSQLite("file::memory:?cache=shared", &gorm.Config{DisableForeignKeyConstraintWhenMigrating: true})
 }
 
 func CloseDB(db *gorm.DB) error {

internal/entity/comment.go

@@ -28,7 +28,7 @@ type Comment struct {
 
 	RootID uint  `gorm:"index"` // Root Node ID (can be derived from `Rid`)
 	Page   *Page `gorm:"foreignKey:page_key;references:key"`
-	User   *User `gorm:"foreignKey:id;references:user_id"`
+	User   *User `gorm:"foreignKey:user_id;references:id"`
 }
 
 func (c Comment) IsEmpty() bool {

internal/entity/page.go

@@ -6,7 +6,7 @@ import (
 
 type Page struct {
 	gorm.Model
-	Key       string `gorm:"index;size:255"` // 页面 Key（一般为不含 hash/query 的完整 url）
+	Key       string `gorm:"uniqueIndex;size:255"` // 页面 Key（一般为不含 hash/query 的完整 url）
 	Title     string
 	AdminOnly bool
 

server/handler/comment_create.go

@@ -201,19 +201,21 @@ func fetchIPRegionForComment(app *core.App, comment entity.CookedComment) entity
 }
 
 func isAllowComment(app *core.App, c *fiber.Ctx, name string, email string, pageAdminOnly bool) (bool, error) {
-	user, err := common.GetUserByReq(app, c)
-	if !errors.Is(err, common.ErrTokenNotProvided) && user.IsEmpty() {
-		return false, common.RespError(c, 401, i18n.T("Login required"), Map{"need_auth_login": true})
-	}
-
+	// if the user is an admin user or page is admin only
 	isAdminUser := app.Dao().IsAdminUserByNameEmail(name, email)
-
-	// 如果用户是管理员，或者当前页只能管理员评论
 	if isAdminUser || pageAdminOnly {
+		// then check has admin access
 		if !common.CheckIsAdminReq(app, c) {
 			return false, common.RespError(c, 403, i18n.T("Admin access required"), Map{"need_login": true})
 		}
 	}
 
+	// if token is provided, then check token is valid
+	user, err := common.GetUserByReq(app, c)
+	if !errors.Is(err, common.ErrTokenNotProvided) && user.IsEmpty() {
+		// need_auth_login is a hook for frontend to show login modal (new Auth api)
+		return false, common.RespError(c, 401, i18n.T("Login required"), Map{"need_auth_login": true})
+	}
+
 	return true, nil
 }

test/app.go

@@ -48,7 +48,8 @@ func NewTestApp() (*TestApp, error) {
 
 	// open a sqlite db
 	dbInstance, err := gorm.Open(sqlite.Open(dbFile), &gorm.Config{
-		Logger: db_logger.NewGormLogger(),
+		Logger:                                   db_logger.NewGormLogger(),
+		DisableForeignKeyConstraintWhenMigrating: true,
 	})
 	if err != nil {
 		return nil, err