Annotations: Improve updating annotation tags queries (#71201)

Annotations: Improve annotation tag updates

pkg/infra/db/db.go

@@ -12,13 +12,30 @@ import (
 )
 
 type DB interface {
+	// WithTransactionalDbSession creates a new SQL transaction to ensure consistency
+	// for the database operations done within the [sqlstore.DBTransactionFunc].
+	// It's better to combine InTransaction and WithDbSession instead, as the context
+	// variable is not updated when using this method.
 	WithTransactionalDbSession(ctx context.Context, callback sqlstore.DBTransactionFunc) error
+	// WithDbSession runs database operations either in an existing transaction available
+	// through [context.Context] or if that's not present, as non-transactional database
+	// operations.
 	WithDbSession(ctx context.Context, callback sqlstore.DBTransactionFunc) error
+	// WithNewDbSession behaves like [DB.WithDbSession] without picking up a transaction
+	// from the context.
 	WithNewDbSession(ctx context.Context, callback sqlstore.DBTransactionFunc) error
+	// GetDialect returns an object that contains information about the peculiarities of
+	// the particular database type available to the runtime.
 	GetDialect() migrator.Dialect
+	// GetDBType returns the name of the database type available to the runtime.
 	GetDBType() core.DbType
+	// GetSqlxSession is an experimental extension to use sqlx instead of xorm to
+	// communicate with the database.
 	GetSqlxSession() *session.SessionDB
+	// InTransaction creates a new SQL transaction that is placed on the context.
+	// Use together with [DB.WithDbSession] to run database operations.
 	InTransaction(ctx context.Context, fn func(ctx context.Context) error) error
+	// Quote wraps an identifier so that it cannot be mistaken for an SQL keyword.
 	Quote(value string) string
 	// RecursiveQueriesAreSupported runs a dummy recursive query and it returns true
 	// if the query runs successfully or false if it fails with mysqlerr.ER_PARSE_ERROR error or any other error

pkg/services/annotations/annotationsimpl/xorm_store.go

@@ -66,7 +66,7 @@ func (r *xormRepositoryImpl) Add(ctx context.Context, item *annotations.Item) er
 		if _, err := sess.Table("annotation").Insert(item); err != nil {
 			return err
 		}
-		return r.synchronizeTags(ctx, item)
+		return r.ensureTags(ctx, item.ID, item.Tags)
 	})
 }
 
@@ -115,7 +115,8 @@ func (r *xormRepositoryImpl) AddMany(ctx context.Context, items []annotations.It
 			if _, err := sess.Table("annotation").Insert(item); err != nil {
 				return err
 			}
-			if err := r.synchronizeTags(ctx, &hasTags[i]); err != nil {
+			itemWithID := &hasTags[i]
+			if err := r.ensureTags(ctx, itemWithID.ID, itemWithID.Tags); err != nil {
 				return err
 			}
 		}
@@ -124,24 +125,6 @@ func (r *xormRepositoryImpl) AddMany(ctx context.Context, items []annotations.It
 	})
 }
 
-func (r *xormRepositoryImpl) synchronizeTags(ctx context.Context, item *annotations.Item) error {
-	// Will re-use session if one has already been opened with the same ctx.
-	return r.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
-		if item.Tags != nil {
-			tags, err := r.tagService.EnsureTagsExist(ctx, tag.ParseTagPairs(item.Tags))
-			if err != nil {
-				return err
-			}
-			for _, tag := range tags {
-				if _, err := sess.Exec("INSERT INTO annotation_tag (annotation_id, tag_id) VALUES(?,?)", item.ID, tag.Id); err != nil {
-					return err
-				}
-			}
-		}
-		return nil
-	})
-}
-
 func (r *xormRepositoryImpl) Update(ctx context.Context, item *annotations.Item) error {
 	return r.db.InTransaction(ctx, func(ctx context.Context) error {
 		return r.update(ctx, item)
@@ -180,18 +163,10 @@ func (r *xormRepositoryImpl) update(ctx context.Context, item *annotations.Item)
 		}
 
 		if item.Tags != nil {
-			tags, err := r.tagService.EnsureTagsExist(ctx, tag.ParseTagPairs(item.Tags))
+			err := r.ensureTags(ctx, existing.ID, item.Tags)
 			if err != nil {
 				return err
 			}
-			if _, err := sess.Exec("DELETE FROM annotation_tag WHERE annotation_id = ?", existing.ID); err != nil {
-				return err
-			}
-			for _, tag := range tags {
-				if _, err := sess.Exec("INSERT INTO annotation_tag (annotation_id, tag_id) VALUES(?,?)", existing.ID, tag.Id); err != nil {
-					return err
-				}
-			}
 		}
 
 		existing.Tags = item.Tags
@@ -205,6 +180,63 @@ func (r *xormRepositoryImpl) update(ctx context.Context, item *annotations.Item)
 	})
 }
 
+func (r *xormRepositoryImpl) ensureTags(ctx context.Context, annotationID int64, tags []string) error {
+	return r.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
+		var tagsInsert []annotationTag
+		var tagsDelete []int64
+
+		expectedTags, err := r.tagService.EnsureTagsExist(ctx, tag.ParseTagPairs(tags))
+		if err != nil {
+			return err
+		}
+		expected := tagSet(func(t *tag.Tag) int64 {
+			return t.Id
+		}, expectedTags)
+
+		existingTags := make([]annotationTag, 0)
+		if err := sess.SQL("SELECT annotation_id, tag_id FROM annotation_tag WHERE annotation_id = ?", annotationID).Find(&existingTags); err != nil {
+			return err
+		}
+		existing := tagSet(func(t annotationTag) int64 {
+			return t.TagID
+		}, existingTags)
+
+		for t := range expected {
+			if _, exists := existing[t]; !exists {
+				tagsInsert = append(tagsInsert, annotationTag{
+					AnnotationID: annotationID,
+					TagID:        t,
+				})
+			}
+		}
+		for t := range existing {
+			if _, exists := expected[t]; !exists {
+				tagsDelete = append(tagsDelete, t)
+			}
+		}
+
+		if len(tagsDelete) != 0 {
+			if _, err := sess.MustCols("annotation_id", "tag_id").In("tag_id", tagsDelete).Delete(annotationTag{AnnotationID: annotationID}); err != nil {
+				return err
+			}
+		}
+		if len(tagsInsert) != 0 {
+			if _, err := sess.InsertMulti(tagsInsert); err != nil {
+				return err
+			}
+		}
+		return nil
+	})
+}
+
+func tagSet[T any](fn func(T) int64, list []T) map[int64]struct{} {
+	set := make(map[int64]struct{}, len(list))
+	for _, item := range list {
+		set[fn(item)] = struct{}{}
+	}
+	return set
+}
+
 func (r *xormRepositoryImpl) Get(ctx context.Context, query *annotations.ItemQuery) ([]*annotations.ItemDTO, error) {
 	var sql bytes.Buffer
 	params := make([]interface{}, 0)
@@ -574,3 +606,8 @@ func (r *xormRepositoryImpl) executeUntilDoneOrCancelled(ctx context.Context, sq
 		}
 	}
 }
+
+type annotationTag struct {
+	AnnotationID int64 `xorm:"annotation_id"`
+	TagID        int64 `xorm:"tag_id"`
+}