From d2f7a7efc08173e0f014f42c8e62e958f60430a2 Mon Sep 17 00:00:00 2001 From: Serge Zaitsev Date: Mon, 15 Dec 2025 19:15:38 +0800 Subject: [PATCH] [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 95e65e258878545856688c57d24a609d5db46734) * fix broken cherry pick * update mysql query --- .../annotations/annotationsimpl/xorm_store.go | 31 +++- .../annotationsimpl/xorm_store_test.go | 145 +++++++++++++++++- .../sqlstore/migrations/annotation_mig.go | 87 +++++++++-- 3 files changed, 240 insertions(+), 23 deletions(-) diff --git a/pkg/services/annotations/annotationsimpl/xorm_store.go b/pkg/services/annotations/annotationsimpl/xorm_store.go index 998236e19ea..b432a8daddd 100644 --- a/pkg/services/annotations/annotationsimpl/xorm_store.go +++ b/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" } diff --git a/pkg/services/annotations/annotationsimpl/xorm_store_test.go b/pkg/services/annotations/annotationsimpl/xorm_store_test.go index 015c42a5075..09825d24a0e 100644 --- a/pkg/services/annotations/annotationsimpl/xorm_store_test.go +++ b/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) } diff --git a/pkg/services/sqlstore/migrations/annotation_mig.go b/pkg/services/sqlstore/migrations/annotation_mig.go index 896ef21aa8d..d1cd33b108f 100644 --- a/pkg/services/sqlstore/migrations/annotation_mig.go +++ b/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 }