Support table and column rename operations preceding drop_multicolumn_constraint operations (#684)

Ensure that `drop_multicolumn_constraint` operations can be preceded by
rename table and rename column operations as in the following example:

```json
{
  "name": "24_drop_constraint",
  "operations": [
    {
      "rename_table": {
        "from": "items",
        "to": "products"
      }
    },
    {
      "rename_column": {
        "table": "products",
        "from": "name",
        "to": "item_name"
      }
    },
    {
      "drop_multicolumn_constraint": {
        "table": "products",
        "name": "name_length",
        "up": {
          "item_name": "item_name"
        },
        "down": {
          "item_name": "SELECT CASE WHEN length(item_name) <= 3 THEN LPAD(item_name, 4, '-') ELSE item_name END"
        }
      }
    }
  ]
}
```

Part of #239.

Author: andrew-farries

pkg/migrations/op_drop_multicolumn_constraint.go

@@ -30,8 +30,13 @@ func (o *OpDropMultiColumnConstraint) Start(ctx context.Context, conn db.DB, lat
 		}
 	}
 
-	// Duplicate each of the columns covered by the constraint to be dropped
+	// Duplicate each of the columns covered by the constraint to be dropped.
+	// Each column is duplicated assuming its final name after the migration is
+	// completed.
 	d := NewColumnDuplicator(conn, table, columns...).WithoutConstraint(o.Name)
+	for _, colName := range constraintColumns {
+		d = d.WithName(table.GetColumn(colName).Name, TemporaryName(colName))
+	}
 	if err := d.Duplicate(ctx); err != nil {
 		return nil, fmt.Errorf("failed to duplicate column: %w", err)
 	}
@@ -45,7 +50,7 @@ func (o *OpDropMultiColumnConstraint) Start(ctx context.Context, conn db.DB, lat
 			Columns:        table.Columns,
 			SchemaName:     s.Name,
 			LatestSchema:   latestSchema,
-			TableName:      o.Table,
+			TableName:      table.Name,
 			PhysicalColumn: TemporaryName(columnName),
 			SQL:            o.upSQL(columnName),
 		})
@@ -55,7 +60,9 @@ func (o *OpDropMultiColumnConstraint) Start(ctx context.Context, conn db.DB, lat
 
 		// Add the new column to the internal schema representation. This is done
 		// here, before creation of the down trigger, so that the trigger can declare
-		// a variable for the new column.
+		// a variable for the new column. Save the old column name for use as the
+		// physical column name in the down trigger first.
+		oldPhysicalColumn := table.GetColumn(columnName).Name
 		table.AddColumn(columnName, &schema.Column{
 			Name: TemporaryName(columnName),
 		})
@@ -67,8 +74,8 @@ func (o *OpDropMultiColumnConstraint) Start(ctx context.Context, conn db.DB, lat
 			Columns:        table.Columns,
 			SchemaName:     s.Name,
 			LatestSchema:   latestSchema,
-			TableName:      o.Table,
-			PhysicalColumn: columnName,
+			TableName:      table.Name,
+			PhysicalColumn: oldPhysicalColumn,
 			SQL:            o.Down[columnName],
 		})
 		if err != nil {
@@ -125,7 +132,7 @@ func (o *OpDropMultiColumnConstraint) Rollback(ctx context.Context, conn db.DB,
 	for _, columnName := range table.GetConstraintColumns(o.Name) {
 		// Drop the new column
 		_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s DROP COLUMN IF EXISTS %s",
-			pq.QuoteIdentifier(o.Table),
+			pq.QuoteIdentifier(table.Name),
 			pq.QuoteIdentifier(TemporaryName(columnName)),
 		))
 		if err != nil {

pkg/migrations/op_drop_multicolumn_constraint_test.go

@@ -410,6 +410,224 @@ func TestDropMultiColumnConstraint(t *testing.T) {
 	})
 }
 
