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

package migrations

import (
	"context"
	"fmt"
	"time"

	"github.com/grafana/authlib/types"
	"github.com/grafana/grafana/pkg/apimachinery/identity"
	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/registry/apis/dashboard/legacy"
	"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"
	"github.com/grafana/grafana/pkg/util/xorm"
	"github.com/grafana/grafana/pkg/util/xorm/core"
	"k8s.io/apimachinery/pkg/runtime/schema"
)

// ValidationFunc is a function that validates migration results.

type Validator interface {
	Name() string
	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.
type ResourceMigration struct {
	migrator.MigrationBase
	migrator    UnifiedMigrator
	resources   []schema.GroupResource
	migrationID string
	validators  []Validator // Optional: custom validation logic for this migration
	log         log.Logger
	cfg         *setting.Cfg
	autoMigrate bool // If true, auto-migrate resource if count is below threshold
	hadErrors   bool // Tracks if errors occurred during migration (used with ignoreErrors)
}

// ResourceMigrationOption is a functional option for configuring ResourceMigration.
type ResourceMigrationOption func(*ResourceMigration)

// WithAutoMigrate configures the migration to auto-migrate resource if count is below threshold.
func WithAutoMigrate(cfg *setting.Cfg) ResourceMigrationOption {
	return func(m *ResourceMigration) {
		m.cfg = cfg
		m.autoMigrate = true
	}
}

// NewResourceMigration creates a new migration for the specified resources.
func NewResourceMigration(
	migrator UnifiedMigrator,
	resources []schema.GroupResource,
	migrationID string,
	validators []Validator,
	opts ...ResourceMigrationOption,
) *ResourceMigration {
	m := &ResourceMigration{
		migrator:    migrator,
		resources:   resources,
		migrationID: migrationID,
		validators:  validators,
		log:         log.New("storage.unified.resource_migration." + migrationID),
	}
	for _, opt := range opts {
		opt(m)
	}
	return m
}

func (m *ResourceMigration) SkipMigrationLog() bool {
	// Skip populating the log table if auto-migrate is enabled and errors occurred
	return m.autoMigrate && m.hadErrors
}

var _ migrator.CodeMigration = (*ResourceMigration)(nil)

// SQL implements migrator.Migration interface. Returns a description string.
func (m *ResourceMigration) SQL(_ migrator.Dialect) string {
	return fmt.Sprintf("unified storage data migration: %s", m.migrationID)
}

// Exec implements migrator.CodeMigration interface. Executes the migration across all organizations.
func (m *ResourceMigration) Exec(sess *xorm.Session, mg *migrator.Migrator) (err error) {
	// Track any errors that occur during migration
	defer func() {
		if err != nil {
			if m.autoMigrate {
				m.log.Warn(
					`[WARN] Resource migration failed and is currently skipped. 
This migration will be enforced in the next major Grafana release, where failures will block startup or resource loading.

This warning is intended to help you detect and report issues early. 
Please investigate the failure and report it to the Grafana team so it can be addressed before the next major release.`,
					"error", err)
			}
			m.hadErrors = true
		}
	}()

	ctx := context.Background()

	orgs, err := m.getAllOrgs(sess)
	if err != nil {
		m.log.Error("failed to get organizations", "error", err)
		return fmt.Errorf("failed to get organizations: %w", err)
	}

	if len(orgs) == 0 {
		m.log.Info("No organizations found to migrate, skipping migration")
		return nil
	}

	m.log.Info("Starting migration for all organizations", "org_count", len(orgs), "resources", m.resources)

	if mg.Dialect.DriverName() == migrator.SQLite {
		// reuse transaction in SQLite to avoid "database is locked" errors
		var tx *core.Tx
		tx, err = sess.Tx()
		if err != nil {
			m.log.Error("Failed to get transaction from session", "error", err)
			return fmt.Errorf("failed to get transaction: %w", err)
		}
		ctx = resource.ContextWithTransaction(ctx, tx.Tx)
		m.log.Info("Stored migrator transaction in context for bulk operations (SQLite compatibility)")
	}

	for _, org := range orgs {
		if err = m.migrateOrg(ctx, sess, org); err != nil {
			return err
		}
	}

	// Auto-enable mode 5 for resources after successful migration
	// TODO: remove this before Grafana 13 GA: https://github.com/grafana/search-and-storage-team/issues/613
	if m.autoMigrate {
		for _, gr := range m.resources {
			m.log.Info("Auto-enabling mode 5 for resource", "resource", gr.Resource+"."+gr.Group)
			m.cfg.EnableMode5(gr.Resource + "." + gr.Group)
		}
	}

	m.log.Info("Migration completed successfully for all organizations", "org_count", len(orgs))

	return nil
}

// migrateOrg handles migration for a single organization
func (m *ResourceMigration) migrateOrg(ctx context.Context, sess *xorm.Session, org orgInfo) error {
	namespace := types.OrgNamespaceFormatter(org.ID)
	m.log.Info("Migrating organization", "org_id", org.ID, "org_name", org.Name, "namespace", namespace)

	// Create a service identity context for this namespace to authenticate with unified storage
	migrationCtx, _ := identity.WithServiceIdentityForSingleNamespace(ctx, namespace)

	startTime := time.Now()

	opts := legacy.MigrateOptions{
		Namespace:   namespace,
		Resources:   m.resources,
		WithHistory: true, // Migrate with full history
		Progress: func(count int, msg string) {
			m.log.Info("Migration progress", "org_id", org.ID, "count", count, "message", msg)
		},
	}

	// Execute the migration via legacy migrator
	response, err := m.migrator.Migrate(migrationCtx, opts)
	if err != nil {
		m.log.Error("Migration failed", "org_id", org.ID, "error", err, "duration", time.Since(startTime))
		return fmt.Errorf("migration failed for org %d (%s): %w", org.ID, org.Name, err)
	}
	if response.Error != nil {
		m.log.Error("Migration reported error", "org_id", org.ID, "error", response.Error.String(), "duration", time.Since(startTime))
		return fmt.Errorf("migration failed for org %d (%s): %w", org.ID, org.Name, fmt.Errorf("migration error: %s", response.Error.Message))
	}

	// Validate the migration results
	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)
	}

	m.log.Info("Migration completed for organization",
		"org_id", org.ID,
		"duration", time.Since(startTime),
		"processed", response.Processed,
		"summaries", len(response.Summary),
		"rejected", len(response.Rejected))

	return nil
}

// validateMigration runs all validators in sequence
func (m *ResourceMigration) validateMigration(ctx context.Context, sess *xorm.Session, response *resourcepb.BulkResponse) error {
	if len(m.validators) == 0 {
		m.log.Debug("No validators provided, skipping validation")
		return nil
	}

	for _, validator := range m.validators {
		m.log.Debug("Running validator", "name", validator.Name(), "total", len(m.validators))
		if err := validator.Validate(ctx, sess, response, m.log); err != nil {
			return fmt.Errorf("validator %s failed: %w", validator.Name(), err)
		}
	}

	m.log.Debug("All validators passed", "count", len(m.validators))
	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)
	if err != nil {
		return 0, fmt.Errorf("failed to parse namespace: %w", err)
	}
	return info.OrgID, nil
}

// orgInfo represents basic organization information
type orgInfo struct {
	ID   int64  `xorm:"id"`
	Name string `xorm:"name"`
}

// getAllOrgs retrieves all organizations from the database
func (m *ResourceMigration) getAllOrgs(sess *xorm.Session) ([]orgInfo, error) {
	var orgs []orgInfo
	err := sess.Table("org").Cols("id", "name").Find(&orgs)
	if err != nil {
		return nil, err
	}
	return orgs, nil
}