From 04370f92dc7acfad01363a4a3ecd31d64ca72abb Mon Sep 17 00:00:00 2001
From: Sofia Papagiannaki <1632407+papagian@users.noreply.github.com>
Date: Tue, 18 Apr 2023 13:27:40 +0300
Subject: [PATCH] Annotations: Improve get tags query performance (#66182)

* Add benchmark

* Add fix
---
 .../annotations/annotationsimpl/xorm_store.go |  3 +-
 .../annotationsimpl/xorm_store_test.go        | 89 +++++++++++++++++++
 2 files changed, 91 insertions(+), 1 deletion(-)

diff --git a/pkg/services/annotations/annotationsimpl/xorm_store.go b/pkg/services/annotations/annotationsimpl/xorm_store.go
index a78584f07a2..06e8e79b193 100644
--- a/pkg/services/annotations/annotationsimpl/xorm_store.go
+++ b/pkg/services/annotations/annotationsimpl/xorm_store.go
@@ -447,9 +447,10 @@ func (r *xormRepositoryImpl) GetTags(ctx context.Context, query *annotations.Tag
 			count(*) as count
 		FROM tag
 		INNER JOIN annotation_tag ON tag.id = annotation_tag.tag_id
+		INNER JOIN annotation ON annotation.id = annotation_tag.annotation_id
 `)
 
-		sql.WriteString(`WHERE EXISTS(SELECT 1 FROM annotation WHERE annotation.id = annotation_tag.annotation_id AND annotation.org_id = ?)`)
+		sql.WriteString(`WHERE annotation.org_id = ?`)
 		params = append(params, query.OrgID)
 
 		sql.WriteString(` AND (` + tagKey + ` ` + r.db.GetDialect().LikeStr() + ` ? OR ` + tagValue + ` ` + r.db.GetDialect().LikeStr() + ` ?)`)
diff --git a/pkg/services/annotations/annotationsimpl/xorm_store_test.go b/pkg/services/annotations/annotationsimpl/xorm_store_test.go
index 630e266911e..4bbd98b907c 100644
--- a/pkg/services/annotations/annotationsimpl/xorm_store_test.go
+++ b/pkg/services/annotations/annotationsimpl/xorm_store_test.go
@@ -27,6 +27,7 @@ import (
 	"github.com/grafana/grafana/pkg/services/guardian"
 	"github.com/grafana/grafana/pkg/services/quota/quotatest"
 	"github.com/grafana/grafana/pkg/services/sqlstore"
+	"github.com/grafana/grafana/pkg/services/tag"
 	"github.com/grafana/grafana/pkg/services/tag/tagimpl"
 	"github.com/grafana/grafana/pkg/services/user"
 	"github.com/grafana/grafana/pkg/setting"
@@ -847,3 +848,91 @@ func setupRBACPermission(t *testing.T, db *sqlstore.SQLStore, role *accesscontro
 
 	require.NoError(t, err)
 }
+
+func BenchmarkFindTags_10k(b *testing.B) {
+	benchmarkFindTags(b, 10000)
+}
+
+func BenchmarkFindTags_100k(b *testing.B) {
+	benchmarkFindTags(b, 100000)
+}
+
+func benchmarkFindTags(b *testing.B, numAnnotations int) {
+	sql := db.InitTestDB(b)
+	var maximumTagsLength int64 = 60
+	repo := xormRepositoryImpl{db: sql, cfg: setting.NewCfg(), log: log.New("annotation.test"), tagService: tagimpl.ProvideService(sql, sql.Cfg), maximumTagsLength: maximumTagsLength}
+
+	type annotationTag struct {
+		ID           int64 `xorm:"pk autoincr 'id'"`
+		AnnotationID int64 `xorm:"annotation_id"`
+		TagID        int64 `xorm:"tag_id"`
+	}
+	newAnnotations := make([]annotations.Item, 0, numAnnotations)
+	newTags := make([]tag.Tag, 0, numAnnotations)
+	newAnnotationTags := make([]annotationTag, 0, numAnnotations)
+	for i := 0; i < numAnnotations; i++ {
+		newAnnotations = append(newAnnotations, annotations.Item{
+			ID:          int64(i),
+			OrgID:       1,
+			UserID:      1,
+			DashboardID: int64(i),
+			Text:        "hello",
+			Type:        "alert",
+			Epoch:       10,
+			Data:        simplejson.NewFromAny(map[string]interface{}{"data1": "I am a cool data", "data2": "I am another cool data"}),
+		})
+		newTags = append(newTags, tag.Tag{
+			Id:  int64(i),
+			Key: fmt.Sprintf("tag%d", i),
+		})
+
+		newAnnotationTags = append(newAnnotationTags, annotationTag{
+			AnnotationID: int64(i),
+			TagID:        int64(1),
+		})
+	}
+	err := sql.WithDbSession(context.Background(), func(sess *sqlstore.DBSession) error {
+		batchSize := 1000
+		numOfBatches := numAnnotations / batchSize
+		for i := 0; i < numOfBatches; i++ {
+			_, err := sess.Insert(newAnnotations[i*batchSize : (i+1)*batchSize-1])
+			require.NoError(b, err)
+
+			_, err = sess.Insert(newTags[i*batchSize : (i+1)*batchSize-1])
+			require.NoError(b, err)
+
+			_, err = sess.Insert(newAnnotationTags[i*batchSize : (i+1)*batchSize-1])
+			require.NoError(b, err)
+		}
+		return nil
+	})
+	require.NoError(b, err)
+
+	annotationWithTheTag := annotations.Item{
+		ID:          int64(numAnnotations) + 1,
+		OrgID:       1,
+		UserID:      1,
+		DashboardID: int64(1),
+		Text:        "hello",
+		Type:        "alert",
+		Epoch:       10,
+		Tags:        []string{"outage", "error", "type:outage", "server:server-1"},
+		Data:        simplejson.NewFromAny(map[string]interface{}{"data1": "I am a cool data", "data2": "I am another cool data"}),
+	}
+	err = repo.Add(context.Background(), &annotationWithTheTag)
+	require.NoError(b, err)
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		result, err := repo.GetTags(context.Background(), &annotations.TagsQuery{
+			OrgID: 1,
+			Tag:   "outage",
+		})
+		require.NoError(b, err)
+		require.Len(b, result.Tags, 2)
+		require.Equal(b, "outage", result.Tags[0].Tag)
+		require.Equal(b, "type:outage", result.Tags[1].Tag)
+		require.Equal(b, int64(1), result.Tags[0].Count)
+		require.Equal(b, int64(1), result.Tags[1].Count)
+	}
+}