Add support for alter_table set NOT NULL operations with rename_table operations (#606)

Ensure that multi-operation migrations combining `alter_column` `SET NOT
NULL` and `rename_table` operations work as expected.

```json
{
  "name": "06_multi_operation",
  "operations": [
    {
      "rename_table": {
        "from": "items",
        "to": "products"
      }
    },
    {
      "alter_column": {
        "table": "products",
        "column": "name",
        "nullable": false,
        "up": "SELECT CASE WHEN name IS NULL THEN 'anonymous' ELSE name END",
        "down": "name || '_from_down_trigger'"
      }
    }
  ]
}
```

This migration renames a table and then sets a column's `nullability` on
the renamed table.

Previously the migration would fail as the `alter_column` operation was
unaware of the changes made by the preceding operation.

Part of #239

Author: andrew-farries

pkg/migrations/op_alter_column.go

@@ -39,7 +39,7 @@ func (o *OpAlterColumn) Start(ctx context.Context, conn db.DB, latestSchema stri
 		Columns:        table.Columns,
 		SchemaName:     s.Name,
 		LatestSchema:   latestSchema,
-		TableName:      o.Table,
+		TableName:      table.Name,
 		PhysicalColumn: TemporaryName(o.Column),
 		SQL:            o.upSQLForOperations(ops),
 	})
@@ -61,7 +61,7 @@ func (o *OpAlterColumn) Start(ctx context.Context, conn db.DB, latestSchema stri
 		Columns:        table.Columns,
 		LatestSchema:   latestSchema,
 		SchemaName:     s.Name,
-		TableName:      o.Table,
+		TableName:      table.Name,
 		PhysicalColumn: o.Column,
 		SQL:            o.downSQLForOperations(ops),
 	})
@@ -115,9 +115,10 @@ func (o *OpAlterColumn) Complete(ctx context.Context, conn db.DB, tr SQLTransfor
 }
 
 func (o *OpAlterColumn) Rollback(ctx context.Context, conn db.DB, tr SQLTransformer, s *schema.Schema) error {
-	ops := o.subOperations()
+	table := s.GetTable(o.Table)
 
 	// Perform any operation specific rollback steps
+	ops := o.subOperations()
 	for _, ops := range ops {
 		if err := ops.Rollback(ctx, conn, tr, nil); err != nil {
 			return err
@@ -126,7 +127,7 @@ func (o *OpAlterColumn) Rollback(ctx context.Context, conn db.DB, tr SQLTransfor
 
 	// 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(o.Column)),
 	))
 	if err != nil {

pkg/migrations/op_create_table.go

@@ -59,7 +59,9 @@ func (o *OpCreateTable) Start(ctx context.Context, conn db.DB, latestSchema stri
 }
 
 func (o *OpCreateTable) Complete(ctx context.Context, conn db.DB, tr SQLTransformer, s *schema.Schema) error {
-	// No-op
+	// Update the in-memory schema representation with the new table
+	o.updateSchema(s)
+
 	return nil
 }
 

pkg/migrations/op_rename_table.go

@@ -22,6 +22,16 @@ func (o *OpRenameTable) Complete(ctx context.Context, conn db.DB, tr SQLTransfor
 	_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s RENAME TO %s",
 		pq.QuoteIdentifier(o.From),
 		pq.QuoteIdentifier(o.To)))
+
+	// Rename the table in the virtual schema so that the `Complete` methods
+	// of subsequent operations in the same migration can find it.
+	s.RenameTable(o.From, o.To)
+
+	// Update the physical name of the table in the virtual schema now that it
+	// has really been renamed.
+	table := s.GetTable(o.To)
+	table.Name = o.To
+
 	return err
 }
 

pkg/migrations/op_set_notnull.go

@@ -24,7 +24,7 @@ func (o *OpSetNotNull) Start(ctx context.Context, conn db.DB, latestSchema strin
 	table := s.GetTable(o.Table)
 
 	// Add an unchecked NOT NULL constraint to the new column.
-	if err := addNotNullConstraint(ctx, conn, o.Table, o.Column, TemporaryName(o.Column)); err != nil {
+	if err := addNotNullConstraint(ctx, conn, table.Name, o.Column, TemporaryName(o.Column)); err != nil {
 		return nil, fmt.Errorf("failed to add not null constraint: %w", err)
 	}
 

pkg/migrations/op_set_notnull_test.go

@@ -531,6 +531,112 @@ func TestSetNotNull(t *testing.T) {
 	})
 }
 
+func TestSetNotNullInMultiOperationMigrations(t *testing.T) {
+	t.Parallel()
+
+	ExecuteTests(t, TestCases{
+		{
+			name: "rename table, set not null",
+			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,
+								},
+							},
+						},
+					},
+				},
+				{
+					Name: "02_multi_operation",
+					Operations: migrations.Operations{
+						&migrations.OpRenameTable{
+							From: "items",
+							To:   "products",
+						},
+						&migrations.OpAlterColumn{
+							Table:    "products",
+							Column:   "name",
+							Nullable: ptr(false),
+							Up:       "SELECT CASE WHEN name IS NULL THEN 'unknown' ELSE name END",
+							Down:     "name",
+						},
+					},
+				},
+			},
+			afterStart: func(t *testing.T, db *sql.DB, schema string) {
+				// Can insert a row into the new schema that meets the constraint
+				MustInsert(t, db, schema, "02_multi_operation", "products", map[string]string{
+					"id":   "1",
+					"name": "apple",
+				})
+
+				// Can't insert a row into the new schema that violates the constraint
+				MustNotInsert(t, db, schema, "02_multi_operation", "products", map[string]string{
+					"id": "2",
+				}, testutils.CheckViolationErrorCode)
+
+				// Can insert a row into the old schema that violates the constraint
+				MustInsert(t, db, schema, "01_create_table", "items", map[string]string{
+					"id": "2",
+				})
+
+				// 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": "apple"},
+					{"id": 2, "name": "unknown"},
+				}, 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": "apple"},
+					{"id": 2, "name": nil},
+				}, 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 meets the constraint
+				MustInsert(t, db, schema, "02_multi_operation", "products", map[string]string{
+					"id":   "3",
+					"name": "carrot",
+				})
+
+				// Can't insert a row into the new schema that violates the constraint
+				MustNotInsert(t, db, schema, "02_multi_operation", "products", map[string]string{
+					"id": "3",
+				}, testutils.NotNullViolationErrorCode)
+
+				// 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": "apple"},
+					{"id": 2, "name": "unknown"},
+					{"id": 3, "name": "carrot"},
+				}, rows)
+
+				// The table has been cleaned up
+				TableMustBeCleanedUp(t, db, schema, "products", "name")
+			},
+		},
+	})
+}
+
 func TestSetNotNullValidation(t *testing.T) {
 	t.Parallel()