[v2] DB migration with version (#7103)

Signed-off-by: machichima <nary12321@gmail.com>

Author: machichima

cache_service/manager/manager_test.go

@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/glebarez/sqlite"
+	"github.com/go-gormigrate/gormigrate/v2"
 	"github.com/stretchr/testify/require"
 	durationpb "google.golang.org/protobuf/types/known/durationpb"
 	"gorm.io/gorm"
@@ -21,7 +22,8 @@ func newTestManager(t *testing.T) *Manager {
 
 	db, err := gorm.Open(sqlite.Open("file::memory:?cache=shared"), &gorm.Config{})
 	require.NoError(t, err)
-	require.NoError(t, migrations.RunMigrations(db))
+	m := gormigrate.New(db, gormigrate.DefaultOptions, migrations.CacheServiceMigrations)
+	require.NoError(t, m.Migrate())
 
 	cfg := &cacheconfig.Config{
 		HeartbeatGracePeriodMultiplier: 3,

cache_service/migrations/migrations.go

@@ -1,25 +1,38 @@
 package migrations
 
 import (
-	"context"
 	"fmt"
 
+	"github.com/go-gormigrate/gormigrate/v2"
 	"gorm.io/gorm"
 
 	"github.com/flyteorg/flyte/v2/cache_service/repository/models"
-	"github.com/flyteorg/flyte/v2/flytestdlib/logger"
 )
 
 var allModels = []interface{}{
 	&models.CachedOutput{},
 	&models.Reservation{},
 }
 
-func RunMigrations(db *gorm.DB) error {
+const MigrationIDInitSchema = "20260327_cache_service_init_schema"
+
+var CacheServiceMigrations = []*gormigrate.Migration{
+	{
+		ID: MigrationIDInitSchema,
+		Migrate: func(tx *gorm.DB) error {
+			return migrateInitSchema(tx)
+		},
+		Rollback: func(tx *gorm.DB) error {
+			// Intentionally no-op for now; this migration can contain destructive steps.
+			return nil
+		},
+	},
+}
+
+// migrateInitSchema initializes the cache service database schema.
+func migrateInitSchema(db *gorm.DB) error {
 	if err := db.AutoMigrate(allModels...); err != nil {
-		return fmt.Errorf("failed to run cache service migrations: %w", err)
+		return fmt.Errorf("failed to initialize cache service schema: %w", err)
 	}
-
-	logger.Infof(context.Background(), "Cache service migrations completed successfully")
 	return nil
 }

cache_service/setup.go

@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"net/http"
 
+	"github.com/go-gormigrate/gormigrate/v2"
+
 	"github.com/flyteorg/flyte/v2/app"
 	"github.com/flyteorg/flyte/v2/cache_service/config"
 	"github.com/flyteorg/flyte/v2/cache_service/migrations"
@@ -17,8 +19,9 @@ import (
 // Requires sc.DB and sc.DataStore to be set by the standalone binary.
 func Setup(ctx context.Context, sc *app.SetupContext) error {
 	cfg := config.GetConfig()
-	if err := migrations.RunMigrations(sc.DB); err != nil {
-		return err
+	m := gormigrate.New(sc.DB, gormigrate.DefaultOptions, migrations.CacheServiceMigrations)
+	if err := m.Migrate(); err != nil {
+		return fmt.Errorf("cache_service: failed to run migrations: %w", err)
 	}
 
 	path, handler := v2connect.NewCacheServiceHandler(service.NewCacheService(cfg, sc.DB))

manager/README.md

@@ -266,11 +266,10 @@ manager:
 
 ### Database Issues
 
-**Error:** "failed to run migrations"
+**Error:** relation/table not found (runs schema missing)
 
 ```bash
-# Remove and recreate database
-rm flyte.db
+# Restart manager so startup migrations run
 ./bin/flyte-manager --config config.yaml
 ```
 

runs/README.md

@@ -409,8 +409,13 @@ No code changes needed, just update `config.yaml`!
 
 ### Database migration
 
-We are using gorm for auto migration, please ensure addding `gorm` tags in `/runs/repository/models/` when you
-add any more columns/models.
+Runs schema changes are managed through versioned gormigrate migrations.
+
+Migrations run automatically during service startup.
+
+```bash
+make run
+```
 
 ### Adding new features
 

runs/migrations/migrations.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 
+	"github.com/go-gormigrate/gormigrate/v2"
 	"gorm.io/gorm"
 
 	"github.com/flyteorg/flyte/v2/flytestdlib/logger"
@@ -19,8 +20,23 @@ var AllModels = []interface{}{
 	&models.TaskSpec{},
 }
 
-// RunMigrations runs all database migrations for the runs service
-func RunMigrations(db *gorm.DB) error {
+const MigrationIDInitSchema = "20260327_runs_init_schema"
+
+var RunsMigrations = []*gormigrate.Migration{
+	{
+		ID: MigrationIDInitSchema,
+		Migrate: func(tx *gorm.DB) error {
+			return migrateInitSchema(tx)
+		},
+		Rollback: func(tx *gorm.DB) error {
+			// Intentionally no-op for now; this migration can contain destructive steps.
+			return nil
+		},
+	},
+}
+
+// migrateInitSchema initializes the runs service database schema.
+func migrateInitSchema(db *gorm.DB) error {
 	ctx := context.Background()
 
 	// Drop stale tables from previous schema versions that are no longer used.

runs/setup.go

@@ -7,6 +7,8 @@ import (
 	"net/http"
 	"time"
 
+	"github.com/go-gormigrate/gormigrate/v2"
+
 	"github.com/flyteorg/flyte/v2/app"
 	"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/actions/actionsconnect"
 	flyteappconnect "github.com/flyteorg/flyte/v2/gen/go/flyteidl2/app/appconnect"
@@ -32,8 +34,8 @@ import (
 // RunLogsService is also mounted to enable pod log streaming.
 func Setup(ctx context.Context, sc *app.SetupContext) error {
 	cfg := config.GetConfig()
-
-	if err := migrations.RunMigrations(sc.DB); err != nil {
+	m := gormigrate.New(sc.DB, gormigrate.DefaultOptions, migrations.RunsMigrations)
+	if err := m.Migrate(); err != nil {
 		return fmt.Errorf("runs: failed to run migrations: %w", err)
 	}
 

runs/test/api/setup_test.go

@@ -78,7 +78,7 @@ func TestMain(m *testing.M) {
 	log.Println("Database initialized")
 
 	// Run migrations
-	if err := migrations.RunMigrations(testDB); err != nil {
+	if err := database.Migrate(ctx, dbConfig, migrations.RunsMigrations); err != nil {
 		log.Printf("Failed to run migrations: %v", err)
 		exitCode = 1
 		return