feat: inject unified data migrations in dual writer (#114138)

* feat: draft changes for on-prem unified migration * feat: further draft changes for on-prem unified migration * fix: remove some tbis * refactor: rename * fix: another approach * fix: background service related issues * fix: address comments * fix: make gen-go * fix: background service related issues * feat: refactor dual writer and legacy migrator * fix: minor issues * feat: working version in oss * fix: wire * fix: revert test data override * fix: enterprise related issues * chore: add todo * fix: revert dual writer method * fix: lint * chore: logger format * fix: reduce log level * fix: log change * fix: disable * fix: address comments * fix: return error on dual writer service * fix: merge conflict --------- Co-authored-by: Rafael Paulovic <rafael.paulovic@grafana.com>
2025-11-20 16:40:20 +01:00
parent cb06bba243
commit 30c04ab3fc
40 changed files with 1779 additions and 1348 deletions
@@ -0,0 +1,118 @@
+package migrations
+
+import (
+	"context"
+	"fmt"
+	"os"
+
+	"github.com/grafana/grafana/pkg/infra/db"
+	"github.com/grafana/grafana/pkg/infra/kvstore"
+	"github.com/grafana/grafana/pkg/infra/log"
+	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/prometheus/client_golang/prometheus"
+	"go.opentelemetry.io/otel"
+	"k8s.io/apimachinery/pkg/runtime/schema"
+)
+
+var tracer = otel.Tracer("github.com/grafana/grafana/pkg/storage/unified/migrations")
+var logger = log.New("storage.unified.migrations")
+
+type UnifiedStorageMigrationServiceImpl struct {
+	migrator UnifiedMigrator
+	cfg      *setting.Cfg
+	sqlStore db.DB
+	kv       kvstore.KVStore
+}
+
+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,
+) contract.UnifiedStorageMigrationService {
+	return &UnifiedStorageMigrationServiceImpl{
+		migrator: migrator,
+		cfg:      cfg,
+		sqlStore: sqlStore,
+		kv:       kv,
+	}
+}
+
+// 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") != "" {
+		return nil
+	}
+
+	// skip migrations if disabled in config
+	if p.cfg.DisableDataMigrations {
+		logger.Info("Data migrations are disabled, skipping")
+		return nil
+	}
+
+	// TODO: Re-enable once migrations are ready
+	// TODO: add guarantee that this only runs once
+	// return RegisterMigrations(p.migrator, p.cfg, p.sqlStore)
+	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,
+) error {
+	ctx, span := tracer.Start(context.Background(), "storage.unified.RegisterMigrations")
+	defer span.End()
+	mg := sqlstoremigrator.NewScopedMigrator(sqlStore.GetEngine(), cfg, "unifiedstorage")
+	mg.AddCreateMigration()
+
+	if err := prometheus.Register(mg); err != nil {
+		logger.Warn("Failed to register migrator metrics", "error", err)
+	}
+
+	// Register resource migrations
+	// To add a new resource type, simply add another migration here with the appropriate resources
+	registerResourceMigrations(mg, migrator)
+
+	// Run all registered migrations (blocking)
+	sec := cfg.Raw.Section("database")
+	if err := mg.RunMigrations(ctx,
+		sec.Key("migration_locking").MustBool(true),
+		sec.Key("locking_attempt_timeout_sec").MustInt()); err != nil {
+		return fmt.Errorf("unified storage data migration failed: %w", err)
+	}
+
+	logger.Info("Unified storage migrations completed successfully")
+	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) {
+	dashboardsAndFolders := NewResourceMigration(
+		migrator,
+		[]schema.GroupResource{
+			{Group: "folder.grafana.app", Resource: "folders"},
+			{Group: "dashboard.grafana.app", Resource: "dashboards"},
+		},
+		"folders-dashboards",
+		NewLegacyTableCountValidator(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"},
+		}),
+	)
+	mg.AddMigration("folders and dashboards migration", dashboardsAndFolders)
+}