From 2491530ffe20eee157633a4896ecc802ce43227d Mon Sep 17 00:00:00 2001
From: Ryan Slade <ryanslade@gmail.com>
Date: Wed, 18 Dec 2024 10:42:44 +0100
Subject: [PATCH] Add more test cases

---
 pkg/sql2pgroll/alter_table.go      | 27 ++++++++++++++++++++++-----
 pkg/sql2pgroll/alter_table_test.go | 15 +++++++++++++++
 2 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/pkg/sql2pgroll/alter_table.go b/pkg/sql2pgroll/alter_table.go
index 14ca6ee2..33c40ae0 100644
--- a/pkg/sql2pgroll/alter_table.go
+++ b/pkg/sql2pgroll/alter_table.go
@@ -467,13 +467,30 @@ func convertAlterTableAddColumn(stmt *pgq.AlterTableStmt, cmd *pgq.AlterTableCmd
 }
 
 func canConvertAddColumn(cmd *pgq.AlterTableCmd) bool {
+	if cmd.GetMissingOk() {
+		return false
+	}
 	for _, constraint := range cmd.GetDef().GetColumnDef().GetConstraints() {
-		switch constraint.GetConstraint().GetFkUpdAction() {
-		case "r", "c", "n", "d":
-			// RESTRICT, CASCADE, SET NULL, SET DEFAULT
+		switch constraint.GetConstraint().GetContype() {
+		case
+			pgq.ConstrType_CONSTR_DEFAULT,
+			pgq.ConstrType_CONSTR_NULL,
+			pgq.ConstrType_CONSTR_NOTNULL,
+			pgq.ConstrType_CONSTR_PRIMARY,
+			pgq.ConstrType_CONSTR_UNIQUE,
+			pgq.ConstrType_CONSTR_FOREIGN,
+			pgq.ConstrType_CONSTR_CHECK:
+			switch constraint.GetConstraint().GetFkUpdAction() {
+			case "r", "c", "n", "d":
+				// RESTRICT, CASCADE, SET NULL, SET DEFAULT
+				return false
+			case "a":
+				// NO ACTION, the default
+				break
+			}
+		case pgq.ConstrType_CONSTR_ATTR_DEFERRABLE, pgq.ConstrType_CONSTR_ATTR_DEFERRED:
 			return false
-		case "a":
-			// NO ACTION, the default
+		case pgq.ConstrType_CONSTR_ATTR_NOT_DEFERRABLE, pgq.ConstrType_CONSTR_ATTR_IMMEDIATE:
 			break
 		}
 	}
diff --git a/pkg/sql2pgroll/alter_table_test.go b/pkg/sql2pgroll/alter_table_test.go
index 28b3c42b..e28cfbc6 100644
--- a/pkg/sql2pgroll/alter_table_test.go
+++ b/pkg/sql2pgroll/alter_table_test.go
@@ -146,6 +146,10 @@ func TestConvertAlterTableStatements(t *testing.T) {
 			sql:        "ALTER TABLE foo ADD COLUMN bar int",
 			expectedOp: expect.AddColumnOp1,
 		},
+		{
+			sql:        "ALTER TABLE foo ADD COLUMN bar int NOT NULL",
+			expectedOp: expect.AddColumnOp1,
+		},
 		{
 			sql:        "ALTER TABLE schema.foo ADD COLUMN bar int",
 			expectedOp: expect.AddColumnOp2,
@@ -170,6 +174,14 @@ func TestConvertAlterTableStatements(t *testing.T) {
 			sql:        "ALTER TABLE foo ADD COLUMN bar int UNIQUE",
 			expectedOp: expect.AddColumnOp4,
 		},
+		{
+			sql:        "ALTER TABLE foo ADD COLUMN bar int UNIQUE NOT DEFERRABLE",
+			expectedOp: expect.AddColumnOp4,
+		},
+		{
+			sql:        "ALTER TABLE foo ADD COLUMN bar int UNIQUE INITIALLY IMMEDIATE",
+			expectedOp: expect.AddColumnOp4,
+		},
 		{
 			sql:        "ALTER TABLE foo ADD COLUMN bar int PRIMARY KEY",
 			expectedOp: expect.AddColumnOp5,
@@ -267,6 +279,9 @@ func TestUnconvertableAlterTableStatements(t *testing.T) {
 		"ALTER TABLE foo ADD COLUMN bar int REFERENCES bar (c) ON UPDATE CASCADE",
 		"ALTER TABLE foo ADD COLUMN bar int REFERENCES bar (c) ON UPDATE SET NULL",
 		"ALTER TABLE foo ADD COLUMN bar int REFERENCES bar (c) ON UPDATE SET DEFAULT",
+		"ALTER TABLE foo ADD COLUMN IF NOT EXISTS bar int",
+		"ALTER TABLE foo ADD COLUMN bar int UNIQUE DEFERRABLE",
+		"ALTER TABLE foo ADD COLUMN bar int UNIQUE INITIALLY DEFERRED",
 	}
 
 	for _, sql := range tests {