User badge model fixes for RemoveUserBadges and AddUserBadges

Author: komarov

models/fixtures/badge.yml

@@ -0,0 +1,5 @@
+-
+  id: 1
+  slug: badge1
+  description: just a test badge
+  image_url: badge1.png

models/migrations/migrations.go

@@ -24,6 +24,7 @@ import (
 	"code.gitea.io/gitea/models/migrations/v1_22"
 	"code.gitea.io/gitea/models/migrations/v1_23"
 	"code.gitea.io/gitea/models/migrations/v1_24"
+	"code.gitea.io/gitea/models/migrations/v1_25"
 	"code.gitea.io/gitea/models/migrations/v1_6"
 	"code.gitea.io/gitea/models/migrations/v1_7"
 	"code.gitea.io/gitea/models/migrations/v1_8"
@@ -382,6 +383,9 @@ func prepareMigrationTasks() []*migration {
 		newMigration(318, "Add anonymous_access_mode for repo_unit", v1_24.AddRepoUnitAnonymousAccessMode),
 		newMigration(319, "Add ExclusiveOrder to Label table", v1_24.AddExclusiveOrderColumnToLabelTable),
 		newMigration(320, "Migrate two_factor_policy to login_source table", v1_24.MigrateSkipTwoFactor),
+
+		// Gitea 1.24.0 ends at database version 321 (database version 322)
+		newMigration(321, "Add unique constraint for user badge", v1_25.AddUniqueIndexForUserBadge),
 	}
 	return preparedMigrations
 }

models/migrations/v1_25/v321.go

@@ -0,0 +1,62 @@
+// Copyright 2025 The Gitea Authors. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+package v1_25 //nolint
+
+import (
+	"fmt"
+
+	"xorm.io/xorm"
+	"xorm.io/xorm/schemas"
+)
+
+type UserBadge struct { //revive:disable-line:exported
+	ID      int64 `xorm:"pk autoincr"`
+	BadgeID int64
+	UserID  int64
+}
+
+// TableIndices implements xorm's TableIndices interface
+func (n *UserBadge) TableIndices() []*schemas.Index {
+	indices := make([]*schemas.Index, 0, 1)
+	ubUnique := schemas.NewIndex("unique_user_badge", schemas.UniqueType)
+	ubUnique.AddColumn("user_id", "badge_id")
+	indices = append(indices, ubUnique)
+	return indices
+}
+
+// AddUniqueIndexForUserBadge adds a compound unique indexes for user badge table
+// and it replaces an old index on user_id
+func AddUniqueIndexForUserBadge(x *xorm.Engine) error {
+	// remove possible duplicated records in table user_badge
+	type result struct {
+		UserID  int64
+		BadgeID int64
+		Cnt     int
+	}
+	var results []result
+	if err := x.Select("user_id, badge_id, count(*) as cnt").
+		Table("user_badge").
+		GroupBy("user_id, badge_id").
+		Having("count(*) > 1").
+		Find(&results); err != nil {
+		return err
+	}
+	for _, r := range results {
+		if x.Dialect().URI().DBType == schemas.MSSQL {
+			if _, err := x.Exec(fmt.Sprintf("delete from user_badge where id in (SELECT top %d id FROM user_badge WHERE user_id = ? and badge_id = ?)", r.Cnt-1), r.UserID, r.BadgeID); err != nil {
+				return err
+			}
+		} else {
+			var ids []int64
+			if err := x.SQL("SELECT id FROM user_badge WHERE user_id = ? and badge_id = ? limit ?", r.UserID, r.BadgeID, r.Cnt-1).Find(&ids); err != nil {
+				return err
+			}
+			if _, err := x.Table("user_badge").In("id", ids).Delete(); err != nil {
+				return err
+			}
+		}
+	}
+
+	return x.Sync(new(UserBadge))
+}

models/user/badge.go

@@ -12,6 +12,7 @@ import (
 	"code.gitea.io/gitea/modules/util"
 
 	"xorm.io/builder"
+	"xorm.io/xorm/schemas"
 )
 
 // Badge represents a user badge
@@ -29,6 +30,15 @@ type UserBadge struct { //nolint:revive // export stutter
 	UserID  int64 `xorm:"INDEX"`
 }
 
