grafana/pkg/storage/unified/migrations/resources.go

package migrations

import (
	"context"
	"fmt"

	v1beta1 "github.com/grafana/grafana/apps/dashboard/pkg/apis/dashboard/v1beta1"
	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/infra/db"
	"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"
)

type resourceDefinition struct {
	groupResource schema.GroupResource
	migratorFunc  string // Name of the method: "MigrateFolders", "MigrateDashboards", etc.
}

type migrationDefinition struct {
	name         string
	migrationID  string // The ID stored in the migration log table (e.g., "playlists migration")
	resources    []string
	registerFunc func(mg *sqlstoremigrator.Migrator, migrator UnifiedMigrator, client resource.ResourceClient, opts ...ResourceMigrationOption)
}

var resourceRegistry = []resourceDefinition{
	{
		groupResource: schema.GroupResource{Group: folders.GROUP, Resource: folders.RESOURCE},
		migratorFunc:  "MigrateFolders",
	},
	{
		groupResource: schema.GroupResource{Group: v1beta1.GROUP, Resource: v1beta1.LIBRARY_PANEL_RESOURCE},
		migratorFunc:  "MigrateLibraryPanels",
	},
	{
		groupResource: schema.GroupResource{Group: v1beta1.GROUP, Resource: v1beta1.DASHBOARD_RESOURCE},
		migratorFunc:  "MigrateDashboards",
	},
	{
		groupResource: schema.GroupResource{Group: playlists.APIGroup, Resource: "playlists"},
		migratorFunc:  "MigratePlaylists",
	},
}

var migrationRegistry = []migrationDefinition{
	{
		name:         "playlists",
		migrationID:  "playlists migration",
		resources:    []string{setting.PlaylistResource},
		registerFunc: registerPlaylistMigration,
	},
	{
		name:         "folders and dashboards",
		migrationID:  "folders and dashboards migration",
		resources:    []string{setting.FolderResource, setting.DashboardResource},
		registerFunc: registerDashboardAndFolderMigration,
	},
}

func registerMigrations(ctx context.Context,
	cfg *setting.Cfg,
	mg *sqlstoremigrator.Migrator,
	migrator UnifiedMigrator,
	client resource.ResourceClient,
	sqlStore db.DB,
) error {
	for _, migration := range migrationRegistry {
		if shouldAutoMigrate(ctx, migration, cfg, sqlStore) {
			migration.registerFunc(mg, migrator, client, WithAutoMigrate(cfg))
			continue
		}

		enabled, err := isMigrationEnabled(migration, cfg)
		if err != nil {
			return err
		}
		if !enabled {
			logger.Info("Migration is disabled in config, skipping", "migration", migration.name)
			continue
		}
		migration.registerFunc(mg, migrator, client)
	}
	return nil
}

func registerDashboardAndFolderMigration(mg *sqlstoremigrator.Migrator,
	migrator UnifiedMigrator,
	client resource.ResourceClient,
	opts ...ResourceMigrationOption,
) {
	foldersDef := getResourceDefinition("folder.grafana.app", "folders")
	dashboardsDef := getResourceDefinition("dashboard.grafana.app", "dashboards")
	driverName := mg.Dialect.DriverName()

	folderCountValidator := NewCountValidator(
		client,
		foldersDef.groupResource,
		"dashboard",
		"org_id = ? and is_folder = true",
		driverName,
	)

	dashboardCountValidator := NewCountValidator(
		client,
		dashboardsDef.groupResource,
		"dashboard",
		"org_id = ? and is_folder = false",
		driverName,
	)

	folderTreeValidator := NewFolderTreeValidator(client, foldersDef.groupResource, driverName)

	dashboardsAndFolders := NewResourceMigration(
		migrator,
		[]schema.GroupResource{foldersDef.groupResource, dashboardsDef.groupResource},
		"folders-dashboards",
		[]Validator{folderCountValidator, dashboardCountValidator, folderTreeValidator},
		opts...,
	)
	mg.AddMigration("folders and dashboards migration", dashboardsAndFolders)
}

func registerPlaylistMigration(mg *sqlstoremigrator.Migrator,
	migrator UnifiedMigrator,
	client resource.ResourceClient,
	opts ...ResourceMigrationOption,
) {
	playlistsDef := getResourceDefinition("playlist.grafana.app", "playlists")
	driverName := mg.Dialect.DriverName()

	playlistCountValidator := NewCountValidator(
		client,
		playlistsDef.groupResource,
		"playlist",
		"org_id = ?",
		driverName,
	)

	playlistsMigration := NewResourceMigration(
		migrator,
		[]schema.GroupResource{playlistsDef.groupResource},
		"playlists",
		[]Validator{playlistCountValidator},
		opts...,
	)
	mg.AddMigration("playlists migration", playlistsMigration)
}

