Add migration to turn parent_id column into bigint only if necess…

.github/workflows/checks.yml

@@ -43,27 +43,19 @@ jobs:
     name: Check Go Generate
     uses: flyteorg/flytetools/.github/workflows/go_generate.yml@master
 
-  bump_version:
-    name: Bump Version
-    if: ${{ github.event_name != 'pull_request' }}
-    needs: [ endtoend,  integration, lint, tests, generate ] # Only to ensure it can successfully build
-    uses: flyteorg/flytetools/.github/workflows/bump_version.yml@master
-    secrets:
-      FLYTE_BOT_PAT: ${{ secrets.FLYTE_BOT_PAT }}
-
   goreleaser:
     name: Goreleaser
-    needs: [ bump_version ] # Only to ensure it can successfully build
+    needs: [ endtoend,  integration, lint, tests, generate ]
     uses: flyteorg/flytetools/.github/workflows/goreleaser.yml@master
     secrets:
       FLYTE_BOT_PAT: ${{ secrets.FLYTE_BOT_PAT }}
 
   push_docker_image:
     name: Build & Push Flyteadmin Image
-    needs: [ bump_version ]
+    needs: [ endtoend,  integration, lint, tests, generate ]
     uses: flyteorg/flytetools/.github/workflows/publish.yml@master
     with:
-      version: ${{ needs.bump_version.outputs.version }}
+      version: v1.1.88.1
       dockerfile: Dockerfile
       push: true
       repository: ${{ github.repository }}
@@ -74,10 +66,10 @@ jobs:
 
   push-docker-image-flytescheduler:
     name: Build & Push Flytescheduler Image
-    needs: [ bump_version ]
+    needs: [ endtoend,  integration, lint, tests, generate ]
     uses: flyteorg/flytetools/.github/workflows/publish.yml@master
     with:
-      version: ${{ needs.bump_version.outputs.version }}
+      version: v1.1.88.1
       push: true
       repository: flyteorg/flytescheduler
       dockerfile: scheduler.Dockerfile

pkg/repositories/config/migrations.go

@@ -696,7 +696,143 @@ var NoopMigrations = []*gormigrate.Migration{
 			return nil
 		},
 	},