+func TestDropMultiColumnConstraintInMultiOperationMigrations(t *testing.T) {
+	t.Parallel()
+
+	ExecuteTests(t, TestCases{
+		{
+			name: "rename table, drop constraint",
+			migrations: []migrations.Migration{
+				{
+					Name: "01_create_table",
+					Operations: migrations.Operations{
+						&migrations.OpCreateTable{
+							Name: "items",
+							Columns: []migrations.Column{
+								{
+									Name: "id",
+									Type: "int",
+									Pk:   true,
+								},
+								{
+									Name:     "name",
+									Type:     "varchar(255)",
+									Nullable: true,
+									Check: &migrations.CheckConstraint{
+										Name:       "check_name_length",
+										Constraint: "length(name) > 3",
+									},
+								},
+							},
+						},
+					},
+				},
+				{
+					Name: "02_multi_operation",
+					Operations: migrations.Operations{
+						&migrations.OpRenameTable{
+							From: "items",
+							To:   "products",
+						},
+						&migrations.OpDropMultiColumnConstraint{
+							Table: "products",
+							Name:  "check_name_length",
+							Down: map[string]string{
+								"name": "SELECT CASE WHEN length(name) <= 3 THEN LPAD(name, 4, '-') ELSE name END",
+							},
+							Up: map[string]string{
+								"name": "name",
+							},
+						},
+					},
+				},
+			},
+			afterStart: func(t *testing.T, db *sql.DB, schema string) {
+				// Can insert a row into the new schema that violates the constraint
+				MustInsert(t, db, schema, "02_multi_operation", "products", map[string]string{
+					"id":   "1",
+					"name": "a",
+				})
+
+				// Can't insert a row into the old schema that violates the constraint
+				MustNotInsert(t, db, schema, "01_create_table", "items", map[string]string{
+					"id":   "2",
+					"name": "b",
+				}, testutils.CheckViolationErrorCode)
+
+				// Can insert a row into the old schema that meets the constraint
+				MustInsert(t, db, schema, "01_create_table", "items", map[string]string{
+					"id":   "2",
+					"name": "bananas",
+				})
+
+				// The new view has the expected rows
+				rows := MustSelect(t, db, schema, "02_multi_operation", "products")
+				assert.Equal(t, []map[string]any{
+					{"id": 1, "name": "a"},
+					{"id": 2, "name": "bananas"},
+				}, rows)
+
+				// The old view has the expected rows
+				rows = MustSelect(t, db, schema, "01_create_table", "items")
+				assert.Equal(t, []map[string]any{
+					{"id": 1, "name": "---a"}, // rewritten by the down migration
+					{"id": 2, "name": "bananas"},
+				}, rows)
+			},
+			afterRollback: func(t *testing.T, db *sql.DB, schema string) {
+				// The table has been cleaned up
+				TableMustBeCleanedUp(t, db, schema, "items", "name")
+			},
+			afterComplete: func(t *testing.T, db *sql.DB, schema string) {
+				// Can insert a row into the new schema that violates the constraint
+				MustInsert(t, db, schema, "02_multi_operation", "products", map[string]string{
+					"id":   "3",
+					"name": "c",
+				})
+
+				// The new view has the expected rows
+				rows := MustSelect(t, db, schema, "02_multi_operation", "products")
+				assert.Equal(t, []map[string]any{
+					{"id": 1, "name": "---a"},
+					{"id": 2, "name": "bananas"},
+					{"id": 3, "name": "c"},
+				}, rows)
+
+				// The table has been cleaned up
+				TableMustBeCleanedUp(t, db, schema, "products", "name")
+			},
+		},
+		{
+			name: "rename table, rename column, drop constraint",
+			migrations: []migrations.Migration{
+				{
+					Name: "01_create_table",
+					Operations: migrations.Operations{
+						&migrations.OpCreateTable{
+							Name: "items",
+							Columns: []migrations.Column{
+								{
+									Name: "id",
+									Type: "int",
+									Pk:   true,
+								},
+								{
+									Name:     "name",
+									Type:     "varchar(255)",
+									Nullable: true,
+									Check: &migrations.CheckConstraint{
+										Name:       "check_name_length",
+										Constraint: "length(name) > 3",
+									},
+								},
+							},
+						},
+					},
+				},
+				{
+					Name: "02_multi_operation",
+					Operations: migrations.Operations{
+						&migrations.OpRenameTable{
+							From: "items",
+							To:   "products",
+						},
+						&migrations.OpRenameColumn{
+							Table: "products",
+							From:  "name",
+							To:    "item_name",
+						},
+						&migrations.OpDropMultiColumnConstraint{
+							Table: "products",
+							Name:  "check_name_length",
+							Down: map[string]string{
+								"item_name": "SELECT CASE WHEN length(item_name) <= 3 THEN LPAD(item_name, 4, '-') ELSE item_name END",
+							},
+							Up: map[string]string{
+								"item_name": "item_name",
+							},
+						},
+					},
+				},
+			},
+			afterStart: func(t *testing.T, db *sql.DB, schema string) {
+				// Can insert a row into the new schema that violates the constraint
+				MustInsert(t, db, schema, "02_multi_operation", "products", map[string]string{
+					"id":        "1",
+					"item_name": "a",
+				})
+
+				// Can't insert a row into the old schema that violates the constraint
+				MustNotInsert(t, db, schema, "01_create_table", "items", map[string]string{
+					"id":   "2",
+					"name": "b",
+				}, testutils.CheckViolationErrorCode)
+
+				// Can insert a row into the old schema that meets the constraint
+				MustInsert(t, db, schema, "01_create_table", "items", map[string]string{
+					"id":   "2",
+					"name": "bananas",
+				})
+
+				// The new view has the expected rows
+				rows := MustSelect(t, db, schema, "02_multi_operation", "products")
+				assert.Equal(t, []map[string]any{
+					{"id": 1, "item_name": "a"},
+					{"id": 2, "item_name": "bananas"},
+				}, rows)
+
+				// The old view has the expected rows
+				rows = MustSelect(t, db, schema, "01_create_table", "items")
+				assert.Equal(t, []map[string]any{
+					{"id": 1, "name": "---a"}, // rewritten by the down migration
+					{"id": 2, "name": "bananas"},
+				}, rows)
+			},
+			afterRollback: func(t *testing.T, db *sql.DB, schema string) {
+				// The table has been cleaned up
+				TableMustBeCleanedUp(t, db, schema, "items", "name")
+			},
+			afterComplete: func(t *testing.T, db *sql.DB, schema string) {
+				// Can insert a row into the new schema that violates the constraint
+				MustInsert(t, db, schema, "02_multi_operation", "products", map[string]string{
+					"id":        "3",
+					"item_name": "c",
+				})
+
+				// The new view has the expected rows
+				rows := MustSelect(t, db, schema, "02_multi_operation", "products")
+				assert.Equal(t, []map[string]any{
+					{"id": 1, "item_name": "---a"},
+					{"id": 2, "item_name": "bananas"},
+					{"id": 3, "item_name": "c"},
+				}, rows)
+
+				// The table has been cleaned up
+				TableMustBeCleanedUp(t, db, schema, "products", "name")
+			},
+		},
+	})
+}
+
 func TestDropMultiColumnConstraintValidation(t *testing.T) {
 	t.Parallel()
 

pkg/migrations/op_rename_column.go

@@ -24,6 +24,10 @@ func (o *OpRenameColumn) Start(ctx context.Context, conn db.DB, latestSchema str
 	}
 	table.RenameColumn(o.From, o.To)
 
+	// Update the name of the column in any constraints that reference the
+	// renamed column.
+	table.RenameConstraintColumns(o.From, o.To)
+
 	return nil, nil
 }
 
