Alter PK column correctly even if it is part of a foreign key

Author: kvch

pkg/migrations/dbactions.go

@@ -736,3 +736,75 @@ func (a *rawSQLAction) Execute(ctx context.Context) error {
 	_, err := a.conn.ExecContext(ctx, a.sql)
 	return err
 }
+
+type alterReferencesAction struct {
+	conn   db.DB
+	table  string
+	column string
+}
+
+func NewAlterReferencesAction(conn db.DB, table, column string) *alterReferencesAction {
+	return &alterReferencesAction{
+		conn:   conn,
+		table:  table,
+		column: column,
+	}
+}
+
+func (a *alterReferencesAction) Execute(ctx context.Context) error {
+	definitions, err := a.constraintDefinitions(ctx)
+	if err != nil {
+		return err
+	}
+	for _, def := range definitions {
+		// Drop the existing constraint
+		_, err := a.conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s DROP CONSTRAINT %s",
+			pq.QuoteIdentifier(def.table),
+			pq.QuoteIdentifier(def.name),
+		))
+		if err != nil {
+			return fmt.Errorf("dropping constraint %s on %s: %w", def.name, def.table, err)
+		}
+
+		// Recreate the constraint with the table and new column
+		newDef := strings.ReplaceAll(def.def, a.column, TemporaryName(a.column))
+		_, err = a.conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s ADD CONSTRAINT %s %s",
+			pq.QuoteIdentifier(def.table),
+			pq.QuoteIdentifier(def.name),
+			newDef,
+		))
+		if err != nil {
+			return fmt.Errorf("altering references for %s.%s: %w", a.table, a.column, err)
+		}
+	}
+	return nil
+}
+
+type constraintDefinition struct {
+	name  string
+	table string
+	def   string
+}
+
+func (a *alterReferencesAction) constraintDefinitions(ctx context.Context) ([]constraintDefinition, error) {
+	rows, err := a.conn.QueryContext(ctx, fmt.Sprintf(`
+SELECT conname, r.conrelid::regclass, pg_catalog.pg_get_constraintdef(r.oid, true) as condef
+FROM pg_catalog.pg_constraint r
+WHERE confrelid = %s::regclass AND r.contype = 'f'`,
+		pq.QuoteLiteral(a.table),
+	))
+	if err != nil {
+		return nil, fmt.Errorf("querying referencing constraints for %s.%s: %w", a.table, a.column, err)
+	}
+	defer rows.Close()
+
+	defs := make([]constraintDefinition, 0)
+	for rows.Next() {
+		var def constraintDefinition
+		if err := rows.Scan(&def.name, &def.table, &def.def); err != nil {
+			return nil, fmt.Errorf("scanning referencing constraints for %s.%s: %w", a.table, a.column, err)
+		}
+		defs = append(defs, def)
+	}
+	return defs, rows.Err()
+}

pkg/migrations/op_alter_column.go