// TODO: remove this before Grafana 13 GA: https://github.com/grafana/search-and-storage-team/issues/613
func shouldAutoMigrate(ctx context.Context, migration migrationDefinition, cfg *setting.Cfg, sqlStore db.DB) bool {
	autoMigrate := false

	for _, res := range migration.resources {
		config := cfg.UnifiedStorageConfig(res)

		if config.DualWriterMode == 5 {
			return false
		}

		if !setting.AutoMigratedUnifiedResources[res] {
			continue
		}

		if checkIfAlreadyMigrated(ctx, migration, sqlStore) {
			for _, res := range migration.resources {
				cfg.EnableMode5(res)
			}
			logger.Info("Auto-migration already completed, enabling mode 5 for resources", "migration", migration.name)
			return true
		}

		autoMigrate = true
		threshold := int64(setting.DefaultAutoMigrationThreshold)
		if config.AutoMigrationThreshold > 0 {
			threshold = int64(config.AutoMigrationThreshold)
		}

		count, err := countResource(ctx, sqlStore, res)
		if err != nil {
			logger.Warn("Failed to count resource for auto migration check", "resource", res, "error", err)
			return false
		}

		logger.Info("Resource count for auto migration check", "resource", res, "count", count, "threshold", threshold)

		if count > threshold {
			return false
		}
	}

	if !autoMigrate {
		return false
	}

	logger.Info("Auto-migration enabled for migration", "migration", migration.name)
	return true
}

func checkIfAlreadyMigrated(ctx context.Context, migration migrationDefinition, sqlStore db.DB) bool {
	if migration.migrationID == "" {
		return false
	}

	exists, err := migrationExists(ctx, sqlStore, migration.migrationID)
	if err != nil {
		logger.Warn("Failed to check if migration exists", "migration", migration.name, "error", err)
		return false
	}

	return exists
}

func isMigrationEnabled(migration migrationDefinition, cfg *setting.Cfg) (bool, error) {
	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 false, fmt.Errorf("cannot migrate resources separately: %v migration must be either all enabled or all disabled", migration.resources)
		}
	}

	return allEnabled, nil
}

// TODO: remove this before Grafana 13 GA: https://github.com/grafana/search-and-storage-team/issues/613
func countResource(ctx context.Context, sqlStore db.DB, resourceName string) (int64, error) {
	var count int64
	err := sqlStore.WithDbSession(ctx, func(sess *db.Session) error {
		switch resourceName {
		case setting.DashboardResource:
			var err error
			count, err = sess.Table("dashboard").Where("is_folder = ?", false).Count()
			return err
		case setting.FolderResource:
			var err error
			count, err = sess.Table("dashboard").Where("is_folder = ?", true).Count()
			return err
		default:
			return fmt.Errorf("unknown resource: %s", resourceName)
		}
	})
	return count, err
}

const migrationLogTableName = "unifiedstorage_migration_log"

func migrationExists(ctx context.Context, sqlStore db.DB, migrationID string) (bool, error) {
	var count int64
	err := sqlStore.WithDbSession(ctx, func(sess *db.Session) error {
		var err error
		count, err = sess.Table(migrationLogTableName).Where("migration_id = ?", migrationID).Count()
		return err
	})
	if err != nil {
		return false, fmt.Errorf("failed to check migration existence: %w", err)
	}
	return count > 0, nil
}

func getResourceDefinition(group, resource string) *resourceDefinition {
	for i := range resourceRegistry {
		r := &resourceRegistry[i]
		if r.groupResource.Group == group && r.groupResource.Resource == resource {
			return r
		}
	}
	return nil
}

func buildResourceKey(group, resource, namespace string) *resourcepb.ResourceKey {
	def := getResourceDefinition(group, resource)
	if def == nil {
		return nil
	}
	return &resourcepb.ResourceKey{
		Namespace: namespace,
		Group:     def.groupResource.Group,
		Resource:  def.groupResource.Resource,
	}
}

func getMigratorFunc(accessor legacy.MigrationDashboardAccessor, group, resource string) migratorFunc {
	def := getResourceDefinition(group, resource)
	if def == nil {
		return nil
	}

	switch def.migratorFunc {
	case "MigrateFolders":
		return accessor.MigrateFolders
	case "MigrateLibraryPanels":
		return accessor.MigrateLibraryPanels
	case "MigrateDashboards":
		return accessor.MigrateDashboards
	case "MigratePlaylists":
		return accessor.MigratePlaylists
	default:
		return nil
	}
}

func validateRegisteredResources() error {
	registeredMap := make(map[string]bool)
	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 {
		if !registeredMap[expected] {
			missing = append(missing, expected)
		}
	}

	if len(missing) > 0 {
		return fmt.Errorf("resources declared in setting.MigratedUnifiedResources are not registered for migration: %v", missing)
	}

	return nil
}