From 39035555915794aa4d5760a704d738621b3a01ae Mon Sep 17 00:00:00 2001
From: Andrew Farries <andrew.farries@xata.io>
Date: Fri, 12 Jan 2024 13:09:16 +0000
Subject: [PATCH] Add test to ensure preservation of FK constraints

Add a test that checks that setting a column with an FK constraint to
`NOT NULL` does not remove the FK constraint.
---
 pkg/migrations/op_set_notnull_test.go | 77 +++++++++++++++++++++++++++
 1 file changed, 77 insertions(+)

diff --git a/pkg/migrations/op_set_notnull_test.go b/pkg/migrations/op_set_notnull_test.go
index c78fceea..eec21abc 100644
--- a/pkg/migrations/op_set_notnull_test.go
+++ b/pkg/migrations/op_set_notnull_test.go
@@ -222,6 +222,83 @@ func TestSetNotNull(t *testing.T) {
 			afterComplete: func(t *testing.T, db *sql.DB) {
 			},
 		},
+		{
+			name: "setting a foreign key column to not null retains the foreign key constraint",
+			migrations: []migrations.Migration{
+				{
+					Name: "01_add_departments_table",
+					Operations: migrations.Operations{
+						&migrations.OpCreateTable{
+							Name: "departments",
+							Columns: []migrations.Column{
+								{
+									Name: "id",
+									Type: "serial",
+									Pk:   true,
+								},
+								{
+									Name:     "name",
+									Type:     "text",
+									Nullable: false,
+								},
+							},
+						},
+					},
+				},
+				{
+					Name: "02_add_employees_table",
+					Operations: migrations.Operations{
+						&migrations.OpCreateTable{
+							Name: "employees",
+							Columns: []migrations.Column{
+								{
+									Name: "id",
+									Type: "serial",
+									Pk:   true,
+								},
+								{
+									Name:     "name",
+									Type:     "text",
+									Nullable: false,
+								},
+								{
+									Name:     "department_id",
+									Type:     "integer",
+									Nullable: true,
+									References: &migrations.ForeignKeyReference{
+										Name:   "fk_employee_department",
+										Table:  "departments",
+										Column: "id",
+									},
+								},
+							},
+						},
+					},
+				},
+				{
+					Name: "03_set_not_null",
+					Operations: migrations.Operations{
+						&migrations.OpAlterColumn{
+							Table:    "employees",
+							Column:   "department_id",
+							Nullable: ptr(false),
+							Up:       "(SELECT CASE WHEN department_id IS NULL THEN 1 ELSE department_id END)",
+							Down:     "department_id",
+						},
+					},
+				},
+			},
+			afterStart: func(t *testing.T, db *sql.DB) {
+				// A temporary FK constraint has been created on the temporary column
+				ConstraintMustExist(t, db, "public", "employees", migrations.TemporaryName("fk_employee_department"))
+			},
+			afterRollback: func(t *testing.T, db *sql.DB) {
+			},
+			afterComplete: func(t *testing.T, db *sql.DB) {
+				// The foreign key constraint still exists on the column
+				ConstraintMustExist(t, db, "public", "employees", "fk_employee_department")
+			},
+		},
 	})
 }