Add more test cases

xataio · Dec 18, 2024 · 2491530 · 2491530
1 parent 73906e3
commit 2491530
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 5 deletions.
diff --git 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
@@ -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 {