[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)
 }

pkg/services/sqlstore/migrations/annotation_mig.go

@@ -2,7 +2,10 @@ package migrations
 
 import (
 	"fmt"
+	"os"
+	"strconv"
 
+	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/util/xorm"
 
 	. "github.com/grafana/grafana/pkg/services/sqlstore/migrator"
@@ -243,30 +246,86 @@ func (m *SetDashboardUIDMigration) SQL(dialect Dialect) string {
 }
 
 func (m *SetDashboardUIDMigration) Exec(sess *xorm.Session, mg *Migrator) error {
-	return RunDashboardUIDMigrations(sess, mg.Dialect.DriverName())
+	return RunDashboardUIDMigrations(sess, mg.Dialect.DriverName(), mg.Logger)
 }
 
-func RunDashboardUIDMigrations(sess *xorm.Session, driverName string) error {
-	sql := `UPDATE annotation
-	SET dashboard_uid = (SELECT uid FROM dashboard WHERE dashboard.id = annotation.dashboard_id)
-	WHERE dashboard_uid IS NULL AND dashboard_id != 0 AND EXISTS (SELECT 1 FROM dashboard WHERE dashboard.id = annotation.dashboard_id);`
+func RunDashboardUIDMigrations(sess *xorm.Session, driverName string, logger log.Logger) error {
+	batchSize := 5000
+	if size := os.Getenv("ANNOTATION_DASHBOARD_UID_MIGRATION_BATCH_SIZE"); size != "" {
+		n, err := strconv.ParseInt(size, 10, 64)
+		if err == nil {
+			batchSize = int(n)
+		}
+	}
+
+	logger.Info("Starting batched dashboard_uid migration for annotations (newest first)", "batchSize", batchSize)
+	updateSQL := `UPDATE annotation
+		SET dashboard_uid = (SELECT uid FROM dashboard WHERE dashboard.id = annotation.dashboard_id)
+		WHERE dashboard_uid IS NULL
+		  AND dashboard_id != 0
+		  AND EXISTS (SELECT 1 FROM dashboard WHERE dashboard.id = annotation.dashboard_id)
+		  AND annotation.id IN (
+			SELECT id FROM annotation
+			WHERE dashboard_uid IS NULL AND dashboard_id != 0
+			ORDER BY id DESC
+			LIMIT ?
+		  )`
 	switch driverName {
 	case Postgres:
-		sql = `UPDATE annotation
+		updateSQL = `UPDATE annotation
 		SET dashboard_uid = dashboard.uid
 		FROM dashboard
 		WHERE annotation.dashboard_id = dashboard.id
 		 AND annotation.dashboard_id != 0
-		 AND annotation.dashboard_uid IS NULL;`
+		 AND annotation.dashboard_uid IS NULL
+		 AND annotation.id IN (
+			SELECT id FROM annotation
+			WHERE dashboard_uid IS NULL AND dashboard_id != 0
+			ORDER BY id DESC
+			LIMIT $1
+		 )`
 	case MySQL:
-		sql = `UPDATE annotation
-		LEFT JOIN dashboard ON annotation.dashboard_id = dashboard.id
-		SET annotation.dashboard_uid = dashboard.uid
-		WHERE annotation.dashboard_uid IS NULL and annotation.dashboard_id != 0;`
-	}
-	if _, err := sess.Exec(sql); err != nil {
-		return fmt.Errorf("failed to set dashboard_uid for annotation: %w", err)
+		updateSQL = `UPDATE annotation AS a
+		JOIN dashboard AS d ON a.dashboard_id = d.id
+		JOIN (
+		  SELECT id
+		  FROM annotation
+		  WHERE dashboard_uid IS NULL
+		  AND dashboard_id != 0
+		  ORDER BY id DESC
+		  LIMIT ?
+		) AS batch ON batch.id = a.id
+		SET a.dashboard_uid = d.uid
+		  WHERE a.dashboard_uid IS NULL
+		  AND a.dashboard_id != 0`
 	}
 
+	updatedTotal := int64(0)
+	batchNum := 0
+	for {
+		batchNum++
+		result, err := sess.Exec(updateSQL, batchSize)
+		if err != nil {
+			return fmt.Errorf("failed to set dashboard_uid for annotation batch %d: %w", batchNum, err)
+		}
+
+		rowsAffected, err := result.RowsAffected()
+		if err != nil {
+			return fmt.Errorf("failed to get rows affected for batch %d: %w", batchNum, err)
+		}
+
+		if rowsAffected == 0 {
+			break
+		}
+
+		updatedTotal += rowsAffected
+		logger.Info("Updated annotation batch", "batch", batchNum, "rowsInBatch", rowsAffected, "totalUpdated", updatedTotal)
+
+		if rowsAffected < int64(batchSize) {
+			break
+		}
+	}
+
+	logger.Info("Completed dashboard_uid migration for annotations", "totalUpdated", updatedTotal)
 	return nil
 }