[release-12.3.1] Chore: Run annotation data migration in batches (#115136)

* 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) * update mysql query --------- Co-authored-by: Serge Zaitsev <serge.zaitsev@grafana.com>
2025-12-15 12:12:31 +01:00
parent 1214d51a42
commit 0e38140bd9
3 changed files with 101 additions and 31 deletions
@@ -105,10 +105,16 @@ func triggerAlwaysOnMigrations(cfg *setting.Cfg, l log.Logger, db db.DB) {
 		l.Debug("skipped dashboard UID startup migration")
 		return
 	}
-	err := migrations.RunDashboardUIDMigrations(db.GetEngine().NewSession(), db.GetEngine().DriverName())
-	if err != nil {
-		l.Error("failed to populate dashboard_uid for annotations", "error", err)
-	}
+	// 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 {
@@ -6,6 +6,7 @@ import (
 	"strings"
 	"sync"
 	"testing"
+	"time"

 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stretchr/testify/assert"
@@ -768,24 +769,28 @@ func TestIntegrationAnnotationsAlwaysOnMigrations(t *testing.T) {
 		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"`
 		}
-		err = sql.WithDbSession(context.Background(), func(sess *db.Session) error {
-			has, err := sess.Table("annotation").
-				Where("id = ?", annotation.ID).
-				Get(&result)
+		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 err
+				return false
 			}
-			if !has {
-				return fmt.Errorf("annotation not found")
-			}
-			return nil
-		})
-		require.NoError(t, err)
-		require.NotNil(t, result.DashboardUID, "dashboard_uid should not be NULL when migration runs")
-		assert.Equal(t, "test-run-uid", *result.DashboardUID, "dashboard_uid should be populated when migration runs")
+			return result.DashboardUID != nil && *result.DashboardUID == "test-run-uid"
+		}, 5*time.Second, 100*time.Millisecond, "dashboard_uid should be populated when migration runs")
 	})
 }

@@ -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
 }