From cec0d80e92feea834cf30d721445fd882a70d715 Mon Sep 17 00:00:00 2001
From: Andrew Farries <andyrb@gmail.com>
Date: Mon, 20 Jan 2025 11:15:56 +0000
Subject: [PATCH] Add support for `alter_table` set `NOT NULL` operations with
 `create_table` operations (#608)

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

```json
{
  "name": "06_multi_operation",
  "operations": [
    {
      "create_table": {
        "name": "items",
        "columns": [
          {
            "name": "id",
            "type": "serial",
            "pk": true
          },
          {
            "name": "name",
            "type": "text",
            "nullable": true
          }
        ]
      }
    },
    {
      "alter_column": {
        "table": "items",
        "column": "name",
        "nullable": false,
        "up": "SELECT CASE WHEN name IS NULL THEN 'anonymous' ELSE name END",
        "down": "name || '_from_down_trigger'"
      }
    }
  ]
}
```

This migration creates a table and then sets a column's `nullability`.

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

Part of #239
---
 pkg/migrations/op_create_table.go     |  2 +
 pkg/migrations/op_set_notnull_test.go | 67 +++++++++++++++++++++++++++
 2 files changed, 69 insertions(+)

diff --git a/pkg/migrations/op_create_table.go b/pkg/migrations/op_create_table.go
index 855c1bfe..1f6b39fb 100644
--- a/pkg/migrations/op_create_table.go
+++ b/pkg/migrations/op_create_table.go
@@ -181,11 +181,13 @@ func (o *OpCreateTable) updateSchema(s *schema.Schema) *schema.Schema {
 			Name:     col.Name,
 			Unique:   col.Unique,
 			Nullable: col.Nullable,
+			Type:     col.Type,
 		}
 		if col.Pk {
 			primaryKeys = append(primaryKeys, col.Name)
 		}
 	}
+
 	uniqueConstraints := make(map[string]*schema.UniqueConstraint, 0)
 	checkConstraints := make(map[string]*schema.CheckConstraint, 0)
 	for _, c := range o.Constraints {
diff --git a/pkg/migrations/op_set_notnull_test.go b/pkg/migrations/op_set_notnull_test.go
index 53dc69ac..cdf9aced 100644
--- a/pkg/migrations/op_set_notnull_test.go
+++ b/pkg/migrations/op_set_notnull_test.go
@@ -733,6 +733,73 @@ func TestSetNotNullInMultiOperationMigrations(t *testing.T) {
 				TableMustBeCleanedUp(t, db, schema, "products", "item_name")
 			},
 		},
+		{
+			name: "create table, set not null",
+			migrations: []migrations.Migration{
+				{
+					Name: "01_multi_operation",
+					Operations: migrations.Operations{
+						&migrations.OpCreateTable{
+							Name: "items",
+							Columns: []migrations.Column{
+								{
+									Name: "id",
+									Type: "int",
+									Pk:   true,
+								},
+								{
+									Name:     "name",
+									Type:     "varchar(255)",
+									Nullable: true,
+								},
+							},
+						},
+						&migrations.OpAlterColumn{
+							Table:    "items",
+							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 (only) schema that meets the constraint
+				MustInsert(t, db, schema, "01_multi_operation", "items", map[string]string{
+					"id":   "1",
+					"name": "apple",
+				})
+
+				// Can't insert a row into the new (only) schema that violates the constraint
+				MustNotInsert(t, db, schema, "01_multi_operation", "items", map[string]string{
+					"id": "2",
+				}, testutils.CheckViolationErrorCode)
+
+				// The new view has the expected rows
+				rows := MustSelect(t, db, schema, "01_multi_operation", "items")
+				assert.Equal(t, []map[string]any{
+					{"id": 1, "name": "apple"},
+				}, rows)
+			},
+			afterRollback: func(t *testing.T, db *sql.DB, schema string) {
+				// Tht table has been dropped
+				TableMustNotExist(t, db, schema, "items")
+			},
+			afterComplete: func(t *testing.T, db *sql.DB, schema string) {
+				// Can insert a row into the new (only) schema that meets the constraint
+				MustInsert(t, db, schema, "01_multi_operation", "items", map[string]string{
+					"id":   "1",
+					"name": "banana",
+				})
+
+				// The new view has the expected rows
+				rows := MustSelect(t, db, schema, "01_multi_operation", "items")
+				assert.Equal(t, []map[string]any{
+					{"id": 1, "name": "banana"},
+				}, rows)
+			},
+		},
 	})
 }