@@ -118,6 +118,7 @@ func (o *OpAlterColumn) Complete(l Logger, conn db.DB, s *schema.Schema) ([]DBAc
 
 	return append(dbActions, []DBAction{
 		NewAlterSequenceOwnerAction(conn, o.Table, o.Column, TemporaryName(o.Column)),
+		NewAlterReferencesAction(conn, o.Table, o.Column),
 		NewDropColumnAction(conn, o.Table, o.Column),
 		NewDropFunctionAction(conn,
 			backfill.TriggerFunctionName(o.Table, o.Column),

pkg/migrations/op_alter_column_test.go

@@ -674,6 +674,127 @@ func TestAlterPrimaryKeyColumns(t *testing.T) {
 				})
 			},
 		},
+		{
+			name: "alter a single primary key column when the column is used in a foreign key constraint",
+			migrations: []migrations.Migration{
+				{
+					Name: "01_create_table",
+					Operations: migrations.Operations{
+						&migrations.OpCreateTable{
+							Name: "events",
+							Columns: []migrations.Column{
+								{
+									Name: "id",
+									Type: "serial",
+									Pk:   true,
+								},
+								{
+									Name:     "name",
+									Type:     "varchar(255)",
+									Nullable: true,
+								},
+							},
+						},
+						&migrations.OpCreateTable{
+							Name: "people",
+							Columns: []migrations.Column{
+								{
+									Name: "id",
+									Type: "int",
+									Pk:   true,
+								},
+								{
+									Name:     "name",
+									Type:     "varchar(255)",
+									Nullable: true,
+								},
+								{
+									Name:     "manages",
+									Type:     "serial",
+									Nullable: false,
+									References: &migrations.ForeignKeyReference{
+										Table:  "events",
+										Column: "id",
+										Name:   "person_manages_event_fk",
+									},
+								},
+							},
+						},
+					},
+				},
+				{
+					Name: "02_alter_column",
+					Operations: migrations.Operations{
+						&migrations.OpAlterColumn{
+							Table:  "people",
+							Column: "id",
+							Type:   ptr("bigint"),
+							Up:     "CAST(id AS bigint)",
+							Down:   "SELECT CASE WHEN id < 2147483647 THEN CAST(id AS int) ELSE 0 END",
+						},
+					},
+				},
+			},
+			afterStart: func(t *testing.T, db *sql.DB, schema string) {
+				PrimaryKeyConstraintMustExist(t, db, schema, "people", "people_pkey")
+				ValidatedForeignKeyMustExist(t, db, schema, "people", "person_manages_event_fk")
+
+				bigint := "31474836471" // A value larger than int can hold
+				integer := "100"
+
+				// Inserting a row with integer id into the old schema should succeed
+				MustInsert(t, db, schema, "01_create_table", "events", map[string]string{
+					"name": "pgroll v1 release party",
+				})
+				MustInsert(t, db, schema, "01_create_table", "people", map[string]string{
+					"id":      integer,
+					"name":    "alice",
+					"manages": "1",
+				})
+				// Inserting a row with integer bigint id into the old schema should fail
+				MustNotInsert(t, db, schema, "01_create_table", "people", map[string]string{
+					"id":      bigint,
+					"name":    "alice",
+					"manages": "1",
+				}, testutils.NumericValueOutOfRangeErrorCode)
+
+				// Inserting a row with bigint id into the new schema should succeed
+				MustInsert(t, db, schema, "02_alter_column", "people", map[string]string{
+					"id":      bigint,
+					"name":    "alice",
+					"manages": "1",
+				})
+				// Inserting a row into the `people` table with a `manages` field that
+				// violates the FK constraint fails
+				MustNotInsert(t, db, schema, "02_alter_column", "people", map[string]string{
+					"id":      "10",
+					"name":    "alice",
+					"manages": "2",
+				}, testutils.FKViolationErrorCode)
+			},
+			afterRollback: func(t *testing.T, db *sql.DB, schema string) {
+				PrimaryKeyConstraintMustExist(t, db, schema, "people", "people_pkey")
+				ValidatedForeignKeyMustExist(t, db, schema, "people", "person_manages_event_fk")
+			},
+			afterComplete: func(t *testing.T, db *sql.DB, schema string) {
+				PrimaryKeyConstraintMustExist(t, db, schema, "people", "people_pkey")
+				ValidatedForeignKeyMustExist(t, db, schema, "people", "person_manages_event_fk")
+
+				// Inserting a row with integer bigint into the new schema should succeed
+				MustInsert(t, db, schema, "02_alter_column", "people", map[string]string{
+					"id":      "31474836472",
+					"name":    "bob",
+					"manages": "1",
+				})
+				// Inserting a row into the `people` table with a `manages` field that
+				// violates the FK constraint fails
+				MustNotInsert(t, db, schema, "02_alter_column", "people", map[string]string{
+					"id":      "31474836474",
+					"name":    "alice",
+					"manages": "2",
+				}, testutils.FKViolationErrorCode)
+			},
+		},
 	})
 }