[release-12.1.5] Chore: Run annotation data migration in batches (#115141)

* Chore: Run annotation data migration in batches (#113589)

* run annotation data migration in batches

* how could i miss it

* run in the background, starting from newest annotations

* update tests

* optionally pass batch size via env

(cherry picked from commit 95e65e2588)

* fix broken cherry pick

* update mysql query

pkg/services/annotations/annotationsimpl/xorm_store.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/grafana/grafana/pkg/services/annotations/accesscontrol"
@@ -37,6 +38,8 @@ func validateTimeRange(item *annotations.Item) error {
 	return nil
 }
 
+var xormMigrationTrigger sync.Once
+
 type xormRepositoryImpl struct {
 	cfg        *setting.Cfg
 	db         db.DB
@@ -45,12 +48,9 @@ type xormRepositoryImpl struct {
 }
 
 func NewXormStore(cfg *setting.Cfg, l log.Logger, db db.DB, tagService tag.Service) *xormRepositoryImpl {
-	// populate dashboard_uid at startup, to ensure safe downgrades & upgrades after
-	// the initial migration occurs
-	err := migrations.RunDashboardUIDMigrations(db.GetEngine().NewSession(), db.GetEngine().DriverName())
-	if err != nil {
-		l.Error("failed to populate dashboard_uid for annotations", "error", err)
-	}
+	xormMigrationTrigger.Do(func() {
+		triggerAlwaysOnMigrations(cfg, l, db)
+	})
 
 	return &xormRepositoryImpl{
 		cfg:        cfg,
@@ -60,6 +60,25 @@ func NewXormStore(cfg *setting.Cfg, l log.Logger, db db.DB, tagService tag.Servi
 	}
 }
 
+func triggerAlwaysOnMigrations(cfg *setting.Cfg, l log.Logger, db db.DB) {
+	sec := cfg.Raw.Section("database")
+	skipDashboardUIDMigration := sec.Key("skip_dashboard_uid_migration_on_startup").MustBool(false)
+	if skipDashboardUIDMigration {
+		l.Debug("skipped dashboard UID startup migration")
+		return
+	}
+	// Run migration in a background goroutine to avoid blocking service startup
+	go func() {
+		l.Info("Starting annotation dashboard_uid migration in background")
+		err := migrations.RunDashboardUIDMigrations(db.GetEngine().NewSession(), db.GetEngine().DriverName(), l)
+		if err != nil {
+			l.Error("failed to populate dashboard_uid for annotations", "error", err)
+		} else {
+			l.Info("Annotation dashboard_uid migration completed successfully")
+		}
+	}()
+}
+
 func (r *xormRepositoryImpl) Type() string {
 	return "sql"
 }

pkg/services/annotations/annotationsimpl/xorm_store_test.go

@@ -4,7 +4,9 @@ import (
 	"context"
 	"fmt"
 	"strings"
+	"sync"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -26,9 +28,6 @@ import (
 )
 
 func TestIntegrationAnnotations(t *testing.T) {
-	if testing.Short() {
-		t.Skip("skipping integration test")
-	}
 	sql := db.InitTestDB(t)
 
 	cfg := setting.NewCfg()
@@ -648,6 +647,146 @@ func TestIntegrationAnnotations(t *testing.T) {
 	})
 }
 
+func TestIntegrationAnnotationsAlwaysOnMigrations(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping integration test")
+	}
+
+	sql := db.InitTestDB(t)
+	cfg := setting.NewCfg()
+	cfg.AnnotationMaximumTagsLength = 60
+
+	t.Run("NewXormStore should call triggerAlwaysOnMigrations and skip migrations", func(t *testing.T) {
+		cfg.Raw.Section("database").Key("skip_dashboard_uid_migration_on_startup").SetValue("true")
+
+		l := log.New("annotation.test")
+
+		dashboard := testutil.CreateDashboard(t, sql, cfg, featuremgmt.WithFeatures(), dashboards.SaveDashboardCommand{
+			UserID: 1,
+			OrgID:  1,
+			Dashboard: simplejson.NewFromAny(map[string]any{
+				"title": "Test Skip Dashboard",
+				"uid":   "test-skip-uid",
+			}),
+		})
+
+		tempStore := NewXormStore(cfg, l, sql, tagimpl.ProvideService(sql))
+		annotation := &annotations.Item{
+			OrgID:        1,
+			UserID:       1,
+			DashboardID:  dashboard.ID, // nolint: staticcheck
+			DashboardUID: dashboard.UID,
+			Text:         "test migration skip",
+			Type:         "alert",
+			Epoch:        100,
+		}
+		err := tempStore.Add(context.Background(), annotation)
+		require.NoError(t, err)
+
+		t.Cleanup(func() {
+			err := tempStore.Delete(context.Background(), &annotations.DeleteParams{ID: annotation.ID, OrgID: 1})
+			assert.NoError(t, err)
+		})
+
+		err = sql.WithDbSession(context.Background(), func(sess *db.Session) error {
+			_, err := sess.Exec("UPDATE annotation SET dashboard_uid = NULL WHERE id = ?", annotation.ID)
+			return err
+		})
+		require.NoError(t, err)
+
+		xormMigrationTrigger = sync.Once{}
+		store := NewXormStore(cfg, l, sql, tagimpl.ProvideService(sql))
+
+		require.NotNil(t, store)
+		assert.Equal(t, "sql", store.Type())
+
+		var result struct {
+			DashboardUID *string `xorm:"dashboard_uid"`
+		}
+		err = sql.WithDbSession(context.Background(), func(sess *db.Session) error {
+			has, err := sess.Table("annotation").
+				Where("id = ?", annotation.ID).
+				Get(&result)
+			if err != nil {
+				return err
+			}
+			if !has {
+				return fmt.Errorf("annotation not found")
+			}
+			return nil
+		})
+		require.NoError(t, err)
+		assert.Nil(t, result.DashboardUID, "dashboard_uid should still be NULL when migration is skipped")
+	})
+
+	t.Run("NewXormStore should call triggerAlwaysOnMigrations and run migrations", func(t *testing.T) {
+		cfg.Raw.Section("database").Key("skip_dashboard_uid_migration_on_startup").SetValue("false")
+		l := log.New("annotation.test")
+
+		dashboard := testutil.CreateDashboard(t, sql, cfg, featuremgmt.WithFeatures(), dashboards.SaveDashboardCommand{
+			UserID: 1,
+			OrgID:  1,
+			Dashboard: simplejson.NewFromAny(map[string]any{
+				"title": "Test Run Dashboard",
+				"uid":   "test-run-uid",
+			}),
+		})
+
+		tempStore := NewXormStore(cfg, l, sql, tagimpl.ProvideService(sql))
+		annotation := &annotations.Item{
+			OrgID:        1,
+			UserID:       1,
+			DashboardID:  dashboard.ID, // nolint: staticcheck
+			DashboardUID: dashboard.UID,
+			Text:         "test migration run",
+			Type:         "alert",
+			Epoch:        100,
+		}
+		err := tempStore.Add(context.Background(), annotation)
+		require.NoError(t, err)
+
+		t.Cleanup(func() {
+			err := tempStore.Delete(context.Background(), &annotations.DeleteParams{ID: annotation.ID, OrgID: 1})
+			assert.NoError(t, err)
+		})
+
+		err = sql.WithDbSession(context.Background(), func(sess *db.Session) error {
+			_, err := sess.Exec("UPDATE annotation SET dashboard_uid = NULL WHERE id = ?", annotation.ID)
+			return err
+		})
+		require.NoError(t, err)
+
+		xormMigrationTrigger = sync.Once{}
+		store := NewXormStore(cfg, l, sql, tagimpl.ProvideService(sql))
+
+		require.NotNil(t, store)
+		assert.Equal(t, "sql", store.Type())
+
+		// Wait for the async migration to complete (runs in background goroutine)
+		var result struct {
+			DashboardUID *string `xorm:"dashboard_uid"`
+		}
+		require.Eventually(t, func() bool {
+			err = sql.WithDbSession(context.Background(), func(sess *db.Session) error {
+				has, err := sess.Table("annotation").
+					Where("id = ?", annotation.ID).
+					Get(&result)
+				if err != nil {
+					return err
+				}
+				if !has {
+					return fmt.Errorf("annotation not found")
+				}
+				return nil
+			})
+			if err != nil {
+				return false
+			}
+			return result.DashboardUID != nil && *result.DashboardUID == "test-run-uid"
+		}, 5*time.Second, 100*time.Millisecond, "dashboard_uid should be populated when migration runs")
+	})
+}
+
 func BenchmarkFindTags_10k(b *testing.B) {
 	benchmarkFindTags(b, 10000)
 }