[release-12.3.1] refactor(annotations): Allow skipping always on dashboard UID migrations (#114096)

refactor(annotations): Allow skipping always on dashboard UID migrations (#113780)

(cherry picked from commit b70c6a726f)

Co-authored-by: Andres Torres <janthoe@users.noreply.github.com>

pkg/services/annotations/annotationsimpl/xorm_store.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
@@ -40,6 +41,8 @@ func validateTimeRange(item *annotations.Item) error {
 	return nil
 }
 
+var xormMigrationTrigger sync.Once
+
 type xormRepositoryImpl struct {
 	cfg                *setting.Cfg
 	db                 db.DB
@@ -51,10 +54,9 @@ type xormRepositoryImpl struct {
 }
 
 func NewXormStore(cfg *setting.Cfg, l log.Logger, db db.DB, tagService tag.Service, reg prometheus.Registerer) *xormRepositoryImpl {
-	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)
+	})
 
 	repo := &xormRepositoryImpl{
 		cfg:        cfg,
@@ -96,6 +98,19 @@ func NewXormStore(cfg *setting.Cfg, l log.Logger, db db.DB, tagService tag.Servi
 	return repo
 }
 
+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
+	}
+	err := migrations.RunDashboardUIDMigrations(db.GetEngine().NewSession(), db.GetEngine().DriverName())
+	if err != nil {
+		l.Error("failed to populate dashboard_uid for annotations", "error", err)
+	}
+}
+
 func (r *xormRepositoryImpl) Type() string {
 	return "sql"
 }

pkg/services/annotations/annotationsimpl/xorm_store_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"strings"
+	"sync"
 	"testing"
 
 	"github.com/prometheus/client_golang/prometheus"
@@ -654,6 +655,140 @@ func TestIntegrationAnnotations(t *testing.T) {
 	})
 }
 
+func TestIntegrationAnnotationsAlwaysOnMigrations(t *testing.T) {
+	tutil.SkipIntegrationTestInShortMode(t)
+
+	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), nil)
+		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), prometheus.NewRegistry())
+
+		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), nil)
+		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), prometheus.NewRegistry())
+
+		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)
+		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")
+	})
+}
+
 func BenchmarkFindTags_10k(b *testing.B) {
 	benchmarkFindTags(b, 10000)
 }