Chore: Move annotations cleanup to the annotations service (#55618)

pkg/services/annotations/annotations.go

@@ -19,18 +19,7 @@ type Repository interface {
 	FindTags(ctx context.Context, query *TagsQuery) (FindTagsResult, error)
 }
 
-// AnnotationCleaner is responsible for cleaning up old annotations
-type AnnotationCleaner interface {
-	CleanAnnotations(ctx context.Context, cfg *setting.Cfg) (int64, int64, error)
-}
-
-// var repositoryInstance Repository
-var cleanerInstance AnnotationCleaner
-
-func GetAnnotationCleaner() AnnotationCleaner {
-	return cleanerInstance
-}
-
-func SetAnnotationCleaner(rep AnnotationCleaner) {
-	cleanerInstance = rep
+// Cleaner is responsible for cleaning up old annotations
+type Cleaner interface {
+	Run(ctx context.Context, cfg *setting.Cfg) (int64, int64, error)
 }

pkg/services/annotations/annotationsimpl/annotations.go

@@ -16,7 +16,7 @@ type RepositoryImpl struct {
 
 func ProvideService(db db.DB, cfg *setting.Cfg, tagService tag.Service) *RepositoryImpl {
 	return &RepositoryImpl{
-		store: &SQLAnnotationRepo{
+		store: &xormRepositoryImpl{
 			cfg:        cfg,
 			db:         db,
 			log:        log.New("annotations"),

pkg/services/annotations/annotationsimpl/cleanup.go

@@ -0,0 +1,62 @@
+package annotationsimpl
+
+import (
+	"context"
+
+	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/services/sqlstore/db"
+	"github.com/grafana/grafana/pkg/setting"
+)
+
+// CleanupServiceImpl is responsible for cleaning old annotations.
+type CleanupServiceImpl struct {
+	store store
+}
+
+func ProvideCleanupService(db db.DB, cfg *setting.Cfg) *CleanupServiceImpl {
+	return &CleanupServiceImpl{
+		store: &xormRepositoryImpl{
+			cfg: cfg,
+			db:  db,
+			log: log.New("annotations"),
+		},
+	}
+}
+
+const (
+	alertAnnotationType     = "alert_id <> 0"
+	dashboardAnnotationType = "dashboard_id <> 0 AND alert_id = 0"
+	apiAnnotationType       = "alert_id = 0 AND dashboard_id = 0"
+)
+
+// Run deletes old annotations created by alert rules, API
+// requests and human made in the UI. It subsequently deletes orphaned rows
+// from the annotation_tag table. Cleanup actions are performed in batches
+// so that no query takes too long to complete.
+//
+// Returns the number of annotation and annotation_tag rows deleted. If an
+// error occurs, it returns the number of rows affected so far.
+func (cs *CleanupServiceImpl) Run(ctx context.Context, cfg *setting.Cfg) (int64, int64, error) {
+	var totalCleanedAnnotations int64
+	affected, err := cs.store.CleanAnnotations(ctx, cfg.AlertingAnnotationCleanupSetting, alertAnnotationType)
+	totalCleanedAnnotations += affected
+	if err != nil {
+		return totalCleanedAnnotations, 0, err
+	}
+
+	affected, err = cs.store.CleanAnnotations(ctx, cfg.APIAnnotationCleanupSettings, apiAnnotationType)
+	totalCleanedAnnotations += affected
+	if err != nil {
+		return totalCleanedAnnotations, 0, err
+	}
+
+	affected, err = cs.store.CleanAnnotations(ctx, cfg.DashboardAnnotationCleanupSettings, dashboardAnnotationType)
+	totalCleanedAnnotations += affected
+	if err != nil {
+		return totalCleanedAnnotations, 0, err
+	}
+	if totalCleanedAnnotations > 0 {
+		affected, err = cs.store.CleanOrphanedAnnotationTags(ctx)
+	}
+	return totalCleanedAnnotations, affected, err
+}

pkg/services/annotations/annotationsimpl/cleanup_test.go

@@ -0,0 +1,234 @@
+package annotationsimpl
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/services/annotations"
+	"github.com/grafana/grafana/pkg/services/sqlstore"
+	"github.com/grafana/grafana/pkg/setting"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestAnnotationCleanUp(t *testing.T) {
+	fakeSQL := sqlstore.InitTestDB(t)
+
+	t.Cleanup(func() {
+		err := fakeSQL.WithDbSession(context.Background(), func(session *sqlstore.DBSession) error {
+			_, err := session.Exec("DELETE FROM annotation")
+			return err
+		})
+		assert.NoError(t, err)
+	})
+
+	createTestAnnotations(t, fakeSQL, 21, 6)
+	assertAnnotationCount(t, fakeSQL, "", 21)
+	assertAnnotationTagCount(t, fakeSQL, 42)
+
+	tests := []struct {
+		name                     string
+		cfg                      *setting.Cfg
+		alertAnnotationCount     int64
+		dashboardAnnotationCount int64
+		APIAnnotationCount       int64
+		affectedAnnotations      int64
+	}{
+		{
+			name: "default settings should not delete any annotations",
+			cfg: &setting.Cfg{
+				AlertingAnnotationCleanupSetting:   settingsFn(0, 0),
+				DashboardAnnotationCleanupSettings: settingsFn(0, 0),
+				APIAnnotationCleanupSettings:       settingsFn(0, 0),
+			},
+			alertAnnotationCount:     7,
+			dashboardAnnotationCount: 7,
+			APIAnnotationCount:       7,
+			affectedAnnotations:      0,
+		},
+		{
+			name: "should remove annotations created before cut off point",
+			cfg: &setting.Cfg{
+				AlertingAnnotationCleanupSetting:   settingsFn(time.Hour*48, 0),
+				DashboardAnnotationCleanupSettings: settingsFn(time.Hour*48, 0),
+				APIAnnotationCleanupSettings:       settingsFn(time.Hour*48, 0),
+			},
+			alertAnnotationCount:     5,
+			dashboardAnnotationCount: 5,
+			APIAnnotationCount:       5,
+			affectedAnnotations:      6,
+		},
+		{
+			name: "should only keep three annotations",
+			cfg: &setting.Cfg{
+				AlertingAnnotationCleanupSetting:   settingsFn(0, 3),
+				DashboardAnnotationCleanupSettings: settingsFn(0, 3),
+				APIAnnotationCleanupSettings:       settingsFn(0, 3),
+			},
+			alertAnnotationCount:     3,
+			dashboardAnnotationCount: 3,
+			APIAnnotationCount:       3,
+			affectedAnnotations:      6,
+		},
+		{
+			name: "running the max count delete again should not remove any annotations",
+			cfg: &setting.Cfg{
+				AlertingAnnotationCleanupSetting:   settingsFn(0, 3),
+				DashboardAnnotationCleanupSettings: settingsFn(0, 3),
+				APIAnnotationCleanupSettings:       settingsFn(0, 3),
+			},
+			alertAnnotationCount:     3,
+			dashboardAnnotationCount: 3,
+			APIAnnotationCount:       3,
+			affectedAnnotations:      0,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			cfg := setting.NewCfg()
+			cfg.AnnotationCleanupJobBatchSize = 1
+			cleaner := ProvideCleanupService(fakeSQL, cfg)
+			affectedAnnotations, affectedAnnotationTags, err := cleaner.Run(context.Background(), test.cfg)
+			require.NoError(t, err)
+
+			assert.Equal(t, test.affectedAnnotations, affectedAnnotations)
+			assert.Equal(t, test.affectedAnnotations*2, affectedAnnotationTags)
+
+			assertAnnotationCount(t, fakeSQL, alertAnnotationType, test.alertAnnotationCount)
+			assertAnnotationCount(t, fakeSQL, dashboardAnnotationType, test.dashboardAnnotationCount)
+			assertAnnotationCount(t, fakeSQL, apiAnnotationType, test.APIAnnotationCount)
+
+			// we create two records in annotation_tag for each sample annotation
+			expectedAnnotationTagCount := (test.alertAnnotationCount +
+				test.dashboardAnnotationCount +
+				test.APIAnnotationCount) * 2
+			assertAnnotationTagCount(t, fakeSQL, expectedAnnotationTagCount)
+		})
+	}
+}
+
+func TestOldAnnotationsAreDeletedFirst(t *testing.T) {
+	fakeSQL := sqlstore.InitTestDB(t)
+
+	t.Cleanup(func() {
+		err := fakeSQL.WithDbSession(context.Background(), func(session *sqlstore.DBSession) error {
+			_, err := session.Exec("DELETE FROM annotation")
+			return err
+		})
+		assert.NoError(t, err)
+	})
+
+	// create some test annotations
+	a := annotations.Item{
+		DashboardId: 1,
+		OrgId:       1,
+		UserId:      1,
+		PanelId:     1,
+		AlertId:     10,
+		Text:        "",
+		Created:     time.Now().AddDate(-10, 0, -10).UnixNano() / int64(time.Millisecond),
+	}
+
+	session := fakeSQL.NewSession(context.Background())
+	defer session.Close()
+
+	_, err := session.Insert(a)
+	require.NoError(t, err, "cannot insert annotation")
+	_, err = session.Insert(a)
+	require.NoError(t, err, "cannot insert annotation")
+
+	a.AlertId = 20
+	_, err = session.Insert(a)
+	require.NoError(t, err, "cannot insert annotation")
+
+	// run the clean up task to keep one annotation.
+	cfg := setting.NewCfg()
+	cfg.AnnotationCleanupJobBatchSize = 1
+	cleaner := &xormRepositoryImpl{cfg: cfg, log: log.New("test-logger"), db: fakeSQL}
+	_, err = cleaner.CleanAnnotations(context.Background(), setting.AnnotationCleanupSettings{MaxCount: 1}, alertAnnotationType)
+	require.NoError(t, err)
+
+	// assert that the last annotations were kept
+	countNew, err := session.Where("alert_id = 20").Count(&annotations.Item{})
+	require.NoError(t, err)
+	require.Equal(t, int64(1), countNew, "the last annotations should be kept")
+
+	countOld, err := session.Where("alert_id = 10").Count(&annotations.Item{})
+	require.NoError(t, err)
+	require.Equal(t, int64(0), countOld, "the two first annotations should have been deleted")
+}
+
+func assertAnnotationCount(t *testing.T, fakeSQL *sqlstore.SQLStore, sql string, expectedCount int64) {
+	t.Helper()
+
+	session := fakeSQL.NewSession(context.Background())
+	defer session.Close()
+	count, err := session.Where(sql).Count(&annotations.Item{})
+	require.NoError(t, err)
+	require.Equal(t, expectedCount, count)
+}
+
+func assertAnnotationTagCount(t *testing.T, fakeSQL *sqlstore.SQLStore, expectedCount int64) {
+	t.Helper()
+
+	session := fakeSQL.NewSession(context.Background())
+	defer session.Close()
+
+	count, err := session.SQL("select count(*) from annotation_tag").Count()
+	require.NoError(t, err)
+	require.Equal(t, expectedCount, count)
+}
+
+func createTestAnnotations(t *testing.T, sqlstore *sqlstore.SQLStore, expectedCount int, oldAnnotations int) {
+	t.Helper()
+
+	cutoffDate := time.Now()
+
+	for i := 0; i < expectedCount; i++ {
+		a := &annotations.Item{
+			DashboardId: 1,
+			OrgId:       1,
+			UserId:      1,
+			PanelId:     1,
+			Text:        "",
+		}
+
+		// mark every third as an API annotation
+		// that does not belong to a dashboard
+		if i%3 == 1 {
+			a.DashboardId = 0
+		}
+
+		// mark every third annotation as an alert annotation
+		if i%3 == 0 {
+			a.AlertId = 10
+			a.DashboardId = 2
+		}
+
+		// create epoch as int annotations.go line 40
+		a.Created = cutoffDate.UnixNano() / int64(time.Millisecond)
+
+		// set a really old date for the first six annotations
+		if i < oldAnnotations {
+			a.Created = cutoffDate.AddDate(-10, 0, -10).UnixNano() / int64(time.Millisecond)
+		}
+
+		_, err := sqlstore.NewSession(context.Background()).Insert(a)
+		require.NoError(t, err, "should be able to save annotation", err)
+
+		// mimick the SQL annotation Save logic by writing records to the annotation_tag table
+		// we need to ensure they get deleted when we clean up annotations
+		sess := sqlstore.NewSession(context.Background())
+		for tagID := range []int{1, 2} {
+			_, err = sess.Exec("INSERT INTO annotation_tag (annotation_id, tag_id) VALUES(?,?)", a.Id, tagID)
+			require.NoError(t, err, "should be able to save annotation tag ID", err)
+		}
+	}
+}
+
+func settingsFn(maxAge time.Duration, maxCount int64) setting.AnnotationCleanupSettings {
+	return setting.AnnotationCleanupSettings{MaxAge: maxAge, MaxCount: maxCount}
+}

pkg/services/annotations/annotationsimpl/store.go

@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"github.com/grafana/grafana/pkg/services/annotations"
+	"github.com/grafana/grafana/pkg/setting"
 )
 
 type store interface {
@@ -12,4 +13,6 @@ type store interface {
 	Get(ctx context.Context, query *annotations.ItemQuery) ([]*annotations.ItemDTO, error)
 	Delete(ctx context.Context, params *annotations.DeleteParams) error
 	GetTags(ctx context.Context, query *annotations.TagsQuery) (annotations.FindTagsResult, error)
+	CleanAnnotations(ctx context.Context, cfg setting.AnnotationCleanupSettings, annotationType string) (int64, error)
+	CleanOrphanedAnnotationTags(ctx context.Context) (int64, error)
 }

pkg/services/annotations/annotationsimpl/xorm_store.go

@@ -40,14 +40,14 @@ func validateTimeRange(item *annotations.Item) error {
 	return nil
 }
 
-type SQLAnnotationRepo struct {
+type xormRepositoryImpl struct {
 	cfg        *setting.Cfg
 	db         db.DB
 	log        log.Logger
 	tagService tag.Service
 }
 
-func (r *SQLAnnotationRepo) Add(ctx context.Context, item *annotations.Item) error {
+func (r *xormRepositoryImpl) Add(ctx context.Context, item *annotations.Item) error {
 	tags := tag.ParseTagPairs(item.Tags)
 	item.Tags = tag.JoinTagPairs(tags)
 	item.Created = timeNow().UnixNano() / int64(time.Millisecond)
@@ -79,7 +79,7 @@ func (r *SQLAnnotationRepo) Add(ctx context.Context, item *annotations.Item) err
 	})
 }
 
-func (r *SQLAnnotationRepo) Update(ctx context.Context, item *annotations.Item) error {
+func (r *xormRepositoryImpl) Update(ctx context.Context, item *annotations.Item) error {
 	return r.db.WithTransactionalDbSession(ctx, func(sess *sqlstore.DBSession) error {
 		var (
 			isExist bool
@@ -132,7 +132,7 @@ func (r *SQLAnnotationRepo) Update(ctx context.Context, item *annotations.Item)
 	})
 }
 
-func (r *SQLAnnotationRepo) Get(ctx context.Context, query *annotations.ItemQuery) ([]*annotations.ItemDTO, error) {
+func (r *xormRepositoryImpl) Get(ctx context.Context, query *annotations.ItemQuery) ([]*annotations.ItemDTO, error) {
 	var sql bytes.Buffer
 	params := make([]interface{}, 0)
 	items := make([]*annotations.ItemDTO, 0)
@@ -291,7 +291,7 @@ func getAccessControlFilter(user *user.SignedInUser) (string, []interface{}, err
 	return strings.Join(filters, " OR "), params, nil
 }
 
-func (r *SQLAnnotationRepo) Delete(ctx context.Context, params *annotations.DeleteParams) error {
+func (r *xormRepositoryImpl) Delete(ctx context.Context, params *annotations.DeleteParams) error {
 	return r.db.WithTransactionalDbSession(ctx, func(sess *sqlstore.DBSession) error {
 		var (
 			sql        string
@@ -327,7 +327,7 @@ func (r *SQLAnnotationRepo) Delete(ctx context.Context, params *annotations.Dele
 	})
 }
 
-func (r *SQLAnnotationRepo) GetTags(ctx context.Context, query *annotations.TagsQuery) (annotations.FindTagsResult, error) {
+func (r *xormRepositoryImpl) GetTags(ctx context.Context, query *annotations.TagsQuery) (annotations.FindTagsResult, error) {
 	var items []*annotations.Tag
 	err := r.db.WithDbSession(ctx, func(dbSession *sqlstore.DBSession) error {
 		if query.Limit == 0 {
@@ -378,3 +378,64 @@ func (r *SQLAnnotationRepo) GetTags(ctx context.Context, query *annotations.Tags
 
 	return annotations.FindTagsResult{Tags: tags}, nil
 }
+
+func (r *xormRepositoryImpl) CleanAnnotations(ctx context.Context, cfg setting.AnnotationCleanupSettings, annotationType string) (int64, error) {
+	var totalAffected int64
+	if cfg.MaxAge > 0 {
+		cutoffDate := time.Now().Add(-cfg.MaxAge).UnixNano() / int64(time.Millisecond)
+		deleteQuery := `DELETE FROM annotation WHERE id IN (SELECT id FROM (SELECT id FROM annotation WHERE %s AND created < %v ORDER BY id DESC %s) a)`
+		sql := fmt.Sprintf(deleteQuery, annotationType, cutoffDate, r.db.GetDialect().Limit(r.cfg.AnnotationCleanupJobBatchSize))
+
+		affected, err := r.executeUntilDoneOrCancelled(ctx, sql)
+		totalAffected += affected
+		if err != nil {
+			return totalAffected, err
+		}
+	}
+
+	if cfg.MaxCount > 0 {
+		deleteQuery := `DELETE FROM annotation WHERE id IN (SELECT id FROM (SELECT id FROM annotation WHERE %s ORDER BY id DESC %s) a)`
+		sql := fmt.Sprintf(deleteQuery, annotationType, r.db.GetDialect().LimitOffset(r.cfg.AnnotationCleanupJobBatchSize, cfg.MaxCount))
+		affected, err := r.executeUntilDoneOrCancelled(ctx, sql)
+		totalAffected += affected
+		return totalAffected, err
+	}
+
+	return totalAffected, nil
+}
+
+func (r *xormRepositoryImpl) CleanOrphanedAnnotationTags(ctx context.Context) (int64, error) {
+	deleteQuery := `DELETE FROM annotation_tag WHERE id IN ( SELECT id FROM (SELECT id FROM annotation_tag WHERE NOT EXISTS (SELECT 1 FROM annotation a WHERE annotation_id = a.id) %s) a)`
+	sql := fmt.Sprintf(deleteQuery, r.db.GetDialect().Limit(r.cfg.AnnotationCleanupJobBatchSize))
+	return r.executeUntilDoneOrCancelled(ctx, sql)
+}
+
+func (r *xormRepositoryImpl) executeUntilDoneOrCancelled(ctx context.Context, sql string) (int64, error) {
+	var totalAffected int64
+	for {
+		select {
+		case <-ctx.Done():
+			return totalAffected, ctx.Err()
+		default:
+			var affected int64
+			err := r.db.WithDbSession(ctx, func(session *sqlstore.DBSession) error {
+				res, err := session.Exec(sql)
+				if err != nil {
+					return err
+				}
+
+				affected, err = res.RowsAffected()
+				totalAffected += affected
+
+				return err
+			})
+			if err != nil {
+				return totalAffected, err
+			}
+
+			if affected == 0 {
+				return totalAffected, nil
+			}
+		}
+	}
+}

pkg/services/annotations/annotationsimpl/xorm_store_test.go

@@ -28,7 +28,7 @@ func TestIntegrationAnnotations(t *testing.T) {
 		t.Skip("skipping integration test")
 	}
 	sql := sqlstore.InitTestDB(t)
-	repo := SQLAnnotationRepo{db: sql, cfg: setting.NewCfg(), log: log.New("annotation.test"), tagService: tagimpl.ProvideService(sql)}
+	repo := xormRepositoryImpl{db: sql, cfg: setting.NewCfg(), log: log.New("annotation.test"), tagService: tagimpl.ProvideService(sql)}
 
 	testUser := &user.SignedInUser{
 		OrgID: 1,
@@ -394,7 +394,7 @@ func TestIntegrationAnnotationListingWithRBAC(t *testing.T) {
 		t.Skip("skipping integration test")
 	}
 	sql := sqlstore.InitTestDB(t, sqlstore.InitTestDBOpt{})
-	repo := SQLAnnotationRepo{db: sql, cfg: setting.NewCfg(), log: log.New("annotation.test"), tagService: tagimpl.ProvideService(sql)}
+	repo := xormRepositoryImpl{db: sql, cfg: setting.NewCfg(), log: log.New("annotation.test"), tagService: tagimpl.ProvideService(sql)}
 	dashboardStore := dashboardstore.ProvideDashboardStore(sql, featuremgmt.WithFeatures(), tagimpl.ProvideService(sql))
 
 	testDashboard1 := models.SaveDashboardCommand{

pkg/services/annotations/annotationstest/fake_cleanup.go

@@ -0,0 +1,18 @@
+package annotationstest
+
+import (
+	"context"
+
+	"github.com/grafana/grafana/pkg/setting"
+)
+
+type fakeCleaner struct {
+}
+
+func NewFakeCleaner() *fakeCleaner {
+	return &fakeCleaner{}
+}
+
+func (f *fakeCleaner) Run(ctx context.Context, cfg *setting.Cfg) (int64, int64, error) {
+	return 0, 0, nil
+}