@@ -44,6 +48,10 @@ func (o *OpRenameColumn) Complete(ctx context.Context, conn db.DB, tr SQLTransfo
 	column := table.GetColumn(o.To)
 	column.Name = o.To
 
+	// Update the name of the column in any constraints that reference the
+	// renamed column.
+	table.RenameConstraintColumns(o.From, o.To)
+
 	return err
 }
 
@@ -86,6 +94,7 @@ func (o *OpRenameColumn) Validate(ctx context.Context, s *schema.Schema) error {
 	// Update the in-memory schema to reflect the column rename so that it is
 	// visible to subsequent operations' validation steps.
 	table.RenameColumn(o.From, o.To)
+	table.RenameConstraintColumns(o.From, o.To)
 
 	return nil
 }

pkg/schema/schema.go

@@ -277,6 +277,28 @@ func (t *Table) GetConstraintColumns(name string) []string {
 	return slices.Compact(columns)
 }
 
+// RenameConstraintColumns renames all occurrences of a column name in any
+// constraint on the table from `from` to `to`.
+func (t *Table) RenameConstraintColumns(from, to string) {
+	updateColumns := func(columns []string) {
+		for i, c := range columns {
+			if c == from {
+				columns[i] = to
+			}
+		}
+	}
+
+	for _, cc := range t.CheckConstraints {
+		updateColumns(cc.Columns)
+	}
+	for _, uc := range t.UniqueConstraints {
+		updateColumns(uc.Columns)
+	}
+	for _, fk := range t.ForeignKeys {
+		updateColumns(fk.Columns)
+	}
+}
+
 // GetPrimaryKey returns the columns that make up the primary key
 func (t *Table) GetPrimaryKey() (columns []*Column) {
 	for _, name := range t.PrimaryKey {