Fix ADD PRIMARY KEY after dropping and re-adding a column

Previously, dropping a NOT NULL column, adding a new column with the same
name, and making it a primary key in a single ALTER TABLE command could
fail with "primary key column ... is not marked NOT NULL". This happened
because ALTER TABLE checked for an existing NOT NULL constraint before
dropping the old column. Finding the constraint caused it to skip adding
one for the new column, even though the old constraint would be removed
with the old column.

Defer this check until after column drops and additions have been
executed, so that it examines the columns that will actually form the
primary key. Add any required NOT NULL constraints before creating the
primary-key index.

Backpatch to v18, where this issue was introduced by commit a379061a22.

Reported-by: Ludvig Janiuk <ludvig.janiuk@proton.me>
Author: Fujii Masao <masao.fujii@gmail.com>
Co-authored-by: Manuel Reyes Bravo <manuelreyesbravo@gmail.com>
Discussion: https://postgr.es/m/0a9b4c44-cdff-4ec8-a8fb-b3a8b59e9ba9@proton.me
Backpatch-through: 18

src/backend/commands/tablecmds.c

@@ -538,9 +538,8 @@ static ObjectAddress ATExecDropColumn(List **wqueue, Relation rel, const char *c
 									  bool recurse, bool recursing,
 									  bool missing_ok, LOCKMODE lockmode,
 									  ObjectAddresses *addrs);
-static void ATPrepAddPrimaryKey(List **wqueue, Relation rel, AlterTableCmd *cmd,
-								bool recurse, LOCKMODE lockmode,
-								AlterTableUtilityContext *context);
+static void ATPrepAddPrimaryKey(AlteredTableInfo *tab, Relation rel,
+								AlterTableCmd *cmd, LOCKMODE lockmode);
 static void verifyNotNullPKCompatible(HeapTuple tuple, const char *colname);
 static ObjectAddress ATExecAddIndex(AlteredTableInfo *tab, Relation rel,
 									IndexStmt *stmt, bool is_rebuild, LOCKMODE lockmode);
