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

2025-11-14 11:34:13 -05:00
parent e6eb40a679
commit b70c6a726f
5 changed files with 166 additions and 4 deletions
@@ -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)
 }