feat: add reindex trigger as a migration validation (#114283)

* feat: add search validation

* chore: remove comment

* fix: remove local test

* fix: revert rebuild index portion

* chore: remove unnecessary comments

* refactor: validator

* fix: revert setting

pkg/server/wire_gen.go

@@ -534,7 +534,7 @@ func Initialize(ctx context.Context, cfg *setting.Cfg, opts Options, apiOpts api
 	}
 	migrationDashboardAccessor := legacy.ProvideMigratorDashboardAccessor(legacyDatabaseProvider, stubProvisioningService, accessControl, featureToggles)
 	unifiedMigrator := migrations2.ProvideUnifiedMigrator(migrationDashboardAccessor, resourceClient)
-	unifiedStorageMigrationService := migrations2.ProvideUnifiedStorageMigrationService(unifiedMigrator, cfg, sqlStore, kvStore)
+	unifiedStorageMigrationService := migrations2.ProvideUnifiedStorageMigrationService(unifiedMigrator, cfg, sqlStore, kvStore, resourceClient)
 	dualwriteService, err := dualwrite.ProvideService(featureToggles, kvStore, cfg, unifiedStorageMigrationService)
 	if err != nil {
 		return nil, err
@@ -1181,7 +1181,7 @@ func InitializeForTest(ctx context.Context, t sqlutil.ITestDB, testingT interfac
 	}
 	migrationDashboardAccessor := legacy.ProvideMigratorDashboardAccessor(legacyDatabaseProvider, stubProvisioningService, accessControl, featureToggles)
 	unifiedMigrator := migrations2.ProvideUnifiedMigrator(migrationDashboardAccessor, resourceClient)
-	unifiedStorageMigrationService := migrations2.ProvideUnifiedStorageMigrationService(unifiedMigrator, cfg, sqlStore, kvStore)
+	unifiedStorageMigrationService := migrations2.ProvideUnifiedStorageMigrationService(unifiedMigrator, cfg, sqlStore, kvStore, resourceClient)
 	dualwriteService, err := dualwrite.ProvideService(featureToggles, kvStore, cfg, unifiedStorageMigrationService)
 	if err != nil {
 		return nil, err

pkg/storage/unified/migrations/migrator.go

@@ -144,7 +144,6 @@ func newUnifiedMigrator(
 	}
 }
 
-// migrate function -- works for a single kind
 type migratorFunc = func(ctx context.Context, orgId int64, opts legacy.MigrateOptions, stream resourcepb.BulkStore_BulkProcessClient) (*legacy.BlobStoreInfo, error)
 
 func (m *unifiedMigration) Migrate(ctx context.Context, opts legacy.MigrateOptions) (*resourcepb.BulkResponse, error) {

pkg/storage/unified/migrations/resource_migration.go

@@ -16,65 +16,34 @@ import (
 )
 
 // ValidationFunc is a function that validates migration results.
-// It receives the database session, migration response, and logger for reporting.
-// Return an error if validation fails, nil if validation passes or is skipped.
-type ValidationFunc func(sess *xorm.Session, response *resourcepb.BulkResponse, log log.Logger) error
+
+type Validator interface {
+	Validate(ctx context.Context, sess *xorm.Session, response *resourcepb.BulkResponse, log log.Logger) error
+}
 
 // ResourceMigration handles migration of specific resource types from legacy to unified storage.
-// It implements migrator.CodeMigration and provides a generic, extensible way to migrate any
-// resource type by:
-//
-//  1. Iterating through all organizations
-//  2. For each org, delegating to LegacyMigrator to read from legacy and write to unified storage
-//  3. Validating migration results using the provided validation function (if any)
-//
-// To add a new resource type migration, simply create a new ResourceMigration instance in
-// service.go with the appropriate schema.GroupResource specifications and optional validation function.
 type ResourceMigration struct {
 	migrator.MigrationBase
-	migrator       UnifiedMigrator
-	resources      []schema.GroupResource
-	migrationID    string
-	validationFunc ValidationFunc // Optional: custom validation logic for this migration
-	log            log.Logger
+	migrator    UnifiedMigrator
+	resources   []schema.GroupResource
+	migrationID string
+	validator   Validator // Optional: custom validation logic for this migration
+	log         log.Logger
 }
 
 // NewResourceMigration creates a new migration for the specified resources.
-// This is the primary way to register new resource migrations.
-//
-// Parameters:
-//   - legacyMigrator: handles reading from legacy storage and writing to unified storage
-//   - resources: list of GroupResource to migrate
-//   - migrationID: unique identifier for this migration
-//   - validationFunc: optional validation function to verify migration results.
-//     If nil, no validation will be performed.
-//
-// Example with legacy table count validation:
-//
-//	NewResourceMigration(
-//	    migrator,
-//	    []schema.GroupResource{{Group: "playlist.grafana.app", Resource: "playlists"}},
-//	    "playlists",
-//	    NewLegacyTableCountValidator(map[string]LegacyTableInfo{
-//	        "playlist.grafana.app/playlists": {Table: "playlist", WhereClause: "org_id = ?"},
-//	    }),
-//	)
-//
-// Example without validation:
-//
-//	NewResourceMigration(migrator, resources, "new-resource", nil)
 func NewResourceMigration(
 	migrator UnifiedMigrator,
 	resources []schema.GroupResource,
 	migrationID string,
-	validationFunc ValidationFunc,
+	validator Validator,
 ) *ResourceMigration {
 	return &ResourceMigration{
-		migrator:       migrator,
-		resources:      resources,
-		migrationID:    migrationID,
-		validationFunc: validationFunc,
-		log:            log.New("storage.unified.resource_migration." + migrationID),
+		migrator:    migrator,
+		resources:   resources,
+		migrationID: migrationID,
+		validator:   validator,
+		log:         log.New("storage.unified.resource_migration." + migrationID),
 	}
 }
 
@@ -139,7 +108,7 @@ func (m *ResourceMigration) migrateOrg(ctx context.Context, sess *xorm.Session,
 	}
 
 	// Validate the migration results
-	if err := m.validateMigration(sess, response); err != nil {
+	if err := m.validateMigration(migrationCtx, sess, response); err != nil {
 		m.log.Error("Migration validation failed", "org_id", org.ID, "error", err, "duration", time.Since(startTime))
 		return fmt.Errorf("migration validation failed for org %d (%s): %w", org.ID, org.Name, err)
 	}
@@ -155,13 +124,13 @@ func (m *ResourceMigration) migrateOrg(ctx context.Context, sess *xorm.Session,
 }
 
 // validateMigration calls the custom validation function if provided
-func (m *ResourceMigration) validateMigration(sess *xorm.Session, response *resourcepb.BulkResponse) error {
-	if m.validationFunc == nil {
+func (m *ResourceMigration) validateMigration(ctx context.Context, sess *xorm.Session, response *resourcepb.BulkResponse) error {
+	if m.validator == nil {
 		m.log.Debug("No validation function provided, skipping validation")
 		return nil
 	}
 
-	return m.validationFunc(sess, response, m.log)
+	return m.validator.Validate(ctx, sess, response, m.log)
 }
 
 // LegacyTableInfo defines how to map a unified storage resource to its legacy table
@@ -170,85 +139,6 @@ type LegacyTableInfo struct {
 	WhereClause string // WHERE clause template with org_id parameter (e.g., "org_id = ? and is_folder = false")
 }
 
-// NewLegacyTableCountValidator creates a ValidationFunc that validates migration by comparing
-// counts between legacy tables and unified storage.
-//
-// This is a helper for the common case of validating that all items from legacy tables
-// were successfully migrated to unified storage.
-//
-// Parameters:
-//   - legacyTableMap: maps "group/resource" keys to LegacyTableInfo for validation.
-//     Only resources with mappings will be validated.
-//
-// Example:
-//
-//	validator := NewLegacyTableCountValidator(map[string]LegacyTableInfo{
-//	    "dashboard.grafana.app/dashboards": {Table: "dashboard", WhereClause: "org_id = ? and is_folder = false"},
-//	    "folder.grafana.app/folders": {Table: "dashboard", WhereClause: "org_id = ? and is_folder = true"},
-//	})
-func NewLegacyTableCountValidator(legacyTableMap map[string]LegacyTableInfo) ValidationFunc {
-	return func(sess *xorm.Session, response *resourcepb.BulkResponse, log log.Logger) error {
-		// Check for rejected items
-		if len(response.Rejected) > 0 {
-			log.Warn("Migration had rejected items", "count", len(response.Rejected))
-			for i, rejected := range response.Rejected {
-				if i < 10 { // Log first 10 rejected items
-					log.Warn("Rejected item",
-						"namespace", rejected.Key.Namespace,
-						"group", rejected.Key.Group,
-						"resource", rejected.Key.Resource,
-						"name", rejected.Key.Name,
-						"reason", rejected.Error)
-				}
-			}
-			// Rejections are not fatal - they may be expected for invalid data
-		}
-
-		// Validate counts for each resource type
-		for _, summary := range response.Summary {
-			key := fmt.Sprintf("%s/%s", summary.Group, summary.Resource)
-			tableInfo, ok := legacyTableMap[key]
-			if !ok {
-				log.Debug("No legacy table mapping for resource, skipping count validation",
-					"resource", fmt.Sprintf("%s.%s", summary.Resource, summary.Group),
-					"namespace", summary.Namespace)
-				continue
-			}
-
-			// Get legacy count
-			orgID, err := ParseOrgIDFromNamespace(summary.Namespace)
-			if err != nil {
-				return fmt.Errorf("invalid namespace %s: %w", summary.Namespace, err)
-			}
-
-			legacyCount, err := sess.Table(tableInfo.Table).Where(tableInfo.WhereClause, orgID).Count()
-			if err != nil {
-				return fmt.Errorf("failed to count %s: %w", tableInfo.Table, err)
-			}
-
-			// Account for rejected items in validation
-			expectedCount := summary.Count + int64(len(response.Rejected))
-
-			log.Info("Count validation",
-				"resource", fmt.Sprintf("%s.%s", summary.Resource, summary.Group),
-				"namespace", summary.Namespace,
-				"legacy_count", legacyCount,
-				"unified_count", summary.Count,
-				"rejected", len(response.Rejected),
-				"history", summary.History)
-
-			// Validate that we migrated all items (allowing for rejected items)
-			if legacyCount > expectedCount {
-				return fmt.Errorf("count mismatch for %s.%s in namespace %s: legacy has %d, unified has %d, rejected %d",
-					summary.Resource, summary.Group, summary.Namespace,
-					legacyCount, summary.Count, len(response.Rejected))
-			}
-		}
-
-		return nil
-	}
-}
-
 func ParseOrgIDFromNamespace(namespace string) (int64, error) {
 	// Use authlib to properly parse all namespace formats including "default" for org 1
 	info, err := types.ParseNamespace(namespace)

pkg/storage/unified/migrations/service.go

@@ -11,6 +11,7 @@ import (
 	sqlstoremigrator "github.com/grafana/grafana/pkg/services/sqlstore/migrator"
 	"github.com/grafana/grafana/pkg/setting"
 	"github.com/grafana/grafana/pkg/storage/unified/migrations/contract"
+	"github.com/grafana/grafana/pkg/storage/unified/resource"
 	"github.com/prometheus/client_golang/prometheus"
 	"go.opentelemetry.io/otel"
 	"k8s.io/apimachinery/pkg/runtime/schema"
@@ -24,29 +25,28 @@ type UnifiedStorageMigrationServiceImpl struct {
 	cfg      *setting.Cfg
 	sqlStore db.DB
 	kv       kvstore.KVStore
+	client   resource.ResourceClient
 }
 
 var _ contract.UnifiedStorageMigrationService = (*UnifiedStorageMigrationServiceImpl)(nil)
 
 // ProvideUnifiedStorageMigrationService is a Wire provider that creates the migration service.
-// The service implements registry.BackgroundService and runs migrations during server startup.
 func ProvideUnifiedStorageMigrationService(
 	migrator UnifiedMigrator,
 	cfg *setting.Cfg,
 	sqlStore db.DB,
 	kv kvstore.KVStore,
+	client resource.ResourceClient,
 ) contract.UnifiedStorageMigrationService {
 	return &UnifiedStorageMigrationServiceImpl{
 		migrator: migrator,
 		cfg:      cfg,
 		sqlStore: sqlStore,
 		kv:       kv,
+		client:   client,
 	}
 }
 
-// Run executes unified storage migrations as a background service.
-// This blocks until migrations complete. If migrations fail, an error is returned
-// which will prevent Grafana from starting.
 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") != "" {
@@ -61,18 +61,15 @@ func (p *UnifiedStorageMigrationServiceImpl) Run(ctx context.Context) error {
 
 	// TODO: Re-enable once migrations are ready
 	// TODO: add guarantee that this only runs once
-	// return RegisterMigrations(p.migrator, p.cfg, p.sqlStore)
+	// return RegisterMigrations(p.migrator, p.cfg, p.sqlStore, p.client)
 	return nil
 }
 
-// RegisterMigrations initializes and registers all unified storage migrations.
-// This function is the entry point for all data migrations from legacy storage
-// to unified storage. It returns an error if migrations fail, preventing Grafana
-// from starting with inconsistent data.
 func RegisterMigrations(
 	migrator UnifiedMigrator,
 	cfg *setting.Cfg,
 	sqlStore db.DB,
+	client resource.ResourceClient,
 ) error {
 	ctx, span := tracer.Start(context.Background(), "storage.unified.RegisterMigrations")
 	defer span.End()
@@ -85,7 +82,7 @@ func RegisterMigrations(
 
 	// Register resource migrations
 	// To add a new resource type, simply add another migration here with the appropriate resources
-	registerResourceMigrations(mg, migrator)
+	registerResourceMigrations(mg, migrator, client)
 
 	// Run all registered migrations (blocking)
 	sec := cfg.Raw.Section("database")
@@ -99,9 +96,7 @@ func RegisterMigrations(
 	return nil
 }
 
-// registerResourceMigrations registers all unified storage resource migrations.
-// Add new resource types here by creating additional ResourceMigration instances.
-func registerResourceMigrations(mg *sqlstoremigrator.Migrator, migrator UnifiedMigrator) {
+func registerResourceMigrations(mg *sqlstoremigrator.Migrator, migrator UnifiedMigrator, client resource.ResourceClient) {
 	dashboardsAndFolders := NewResourceMigration(
 		migrator,
 		[]schema.GroupResource{
@@ -109,7 +104,7 @@ func registerResourceMigrations(mg *sqlstoremigrator.Migrator, migrator UnifiedM
 			{Group: "dashboard.grafana.app", Resource: "dashboards"},
 		},
 		"folders-dashboards",
-		NewLegacyTableCountValidator(map[string]LegacyTableInfo{
+		NewCountValidator(client, map[string]LegacyTableInfo{
 			"folder.grafana.app/folders":       {Table: "dashboard", WhereClause: "org_id = ? and is_folder = true"},
 			"dashboard.grafana.app/dashboards": {Table: "dashboard", WhereClause: "org_id = ? and is_folder = false"},
 		}),

pkg/storage/unified/migrations/validator.go

@@ -0,0 +1,99 @@
+package migrations
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/storage/unified/resourcepb"
+	"github.com/grafana/grafana/pkg/util/xorm"
+)
+
+type CountValidator struct {
+	client         resourcepb.ResourceIndexClient
+	legacyTableMap map[string]LegacyTableInfo
+}
+
+func NewCountValidator(client resourcepb.ResourceIndexClient, legacyTableMap map[string]LegacyTableInfo) Validator {
+	return &CountValidator{client: client, legacyTableMap: legacyTableMap}
+}
+
+func (v *CountValidator) Validate(ctx context.Context, sess *xorm.Session, response *resourcepb.BulkResponse, log log.Logger) error {
+	if len(response.Rejected) > 0 {
+		log.Warn("Migration had rejected items", "count", len(response.Rejected))
+		for i, rejected := range response.Rejected {
+			if i < 10 { // Log first 10 rejected items
+				log.Warn("Rejected item",
+					"namespace", rejected.Key.Namespace,
+					"group", rejected.Key.Group,
+					"resource", rejected.Key.Resource,
+					"name", rejected.Key.Name,
+					"reason", rejected.Error)
+			}
+		}
+		// Rejections are not fatal - they may be expected for invalid data
+	}
+
+	// Validate counts for each resource type
+	for _, summary := range response.Summary {
+		key := fmt.Sprintf("%s/%s", summary.Group, summary.Resource)
+		tableInfo, ok := v.legacyTableMap[key]
+		if !ok {
+			log.Debug("No legacy table mapping for resource, skipping count validation",
+				"resource", fmt.Sprintf("%s.%s", summary.Resource, summary.Group),
+				"namespace", summary.Namespace)
+			continue
+		}
+
+		// Get legacy count from database
+		orgID, err := ParseOrgIDFromNamespace(summary.Namespace)
+		if err != nil {
+			return fmt.Errorf("invalid namespace %s: %w", summary.Namespace, err)
+		}
+
+		legacyCount, err := sess.Table(tableInfo.Table).Where(tableInfo.WhereClause, orgID).Count()
+		if err != nil {
+			return fmt.Errorf("failed to count %s: %w", tableInfo.Table, err)
+		}
+
+		// Get unified storage count using GetStats API
+		statsResp, err := v.client.GetStats(ctx, &resourcepb.ResourceStatsRequest{
+			Namespace: summary.Namespace,
+			Kinds:     []string{fmt.Sprintf("%s/%s", summary.Group, summary.Resource)},
+		})
+		if err != nil {
+			return fmt.Errorf("failed to get stats for %s/%s in namespace %s: %w",
+				summary.Group, summary.Resource, summary.Namespace, err)
+		}
+
+		// Find the count for this specific resource type
+		var unifiedCount int64
+		for _, stat := range statsResp.Stats {
+			if stat.Group == summary.Group && stat.Resource == summary.Resource {
+				unifiedCount = stat.Count
+				break
+			}
+		}
+
+		// Account for rejected items in validation
+		expectedCount := unifiedCount + int64(len(response.Rejected))
+
+		log.Info("Count validation",
+			"resource", fmt.Sprintf("%s.%s", summary.Resource, summary.Group),
+			"namespace", summary.Namespace,
+			"legacy_count", legacyCount,
+			"unified_count", unifiedCount,
+			"migration_summary_count", summary.Count,
+			"rejected", len(response.Rejected),
+			"history", summary.History)
+
+		// Validate that we migrated all items (allowing for rejected items)
+		if legacyCount > expectedCount {
+			return fmt.Errorf("count mismatch for %s.%s in namespace %s: legacy has %d, unified has %d, rejected %d",
+				summary.Resource, summary.Group, summary.Namespace,
+				legacyCount, unifiedCount, len(response.Rejected))
+		}
+	}
+
+	return nil
+}