Unified: Run resource data migrations at startup (#114857)

* chore: uncomment unified migration

* chore: adapt and fix tests

* chore: dynamically bump max conns if needed during migration

* chore: copilot suggestions

* chore: pass ctx in RegisterMigration

* chore: make playlists opt-out and dashboards opt-in

* chore: adjust dashboard test

* chore: disable enable log in test

* chore: address review comments

- do not use pointer config
- add migration registry

* chore: more consistent naming

* chore: fix playlist discovery test

pkg/storage/unified/migrations/migrator_test.go

@@ -162,14 +162,49 @@ func runMigrationTestSuite(t *testing.T, testCases []resourceMigratorTestCase) {
 		}
 	})
 
-	t.Run("Step 3: verify data is migrated to unified storage", func(t *testing.T) {
-		// Migrations will run automatically at startup and mode 5 is enforced by the config
+	t.Run("Step 3: verify that opted-out resources are not migrated", func(t *testing.T) {
+		// Build unified storage config for Mode5
+		unifiedConfig := make(map[string]setting.UnifiedStorageConfig)
+		for _, tc := range testCases {
+			for _, gvr := range tc.resources() {
+				resourceKey := fmt.Sprintf("%s.%s", gvr.Resource, gvr.Group)
+				unifiedConfig[resourceKey] = setting.UnifiedStorageConfig{
+					DualWriterMode:  grafanarest.Mode5,
+					EnableMigration: false,
+				}
+			}
+		}
+
+		helper := apis.NewK8sTestHelperWithOpts(t, apis.K8sTestHelperOpts{
+			GrafanaOpts: testinfra.GrafanaOpts{
+				AppModeProduction:    true,
+				DisableAnonymous:     true,
+				DisableDBCleanup:     true,
+				APIServerStorageType: "unified",
+				UnifiedStorageConfig: unifiedConfig,
+			},
+			Org1Users: org1,
+			OrgBUsers: orgB,
+		})
+		t.Cleanup(helper.Shutdown)
+
+		for _, state := range testStates {
+			t.Run(state.tc.name(), func(t *testing.T) {
+				// Verify resources don't exist in unified storage yet
+				state.tc.verify(t, helper, false)
+			})
+		}
+		verifyRegisteredMigrations(t, helper, false, true)
+	})
+
+	t.Run("Step 4: verify data is migrated to unified storage", func(t *testing.T) {
+		// Migrations enabled by default will run automatically at startup and mode 5 is enforced by the config
 		helper := apis.NewK8sTestHelperWithOpts(t, apis.K8sTestHelperOpts{
 			GrafanaOpts: testinfra.GrafanaOpts{
-				// EnableLog:             true,
 				AppModeProduction:     true,
 				DisableAnonymous:      true,
 				DisableDataMigrations: false, // Run migrations at startup
+				DisableDBCleanup:      true,
 				APIServerStorageType:  "unified",
 			},
 			Org1Users: org1,
@@ -183,7 +218,89 @@ func runMigrationTestSuite(t *testing.T, testCases []resourceMigratorTestCase) {
 				state.tc.verify(t, helper, true)
 			})
 		}
+
+		t.Logf("Verifying migrations are correctly registered")
+		verifyRegisteredMigrations(t, helper, true, false)
 	})
+
+	t.Run("Step 5: verify data is migrated for all migrations", func(t *testing.T) {
+		// Trigger migrations that are not enabled by default
+		unifiedConfig := make(map[string]setting.UnifiedStorageConfig)
+		for _, tc := range testCases {
+			for _, gvr := range tc.resources() {
+				resourceKey := fmt.Sprintf("%s.%s", gvr.Resource, gvr.Group)
+				unifiedConfig[resourceKey] = setting.UnifiedStorageConfig{
+					EnableMigration: true,
+				}
+			}
+		}
+		helper := apis.NewK8sTestHelperWithOpts(t, apis.K8sTestHelperOpts{
+			GrafanaOpts: testinfra.GrafanaOpts{
+				// EnableLog:             true,
+				AppModeProduction:     true,
+				DisableAnonymous:      true,
+				DisableDataMigrations: false,
+				APIServerStorageType:  "unified",
+				UnifiedStorageConfig:  unifiedConfig,
+			},
+			Org1Users: org1,
+			OrgBUsers: orgB,
+		})
+		t.Cleanup(helper.Shutdown)
+
+		for _, state := range testStates {
+			t.Run(state.tc.name(), func(t *testing.T) {
+				// Verify resources still exist in unified storage after restart
+				state.tc.verify(t, helper, true)
+			})
+		}
+
+		t.Logf("Verifying migrations are correctly registered")
+		verifyRegisteredMigrations(t, helper, false, false)
+	})
+}
+
+const (
+	migrationScope = "unifiedstorage"
+	migrationTable = migrationScope + "_migration_log"
+
+	playlistsID            = "playlists migration"
+	foldersAndDashboardsID = "folders and dashboards migration"
+)
+
+var migrationIDsToDefault = map[string]bool{
+	playlistsID:            true,
+	foldersAndDashboardsID: false,
+}
+
+func verifyRegisteredMigrations(t *testing.T, helper *apis.K8sTestHelper, onlyDefault bool, optOut bool) {
+	getMigrationsQuery := fmt.Sprintf("SELECT migration_id FROM %s", migrationTable)
+	createTableMigrationID := fmt.Sprintf("create %s table", migrationTable)
+	expectedMigrationIDs := []string{createTableMigrationID}
+	for id, enabled := range migrationIDsToDefault {
+		if onlyDefault && !enabled {
+			continue
+		}
+		if optOut {
+			continue
+		}
+		expectedMigrationIDs = append(expectedMigrationIDs, id)
+	}
+	rows, err := helper.GetEnv().SQLStore.GetEngine().DB().Query(getMigrationsQuery)
+	require.NoError(t, err)
+	defer func() {
+		require.NoError(t, rows.Close())
+	}()
+
+	migrationIDs := make(map[string]struct{})
+	for rows.Next() {
+		var migrationID string
+		require.NoError(t, rows.Scan(&migrationID))
+		require.Contains(t, expectedMigrationIDs, migrationID)
+		migrationIDs[migrationID] = struct{}{}
+	}
+	require.NoError(t, rows.Err())
+	require.Len(t, migrationIDs, len(expectedMigrationIDs))
 }
 
 // verifyResourceCount verifies that the expected number of resources exist in K8s storage

pkg/storage/unified/migrations/resources.go

@@ -7,7 +7,9 @@ import (
 	folders "github.com/grafana/grafana/apps/folder/pkg/apis/folder/v1beta1"
 	playlists "github.com/grafana/grafana/apps/playlist/pkg/apis/playlist/v0alpha1"
 	"github.com/grafana/grafana/pkg/registry/apis/dashboard/legacy"
+	sqlstoremigrator "github.com/grafana/grafana/pkg/services/sqlstore/migrator"
 	"github.com/grafana/grafana/pkg/setting"
+	"github.com/grafana/grafana/pkg/storage/unified/resource"
 	"github.com/grafana/grafana/pkg/storage/unified/resourcepb"
 	"k8s.io/apimachinery/pkg/runtime/schema"
 )
@@ -17,7 +19,13 @@ type ResourceDefinition struct {
 	MigratorFunc  string // Name of the method: "MigrateFolders", "MigrateDashboards", etc.
 }
 
-var registeredResources = []ResourceDefinition{
+type migrationDefinition struct {
+	name         string
+	resources    []string
+	registerFunc func(mg *sqlstoremigrator.Migrator, migrator UnifiedMigrator, client resource.ResourceClient)
+}
+
+var resourceRegistry = []ResourceDefinition{
 	{
 		GroupResource: schema.GroupResource{Group: folders.GROUP, Resource: folders.RESOURCE},
 		MigratorFunc:  "MigrateFolders",
@@ -36,9 +44,50 @@ var registeredResources = []ResourceDefinition{
 	},
 }
 
+var migrationRegistry = []migrationDefinition{
+	{
+		name:         "playlists",
+		resources:    []string{setting.PlaylistResource},
+		registerFunc: registerPlaylistMigration,
+	},
+	{
+		name:         "folders and dashboards",
+		resources:    []string{setting.FolderResource, setting.DashboardResource},
+		registerFunc: registerDashboardAndFolderMigration,
+	},
+}
+
+func registerMigrations(cfg *setting.Cfg, mg *sqlstoremigrator.Migrator, migrator UnifiedMigrator, client resource.ResourceClient) error {
+	for _, migration := range migrationRegistry {
+		var (
+			hasValue   bool
+			allEnabled bool
+		)
+
+		for _, res := range migration.resources {
+			enabled := cfg.UnifiedStorage[res].EnableMigration
+			if !hasValue {
+				allEnabled = enabled
+				hasValue = true
+				continue
+			}
+			if enabled != allEnabled {
+				return fmt.Errorf("cannot migrate resources separately: %v migration must be either all enabled or all disabled", migration.resources)
+			}
+		}
+
+		if !allEnabled {
+			logger.Info("Migration is disabled in config, skipping", "migration", migration.name)
+			continue
+		}
+		migration.registerFunc(mg, migrator, client)
+	}
+	return nil
+}
+
 func getResourceDefinition(group, resource string) *ResourceDefinition {
-	for i := range registeredResources {
-		r := &registeredResources[i]
+	for i := range resourceRegistry {
+		r := &resourceRegistry[i]
 		if r.GroupResource.Group == group && r.GroupResource.Resource == resource {
 			return r
 		}
@@ -80,13 +129,13 @@ func getMigratorFunc(accessor legacy.MigrationDashboardAccessor, group, resource
 
 func validateRegisteredResources() error {
 	registeredMap := make(map[string]bool)
-	for _, gr := range registeredResources {
+	for _, gr := range resourceRegistry {
 		key := fmt.Sprintf("%s.%s", gr.GroupResource.Resource, gr.GroupResource.Group)
 		registeredMap[key] = true
 	}
 
 	var missing []string
-	for _, expected := range setting.MigratedUnifiedResources {
+	for expected := range setting.MigratedUnifiedResources {
 		if !registeredMap[expected] {
 			missing = append(missing, expected)
 		}

pkg/storage/unified/migrations/resources_test.go

@@ -0,0 +1,92 @@
+package migrations
+
+import (
+	"testing"
+
+	sqlstoremigrator "github.com/grafana/grafana/pkg/services/sqlstore/migrator"
+	"github.com/grafana/grafana/pkg/setting"
+	"github.com/grafana/grafana/pkg/storage/unified/resource"
+	"github.com/stretchr/testify/require"
+)
+
+// TestRegisterMigrations exercises registerMigrations with various EnableMigration configs using a table-driven test.
+func TestRegisterMigrations(t *testing.T) {
+	origRegistry := migrationRegistry
+	t.Cleanup(func() { migrationRegistry = origRegistry })
+
+	// helper to build a fake registry with custom register funcs that bump counters
+	makeFakeRegistry := func(migrationCalls map[string]int) []migrationDefinition {
+		return []migrationDefinition{
+			{
+				name:      "playlists",
+				resources: []string{setting.PlaylistResource},
+				registerFunc: func(mg *sqlstoremigrator.Migrator, migrator UnifiedMigrator, client resource.ResourceClient) {
+					migrationCalls["playlists"]++
+				},
+			},
+			{
+				name:      "folders and dashboards",
+				resources: []string{setting.FolderResource, setting.DashboardResource},
+				registerFunc: func(mg *sqlstoremigrator.Migrator, migrator UnifiedMigrator, client resource.ResourceClient) {
+					migrationCalls["folders and dashboards"]++
+				},
+			},
+		}
+	}
+
+	// Build a minimal cfg with UnifiedStorage entries used by registerMigrations
+	makeCfg := func(vals map[string]bool) *setting.Cfg {
+		cfg := &setting.Cfg{UnifiedStorage: make(map[string]setting.UnifiedStorageConfig)}
+		for k, v := range vals {
+			cfg.UnifiedStorage[k] = setting.UnifiedStorageConfig{EnableMigration: v}
+		}
+		return cfg
+	}
+
+	// Table of scenarios
+	tests := []struct {
+		name              string
+		enablePlaylist    bool
+		enableFolder      bool
+		enableDashboard   bool
+		wantPlaylistCalls int
+		wantFDCalls       int
+		wantErr           bool
+	}{
+		{name: "playlists enabled", enablePlaylist: true, wantPlaylistCalls: 1},
+		{name: "playlists disabled", enablePlaylist: false, wantPlaylistCalls: 0},
+		{name: "folders+dashboards both enabled", enableFolder: true, enableDashboard: true, wantFDCalls: 1},
+		{name: "folders enabled, dashboards disabled (mismatch)", enableFolder: true, enableDashboard: false, wantFDCalls: 0, wantErr: true},
+		{name: "folders disabled, dashboards enabled (mismatch)", enableFolder: false, enableDashboard: true, wantFDCalls: 0, wantErr: true},
+		{name: "folders+dashboards both disabled", enableFolder: false, enableDashboard: false, wantFDCalls: 0},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			migrationCalls := map[string]int{
+				"playlists":              0,
+				"folders and dashboards": 0,
+			}
+
+			migrationRegistry = makeFakeRegistry(migrationCalls)
+
+			cfg := makeCfg(map[string]bool{
+				setting.PlaylistResource:  tt.enablePlaylist,
+				setting.FolderResource:    tt.enableFolder,
+				setting.DashboardResource: tt.enableDashboard,
+			})
+
+			// We pass nils for migrator dependencies because our fake registerFuncs don't use them
+			err := registerMigrations(cfg, nil, nil, nil)
+
+			if tt.wantErr {
+				require.Error(t, err, "expected error for mismatched enablement")
+			} else {
+				require.NoError(t, err, "unexpected error")
+			}
+
+			require.Equal(t, tt.wantPlaylistCalls, migrationCalls["playlists"], "playlists register call count")
+			require.Equal(t, tt.wantFDCalls, migrationCalls["folders and dashboards"], "folders+dashboards register call count")
+		})
+	}
+}

pkg/storage/unified/migrations/service.go

@@ -3,7 +3,6 @@ package migrations
 import (
 	"context"
 	"fmt"
-	"os"
 
 	"github.com/grafana/grafana/pkg/infra/db"
 	"github.com/grafana/grafana/pkg/infra/kvstore"
@@ -49,34 +48,25 @@ func ProvideUnifiedStorageMigrationService(
 }
 
 func (p *UnifiedStorageMigrationServiceImpl) Run(ctx context.Context) error {
-	// TODO: temporary skip migrations in test environments to prevent integration test timeouts.
-	if os.Getenv("GRAFANA_TEST_DB") != "" {
-		return nil
-	}
-
 	// skip migrations if disabled in config
 	if p.cfg.DisableDataMigrations {
 		metrics.MUnifiedStorageMigrationStatus.Set(1)
 		logger.Info("Data migrations are disabled, skipping")
 		return nil
-	} else {
-		metrics.MUnifiedStorageMigrationStatus.Set(2)
-		logger.Info("Data migrations not yet enforced, skipping")
 	}
-
-	// TODO: Re-enable once migrations are ready
-	// TODO: add guarantee that this only runs once
-	// return RegisterMigrations(p.migrator, p.cfg, p.sqlStore, p.client)
-	return nil
+	logger.Info("Running migrations for unified storage")
+	metrics.MUnifiedStorageMigrationStatus.Set(3)
+	return RegisterMigrations(ctx, p.migrator, p.cfg, p.sqlStore, p.client)
 }
 
 func RegisterMigrations(
+	ctx context.Context,
 	migrator UnifiedMigrator,
 	cfg *setting.Cfg,
 	sqlStore db.DB,
 	client resource.ResourceClient,
 ) error {
-	ctx, span := tracer.Start(context.Background(), "storage.unified.RegisterMigrations")
+	ctx, span := tracer.Start(ctx, "storage.unified.RegisterMigrations")
 	defer span.End()
 	mg := sqlstoremigrator.NewScopedMigrator(sqlStore.GetEngine(), cfg, "unifiedstorage")
 	mg.AddCreateMigration()
@@ -89,12 +79,19 @@ func RegisterMigrations(
 		return err
 	}
 
-	// Register resource migrations
-	registerDashboardAndFolderMigration(mg, migrator, client)
-	registerPlaylistMigration(mg, migrator, client)
+	if err := registerMigrations(cfg, mg, migrator, client); err != nil {
+		return err
+	}
 
 	// Run all registered migrations (blocking)
 	sec := cfg.Raw.Section("database")
+	db := mg.DBEngine.DB().DB
+	maxOpenConns := db.Stats().MaxOpenConnections
+	if maxOpenConns <= 2 {
+		// migrations require at least 3 connections due to extra GRPC connections
+		db.SetMaxOpenConns(3)
+		defer db.SetMaxOpenConns(maxOpenConns)
+	}
 	if err := mg.RunMigrations(ctx,
 		sec.Key("migration_locking").MustBool(true),
 		sec.Key("locking_attempt_timeout_sec").MustInt()); err != nil {