From 710ba9b55d9fcf38464c9d2e00e2946e051425ab Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81lvaro=20Herrera?= <alvherre@alvh.no-ip.org>
Date: Tue, 15 Apr 2025 20:38:54 +0200
Subject: [PATCH v2] Fix verification of not-null constraints on children
 during PK creation

---
 src/backend/commands/tablecmds.c          | 145 +++++++++++++---------
 src/test/regress/expected/constraints.out |  22 +++-
 src/test/regress/sql/constraints.sql      |  14 +++
 3 files changed, 122 insertions(+), 59 deletions(-)

diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index b3ed69457fc..a485f045890 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -540,6 +540,7 @@ static ObjectAddress ATExecDropColumn(List **wqueue, Relation rel, const char *c
 static void ATPrepAddPrimaryKey(List **wqueue, Relation rel, AlterTableCmd *cmd,
 								bool recurse, LOCKMODE lockmode,
 								AlterTableUtilityContext *context);
+static void verifyNotNullPKCompatible(HeapTuple tuple, const char *colname);
 static ObjectAddress ATExecAddIndex(AlteredTableInfo *tab, Relation rel,
 									IndexStmt *stmt, bool is_rebuild, LOCKMODE lockmode);
 static ObjectAddress ATExecAddStatistics(AlteredTableInfo *tab, Relation rel,
@@ -9438,8 +9439,26 @@ ATExecDropColumn(List **wqueue, Relation rel, const char *colName,
 }
 
 /*
- * Prepare to add a primary key on table, by adding not-null constraints
+ * Prepare to add a primary key on a table, by adding not-null constraints
  * on all columns.
+ *
+ * The not-null constraints for a primary key must cover the whole inheritance
+ * hierarchy (failing to ensure that leads to funny corner cases).  For the
+ * normal case where we're asked to recurse, this routine ensures that the
+ * not-null constraints either exist already, or queues a requirement for them
+ * to be created by phase 2.
+ *
+ * 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
+ * error if not.  Not throwing an error would also work, because a not-null
+ * constraint would be created anyway, but it'd cause a silent scan of the
+ * child table to verify absence of nulls.  We prefer to let the user know so
+ * that they can add the constraint manually without having to hold
+ * AccessExclusiveLock while at it.
+ *
+ * However, it's also important that we do not acquire locks on children if
+ * the not-null constraints already exist on the parent, to avoid risking
+ * deadlocks during parallel pg_restore of PKs on partitioned tables.
  */
 static void
 ATPrepAddPrimaryKey(List **wqueue, Relation rel, AlterTableCmd *cmd,
@@ -9447,42 +9466,13 @@ ATPrepAddPrimaryKey(List **wqueue, Relation rel, AlterTableCmd *cmd,
 					AlterTableUtilityContext *context)
 {
 	Constraint *pkconstr;
+	List	   *children;
+	bool		got_children = false;
 
 	pkconstr = castNode(Constraint, cmd->def);
 	if (pkconstr->contype != CONSTR_PRIMARY)
 		return;
 
-	/*
-	 * If not recursing, we must ensure that all children have a NOT NULL
-	 * constraint on the columns, and error out if not.
-	 */
-	if (!recurse)
-	{
-		List	   *children;
-
-		children = find_inheritance_children(RelationGetRelid(rel),
-											 lockmode);
-		foreach_oid(childrelid, children)
-		{
-			foreach_node(String, attname, pkconstr->keys)
-			{
-				HeapTuple	tup;
-				Form_pg_attribute attrForm;
-
-				tup = SearchSysCacheAttName(childrelid, strVal(attname));
-				if (!tup)
-					elog(ERROR, "cache lookup failed for attribute %s of relation %u",
-						 strVal(attname), childrelid);
-				attrForm = (Form_pg_attribute) GETSTRUCT(tup);
-				if (!attrForm->attnotnull)
-					ereport(ERROR,
-							errmsg("column \"%s\" of table \"%s\" is not marked NOT NULL",
-								   strVal(attname), get_rel_name(childrelid)));
-				ReleaseSysCache(tup);
-			}
-		}
-	}
-
 	/* Verify that columns are not-null, or request that they be made so */
 	foreach_node(String, column, pkconstr->keys)
 	{
@@ -9498,42 +9488,46 @@ ATPrepAddPrimaryKey(List **wqueue, Relation rel, AlterTableCmd *cmd,
 		tuple = findNotNullConstraint(RelationGetRelid(rel), strVal(column));
 		if (tuple != NULL)
 		{
-			Form_pg_constraint conForm = (Form_pg_constraint) GETSTRUCT(tuple);
-
-			/* a NO INHERIT constraint is no good */
-			if (conForm->connoinherit)
-				ereport(ERROR,
-						errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
-						errmsg("cannot create primary key on column \"%s\"",
-							   strVal(column)),
-				/*- translator: third %s is a constraint characteristic such as NOT VALID */
-						errdetail("The constraint \"%s\" on column \"%s\", marked %s, is incompatible with a primary key.",
-								  NameStr(conForm->conname), strVal(column), "NO INHERIT"),
-						errhint("You will need to make it inheritable using %s.",
-								"ALTER TABLE ... ALTER CONSTRAINT ... INHERIT"));
-
-			/* an unvalidated constraint is no good */
-			if (!conForm->convalidated)
-				ereport(ERROR,
-						errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
-						errmsg("cannot create primary key on column \"%s\"",
-							   strVal(column)),
-				/*- translator: third %s is a constraint characteristic such as NOT VALID */
-						errdetail("The constraint \"%s\" on column \"%s\", marked %s, is incompatible with a primary key.",
-								  NameStr(conForm->conname), strVal(column), "NOT VALID"),
-						errhint("You will need to validate it using %s.",
-								"ALTER TABLE ... VALIDATE CONSTRAINT"));
+			verifyNotNullPKCompatible(tuple, strVal(column));
 
 			/* All good with this one; don't request another */
 			heap_freetuple(tuple);
 			continue;
 		}
+		else if (!recurse)
+		{
+			/*
+			 * No constraint on this column.  Asked not to recurse, we won't
+			 * create one here, but verify that all children have one.
+			 */
+			if (!got_children)
+			{
+				children = find_inheritance_children(RelationGetRelid(rel),
+													 lockmode);
+				/* only search for children on the first time through */
+				got_children = true;
+			}
+
+			foreach_oid(childrelid, children)
+			{
+				HeapTuple	tup;
+
+				tup = findNotNullConstraint(childrelid, strVal(column));
+				if (!tup)
+					ereport(ERROR,
+							errmsg("column \"%s\" of table \"%s\" is not marked NOT NULL",
+								   strVal(column), get_rel_name(childrelid)));
+				/* verify it's good enough */
+				verifyNotNullPKCompatible(tup, strVal(column));
+			}
+		}
 
 		/* This column is not already not-null, so add it to the queue */
 		nnconstr = makeNotNullConstraint(column);
 
 		newcmd = makeNode(AlterTableCmd);
 		newcmd->subtype = AT_AddConstraint;
+		/* note we force recurse=true here; see above */
 		newcmd->recurse = true;
 		newcmd->def = (Node *) nnconstr;
 
@@ -9541,6 +9535,43 @@ ATPrepAddPrimaryKey(List **wqueue, Relation rel, AlterTableCmd *cmd,
 	}
 }
 
+/*
+ * Verify whether the given not-null constraint is compatible with a
+ * primary key.  If not, an error is thrown.
+ */
+static void
+verifyNotNullPKCompatible(HeapTuple tuple, const char *colname)
+{
+	Form_pg_constraint conForm = (Form_pg_constraint) GETSTRUCT(tuple);
+
+	if (conForm->contype != CONSTRAINT_NOTNULL)
+		elog(ERROR, "constraint %u is not a not-null constraint", conForm->oid);
+
+	/* a NO INHERIT constraint is no good */
+	if (conForm->connoinherit)
+		ereport(ERROR,
+				errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+				errmsg("cannot create primary key on column \"%s\"", colname),
+		/*- translator: third %s is a constraint characteristic such as NOT VALID */
+				errdetail("The constraint \"%s\" on column \"%s\" of table \"%s\", marked %s, is incompatible with a primary key.",
+						  NameStr(conForm->conname), colname,
+						  get_rel_name(conForm->conrelid), "NO INHERIT"),
+				errhint("You will need to make it inheritable using %s.",
+						"ALTER TABLE ... ALTER CONSTRAINT ... INHERIT"));
+
+	/* an unvalidated constraint is no good */
+	if (!conForm->convalidated)
+		ereport(ERROR,
+				errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+				errmsg("cannot create primary key on column \"%s\"", colname),
+		/*- translator: third %s is a constraint characteristic such as NOT VALID */
+				errdetail("The constraint \"%s\" on column \"%s\" of table \"%s\", marked %s, is incompatible with a primary key.",
+						  NameStr(conForm->conname), colname,
+						  get_rel_name(conForm->conrelid), "NOT VALID"),
+				errhint("You will need to validate it using %s.",
+						"ALTER TABLE ... VALIDATE CONSTRAINT"));
+}
+
 /*
  * ALTER TABLE ADD INDEX
  *
diff --git a/src/test/regress/expected/constraints.out b/src/test/regress/expected/constraints.out
index a8c6495ae01..c151ecf76e6 100644
--- a/src/test/regress/expected/constraints.out
+++ b/src/test/regress/expected/constraints.out
@@ -1233,7 +1233,7 @@ Indexes:
 create table cnn_pk (a int not null no inherit);
 alter table cnn_pk add primary key (a);
 ERROR:  cannot create primary key on column "a"
-DETAIL:  The constraint "cnn_pk_a_not_null" on column "a", marked NO INHERIT, is incompatible with a primary key.
+DETAIL:  The constraint "cnn_pk_a_not_null" on column "a" of table "cnn_pk", marked NO INHERIT, is incompatible with a primary key.
 HINT:  You will need to make it inheritable using ALTER TABLE ... ALTER CONSTRAINT ... INHERIT.
 drop table cnn_pk;
 -- Ensure partitions are scanned for null values when adding a PK
@@ -1395,7 +1395,7 @@ HINT:  You will need to use ALTER TABLE ... VALIDATE CONSTRAINT to validate it.
 -- cannot add primary key on a column with an invalid not-null
 ALTER TABLE notnull_tbl1 ADD PRIMARY KEY (a);
 ERROR:  cannot create primary key on column "a"
-DETAIL:  The constraint "nn" on column "a", marked NOT VALID, is incompatible with a primary key.
+DETAIL:  The constraint "nn" on column "a" of table "notnull_tbl1", marked NOT VALID, is incompatible with a primary key.
 HINT:  You will need to validate it using ALTER TABLE ... VALIDATE CONSTRAINT.
 -- ALTER column SET NOT NULL validates an invalid constraint (but this fails
 -- because of rows with null values)
@@ -1567,6 +1567,24 @@ ERROR:  constraint "nn1" conflicts with NOT VALID constraint on child table "pp_
 ALTER TABLE pp_nn_1 VALIDATE CONSTRAINT nn1;
 ALTER TABLE pp_nn ATTACH PARTITION pp_nn_1 FOR VALUES IN (NULL,5); --ok
 DROP TABLE pp_nn;
+-- Try a partition with an invalid constraint and create a PK on the parent.
+CREATE TABLE pp_nn (a int) PARTITION BY HASH (a);
+CREATE TABLE pp_nn_1 PARTITION OF pp_nn FOR VALUES WITH (MODULUS 2, REMAINDER 0);
+ALTER TABLE pp_nn_1 ADD CONSTRAINT nn NOT NULL a NOT VALID;
+ALTER TABLE ONLY pp_nn ADD PRIMARY KEY (a);
+ERROR:  cannot create primary key on column "a"
+DETAIL:  The constraint "nn" on column "a" of table "pp_nn_1", marked NOT VALID, is incompatible with a primary key.
+HINT:  You will need to validate it using ALTER TABLE ... VALIDATE CONSTRAINT.
+DROP TABLE pp_nn;
+-- same as above, but the constraint is NO INHERIT
+CREATE TABLE pp_nn (a int) PARTITION BY HASH (a);
+CREATE TABLE pp_nn_1 PARTITION OF pp_nn FOR VALUES WITH (MODULUS 2, REMAINDER 0);
+ALTER TABLE pp_nn_1 ADD CONSTRAINT nn NOT NULL a NO INHERIT;
+ALTER TABLE ONLY pp_nn ADD PRIMARY KEY (a);
+ERROR:  cannot create primary key on column "a"
+DETAIL:  The constraint "nn" on column "a" of table "pp_nn_1", marked NO INHERIT, is incompatible with a primary key.
+HINT:  You will need to make it inheritable using ALTER TABLE ... ALTER CONSTRAINT ... INHERIT.
+DROP TABLE pp_nn;
 -- Create table with NOT NULL INVALID constraint, for pg_upgrade.
 CREATE TABLE notnull_tbl1_upg (a int, b int);
 INSERT INTO notnull_tbl1_upg VALUES (NULL, 1), (NULL, 2), (300, 3);
diff --git a/src/test/regress/sql/constraints.sql b/src/test/regress/sql/constraints.sql
index bf8f0aa181d..5d6d749c150 100644
--- a/src/test/regress/sql/constraints.sql
+++ b/src/test/regress/sql/constraints.sql
@@ -940,6 +940,20 @@ ALTER TABLE pp_nn_1 VALIDATE CONSTRAINT nn1;
 ALTER TABLE pp_nn ATTACH PARTITION pp_nn_1 FOR VALUES IN (NULL,5); --ok
 DROP TABLE pp_nn;
 
+-- Try a partition with an invalid constraint and create a PK on the parent.
+CREATE TABLE pp_nn (a int) PARTITION BY HASH (a);
+CREATE TABLE pp_nn_1 PARTITION OF pp_nn FOR VALUES WITH (MODULUS 2, REMAINDER 0);
+ALTER TABLE pp_nn_1 ADD CONSTRAINT nn NOT NULL a NOT VALID;
+ALTER TABLE ONLY pp_nn ADD PRIMARY KEY (a);
+DROP TABLE pp_nn;
+
+-- same as above, but the constraint is NO INHERIT
+CREATE TABLE pp_nn (a int) PARTITION BY HASH (a);
+CREATE TABLE pp_nn_1 PARTITION OF pp_nn FOR VALUES WITH (MODULUS 2, REMAINDER 0);
+ALTER TABLE pp_nn_1 ADD CONSTRAINT nn NOT NULL a NO INHERIT;
+ALTER TABLE ONLY pp_nn ADD PRIMARY KEY (a);
+DROP TABLE pp_nn;
+
 -- Create table with NOT NULL INVALID constraint, for pg_upgrade.
 CREATE TABLE notnull_tbl1_upg (a int, b int);
 INSERT INTO notnull_tbl1_upg VALUES (NULL, 1), (NULL, 2), (300, 3);
-- 
2.39.5