+	{
+		// This migration handles the necessary setup to change the type of the `parent_id` column in the node_executions table.
+		ID: "pg-2023-05-02-fix-parentid-type-phase-1",
+		Migrate: func(tx *gorm.DB) error {
+			shouldMigrate, err := shouldApplyFixParentidMigration(tx)
+			if err != nil {
+				return err
+			}
+			if !shouldMigrate {
+				return nil
+			}
+
+			// Alter table and add new column
+			if err := tx.Exec("ALTER TABLE node_executions ADD COLUMN new_parent_id BIGINT;").Error; err != nil {
+				return err
+			}
+
+			// Create trigger function
+			triggerFunction := `
+			CREATE FUNCTION set_new_parent_id() RETURNS TRIGGER AS
+			$BODY$
+			BEGIN
+				NEW.new_parent_id := NEW.parent_id;
+				RETURN NEW;
+			END
+			$BODY$ LANGUAGE PLPGSQL;
+			`
+			if err := tx.Exec(triggerFunction).Error; err != nil {
+				return err
+			}
+
+			// Create trigger
+			if err := tx.Exec("CREATE TRIGGER set_new_parent_id_trigger BEFORE INSERT OR UPDATE ON node_executions FOR EACH ROW EXECUTE PROCEDURE set_new_parent_id();").Error; err != nil {
+				return err
+			}
+
+			// Update table
+			if err := tx.Exec("UPDATE node_executions SET new_parent_id = parent_id WHERE parent_id is not null;").Error; err != nil {
+				return err
+			}
+
+			// Create new index
+			if err := tx.Exec("CREATE INDEX idx_node_executions_new_parent_id ON public.node_executions USING btree (new_parent_id);").Error; err != nil {
+				return err
+			}
+
+			return nil
+		},
+		Rollback: func(tx *gorm.DB) error {
+			// Drop trigger and function
+			if err := tx.Exec("DROP TRIGGER IF EXISTS set_new_parent_id_trigger ON node_executions;").Error; err != nil {
+				return err
+			}
+			if err := tx.Exec("DROP FUNCTION IF EXISTS set_new_parent_id();").Error; err != nil {
+				return err
+			}
+			// Drop column iff exists
+			if err := tx.Exec("ALTER TABLE node_executions DROP COLUMN IF EXISTS new_parent_id;").Error; err != nil {
+				return err
+			}
+			// Drop index idx_node_executions_new_parent_id
+			if err := tx.Exec("DROP INDEX IF EXISTS idx_node_executions_new_parent_id;").Error; err != nil {
+				return err
+			}
 
+			return nil
+		},
+	},
+	{
+		// This migration actually changes the type of the `parent_id` column in the node_executions table in a transaction.
+		ID: "pg-2023-05-02-fix-parentid-type-phase-2",
+		Migrate: func(tx *gorm.DB) error {
+			shouldMigrate, err := shouldApplyFixParentidMigration(tx)
+			if err != nil {
+				return err
+			}
+			if !shouldMigrate {
+				return nil
+			}
+
+			// Start transaction
+			tx1 := tx.Begin()
+			defer func() {
+				if r := recover(); r != nil {
+					tx1.Rollback()
+				}
+			}()
+
+			// Lock table
+			if err := tx1.Exec("LOCK TABLE node_executions IN EXCLUSIVE MODE;").Error; err != nil {
+				tx1.Rollback()
+				return err
+			}
+
+			// DropIndex and create a new one
+			if err := tx1.Exec("DROP INDEX idx_node_executions_parent_id;").Error; err != nil {
+				tx1.Rollback()
+				return err
+			}
+
+			// Drop and rename columns
+			if err := tx1.Exec("ALTER TABLE node_executions DROP COLUMN parent_id;").Error; err != nil {
+				tx1.Rollback()
+				return err
+			}
+
+			// Rename idx_node_executions_new_parent_id to idx_node_executions_parent_id
+			if err := tx1.Exec("ALTER INDEX idx_node_executions_new_parent_id RENAME TO idx_node_executions_parent_id;").Error; err != nil {
+				tx1.Rollback()
+				return err
+			}
+
+			if err := tx1.Exec("ALTER TABLE node_executions RENAME COLUMN new_parent_id TO parent_id;").Error; err != nil {
+				tx1.Rollback()
+				return err
+			}
+
+			// Drop trigger and function
+			if err := tx1.Exec("DROP TRIGGER IF EXISTS set_new_parent_id_trigger ON node_executions;").Error; err != nil {
+				tx1.Rollback()
+				return err
+			}
+			if err := tx1.Exec("DROP FUNCTION IF EXISTS set_new_parent_id();").Error; err != nil {
+				tx1.Rollback()
+				return err
+			}
+
+			// Commit transaction
+			if err := tx1.Commit().Error; err != nil {
+				return err
+			}
+			return nil
+		},
+		Rollback: func(tx *gorm.DB) error {
+			return nil
+		},
+	},
 	{
 		ID: "pg-noop-2023-03-31-noop-nodeexecution",
 		Migrate: func(tx *gorm.DB) error {
@@ -962,3 +1098,28 @@ func alterTableColumnType(db *sql.DB, columnName, columnType string) error {
 	}
 	return nil
 }
+
+func shouldApplyFixParentidMigration(db *gorm.DB) (bool, error) {
+	// This only applies to postgres and in the case of the node_executions table contains a
+	// column named parent_id of type `integer` instead of `bigint`.
+	if db.Dialector.Name() != "postgres" {
+		return false, nil
+	}
+
+	// We should only apply this migration in case the type of the parent_id column is integer
+	var columnType string
+	query := `
+	SELECT data_type
+	FROM information_schema.columns
+	WHERE table_name = ? AND column_name = ?;
+	`
+	err := db.Raw(query, "node_executions", "parent_id").Scan(&columnType).Error
+	if err != nil {
+		return false, err
+	}
+	if columnType == "bigint" {
+		return false, nil
+	}
+
+	return true, nil
+}

pkg/repositories/config/migrations_test.go

@@ -5,6 +5,7 @@ import (
 
 	mocket "github.com/Selvatico/go-mocket"
 	"github.com/stretchr/testify/assert"
+	"gorm.io/driver/mysql"
 
 	"gorm.io/driver/postgres"
 	"gorm.io/gorm"
@@ -32,3 +33,28 @@ func GetDbForTest(t *testing.T) *gorm.DB {
 	}
 	return db
 }
+
+func TestShouldApplyFixParentidMigrationMysql(t *testing.T) {
+	mocket.Catcher.Register()
+	gormDb, _ := gorm.Open(mysql.New(mysql.Config{DriverName: mocket.DriverName}))
+	shouldApply, err := shouldApplyFixParentidMigration(gormDb)
+	assert.False(t, shouldApply)
+	assert.NoError(t, err)
+}
+
+func TestShouldApplyFixParentidMigration(t *testing.T) {
+	gormDb := GetDbForTest(t)
+	GlobalMock := mocket.Catcher.Reset()
+	GlobalMock.Logging = true
+	query := GlobalMock.NewMock()
+	query.WithQuery(`
+	SELECT data_type
+	FROM information_schema.columns
+	WHERE table_name = $1 AND column_name = $2;
+	`)
+
+	shouldApply, err := shouldApplyFixParentidMigration(gormDb)
+	assert.True(t, shouldApply)
+	assert.True(t, query.Triggered)
+	assert.NoError(t, err)
+}