@@ -5096,7 +5095,6 @@ ATPrepCmd(List **wqueue, Relation rel, AlterTableCmd *cmd,
 		case AT_AddConstraint:	/* ADD CONSTRAINT */
 			ATSimplePermissions(cmd->subtype, rel,
 								ATT_TABLE | ATT_PARTITIONED_TABLE | ATT_FOREIGN_TABLE);
-			ATPrepAddPrimaryKey(wqueue, rel, cmd, recurse, lockmode, context);
 			if (recurse)
 			{
 				/* recurses at exec time; lock descendants and set flag */
@@ -5344,6 +5342,19 @@ ATRewriteCatalogs(List **wqueue, LOCKMODE lockmode,
 						  lfirst_node(AlterTableCmd, lcmd),
 						  lockmode, pass, context);
 
+			/*
+			 * For ADD PRIMARY KEY without USING INDEX, queue any required NOT
+			 * NULL constraints after transforming all explicit constraints.
+			 * This lets a named NOT NULL constraint take precedence
+			 * regardless of subcommand order. USING INDEX is handled during
+			 * transformation.
+			 */
+			if (pass == AT_PASS_ADD_CONSTR)
+			{
+				foreach_node(AlterTableCmd, cmd, subcmds)
+					ATPrepAddPrimaryKey(tab, tab->rel, cmd, lockmode);
+			}
+
 			/*
 			 * After the ALTER TYPE or SET EXPRESSION pass, do cleanup work
 			 * (this is not done in ATExecAlterColumnType since it should be
@@ -9490,7 +9501,7 @@ ATExecDropColumn(List **wqueue, Relation rel, const char *colName,
  * hierarchy (failing to ensure that leads to funny corner cases).  For the
  * normal case where we're asked to recurse, this routine checks if the
  * not-null constraints exist already, and if not queues a requirement for
- * them to be created by phase 2.
+ * them to be created in the AT_PASS_COL_ATTRS pass.
  *
  * For the case where we're asked not to recurse, we verify that a not-null
  * constraint exists on each column of each (direct) child table, throwing an
@@ -9505,16 +9516,17 @@ ATExecDropColumn(List **wqueue, Relation rel, const char *colName,
  * deadlocks during parallel pg_restore of PKs on partitioned tables.
  */
 static void
-ATPrepAddPrimaryKey(List **wqueue, Relation rel, AlterTableCmd *cmd,
-					bool recurse, LOCKMODE lockmode,
-					AlterTableUtilityContext *context)
+ATPrepAddPrimaryKey(AlteredTableInfo *tab, Relation rel, AlterTableCmd *cmd,
+					LOCKMODE lockmode)
 {
 	Constraint *pkconstr;
 	List	   *children = NIL;
 	bool		got_children = false;
 
 	pkconstr = castNode(Constraint, cmd->def);
-	if (pkconstr->contype != CONSTR_PRIMARY)
+
+	/* USING INDEX queues its own NOT NULL constraints during transformation */
+	if (pkconstr->contype != CONSTR_PRIMARY || pkconstr->indexname != NULL)
 		return;
 
 	/* Verify that columns are not-null, or request that they be made so */
@@ -9538,7 +9550,7 @@ ATPrepAddPrimaryKey(List **wqueue, Relation rel, AlterTableCmd *cmd,
 			heap_freetuple(tuple);
 			continue;
 		}
-		else if (!recurse)
+		else if (!cmd->recurse)
 		{
 			/*
 			 * No constraint on this column.  Asked not to recurse, we won't
@@ -9575,7 +9587,8 @@ ATPrepAddPrimaryKey(List **wqueue, Relation rel, AlterTableCmd *cmd,
 		newcmd->recurse = true;
 		newcmd->def = (Node *) nnconstr;
 
-		ATPrepCmd(wqueue, rel, newcmd, true, false, lockmode, context);
+		tab->subcmds[AT_PASS_COL_ATTRS] =
+			lappend(tab->subcmds[AT_PASS_COL_ATTRS], newcmd);
 	}
 }
 

src/test/regress/expected/alter_table.out

@@ -1041,6 +1041,31 @@ ERROR:  column "test2" of relation "atacc1" contains null values
 -- now add a primary key column with a default (succeeds).
 alter table atacc1 add column test2 int default 0 primary key;
 drop table atacc1;
+-- Replacing a NOT NULL column must create a new constraint for the primary key.
+create table atacc1 (id int not null);
+alter table atacc1 drop column id, add column id int, add primary key (id);
+drop table atacc1;
+-- A NOT NULL constraint dropped in the same ALTER TABLE must be replaced.
+create table atacc1 (id int not null);
+alter table atacc1 alter column id drop not null, add primary key (id);
+drop table atacc1;
+create table atacc1 (id int constraint atacc1_id_nn not null);
+alter table atacc1 drop constraint atacc1_id_nn, add primary key (id);
+drop table atacc1;
+-- ADD PRIMARY KEY must prefer an explicitly named replacement NOT NULL.
+create table atacc1 (id int not null);
+alter table atacc1 alter column id drop not null,
+  add primary key (id), add constraint atacc1_id_nn not null id;
+drop table atacc1;
+create table atacc1 (id int not null);
+alter table atacc1 alter column id drop not null,
+  add constraint atacc1_id_nn not null id, add primary key (id);
+drop table atacc1;
+-- The replacement NOT NULL constraint must propagate to the partitions too.
+create table atacc1 (id int not null) partition by range (id);
+create table atacc1_1 partition of atacc1 for values from (0) to (10);
+alter table atacc1 alter column id drop not null, add primary key (id);
+drop table atacc1;
 -- this combination used to have order-of-execution problems (bug #15580)
 create table atacc1 (a int);
 insert into atacc1 values(1);

src/test/regress/sql/alter_table.sql

@@ -782,6 +782,32 @@ alter table atacc1 add column test2 int primary key;
 alter table atacc1 add column test2 int default 0 primary key;
 drop table atacc1;
 
+-- Replacing a NOT NULL column must create a new constraint for the primary key.
+create table atacc1 (id int not null);
+alter table atacc1 drop column id, add column id int, add primary key (id);
+drop table atacc1;
+-- A NOT NULL constraint dropped in the same ALTER TABLE must be replaced.
+create table atacc1 (id int not null);
+alter table atacc1 alter column id drop not null, add primary key (id);
+drop table atacc1;
+create table atacc1 (id int constraint atacc1_id_nn not null);
+alter table atacc1 drop constraint atacc1_id_nn, add primary key (id);
+drop table atacc1;
+-- ADD PRIMARY KEY must prefer an explicitly named replacement NOT NULL.
+create table atacc1 (id int not null);
+alter table atacc1 alter column id drop not null,
+  add primary key (id), add constraint atacc1_id_nn not null id;
+drop table atacc1;
+create table atacc1 (id int not null);
+alter table atacc1 alter column id drop not null,
+  add constraint atacc1_id_nn not null id, add primary key (id);
+drop table atacc1;
+-- The replacement NOT NULL constraint must propagate to the partitions too.
+create table atacc1 (id int not null) partition by range (id);
+create table atacc1_1 partition of atacc1 for values from (0) to (10);
+alter table atacc1 alter column id drop not null, add primary key (id);
+drop table atacc1;
+
 -- this combination used to have order-of-execution problems (bug #15580)
 create table atacc1 (a int);
 insert into atacc1 values(1);