+// TableIndices implements xorm's TableIndices interface
+func (n *UserBadge) TableIndices() []*schemas.Index {
+	indices := make([]*schemas.Index, 0, 1)
+	ubUnique := schemas.NewIndex("unique_user_badge", schemas.UniqueType)
+	ubUnique.AddColumn("user_id", "badge_id")
+	indices = append(indices, ubUnique)
+	return indices
+}
+
 // ErrBadgeAlreadyExist represents a "badge already exists" error.
 type ErrBadgeAlreadyExist struct {
 	Slug string
@@ -189,20 +199,22 @@ func RemoveUserBadge(ctx context.Context, u *User, badge *Badge) error {
 // RemoveUserBadges removes specific badges from a user.
 func RemoveUserBadges(ctx context.Context, u *User, badges []*Badge) error {
 	return db.WithTx(ctx, func(ctx context.Context) error {
-		slugs := make([]string, len(badges))
-		for i, badge := range badges {
-			slugs[i] = badge.Slug
+		badgeSlugs := make([]string, 0, len(badges))
+		for _, badge := range badges {
+			badgeSlugs = append(badgeSlugs, badge.Slug)
 		}
-
-		var badgeIDs []int64
-		if err := db.GetEngine(ctx).Table("badge").In("slug", slugs).Cols("id").Find(&badgeIDs); err != nil {
+		var userBadges []UserBadge
+		if err := db.GetEngine(ctx).Table("user_badge").
+			Join("INNER", "badge", "badge.id = `user_badge`.badge_id").
+			Where("`user_badge`.user_id = ?", u.ID).In("`badge`.slug", badgeSlugs).
+			Find(&userBadges); err != nil {
 			return err
 		}
-
-		if _, err := db.GetEngine(ctx).
-			Where("user_id = ?", u.ID).
-			In("badge_id", badgeIDs).
-			Delete(&UserBadge{}); err != nil {
+		userBadgeIDs := make([]int64, 0, len(userBadges))
+		for _, ub := range userBadges {
+			userBadgeIDs = append(userBadgeIDs, ub.ID)
+		}
+		if _, err := db.GetEngine(ctx).Table("user_badge").In("id", userBadgeIDs).Delete(); err != nil {
 			return err
 		}
 		return nil

models/user/badge_test.go

@@ -33,7 +33,7 @@ func TestGetBadgeUsers(t *testing.T) {
 
 	// Test getting users with pagination
 	opts := &user_model.GetBadgeUsersOptions{
-		Badge: badge,
+		BadgeSlug: badge.Slug,
 		ListOptions: db.ListOptions{
 			Page:     1,
 			PageSize: 1,
@@ -53,9 +53,37 @@ func TestGetBadgeUsers(t *testing.T) {
 	assert.Len(t, users, 1)
 
 	// Test with non-existent badge
-	opts.Badge = &user_model.Badge{Slug: "non-existent"}
+	opts.BadgeSlug = "non-existent"
 	users, count, err = user_model.GetBadgeUsers(db.DefaultContext, opts)
 	assert.NoError(t, err)
 	assert.EqualValues(t, 0, count)
 	assert.Empty(t, users)
 }
+
+func TestAddAndRemoveUserBadges(t *testing.T) {
+	assert.NoError(t, unittest.PrepareTestDatabase())
+	badge1 := unittest.AssertExistsAndLoadBean(t, &user_model.Badge{ID: 1})
+	user1 := unittest.AssertExistsAndLoadBean(t, &user_model.User{ID: 1})
+
+	// Add a badge to user and verify that it is returned in the list
+	assert.NoError(t, user_model.AddUserBadge(db.DefaultContext, user1, badge1))
+	badges, count, err := user_model.GetUserBadges(db.DefaultContext, user1)
+	assert.Equal(t, int64(1), count)
+	assert.Equal(t, badge1.Slug, badges[0].Slug)
+	assert.NoError(t, err)
+
+	// Confirm that it is impossible to duplicate the same badge
+	assert.Error(t, user_model.AddUserBadge(db.DefaultContext, user1, badge1))
+
+	// Nothing happened to the existing badge
+	badges, count, err = user_model.GetUserBadges(db.DefaultContext, user1)
+	assert.Equal(t, int64(1), count)
+	assert.Equal(t, badge1.Slug, badges[0].Slug)
+	assert.NoError(t, err)
+
+	// Remove a badge from user and verify that it is no longer in the list
+	assert.NoError(t, user_model.RemoveUserBadge(db.DefaultContext, user1, badge1))
+	_, count, err = user_model.GetUserBadges(db.DefaultContext, user1)
+	assert.Equal(t, int64(0), count)
+	assert.NoError(t, err)
+}