From 8aacd6aa67ad6891aaea7a16b9304897798fef39 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] Fix verification of not-null constraints on children during
 PK creation

---
 src/backend/commands/tablecmds.c | 80 +++++++++++++++++++-------------
 1 file changed, 49 insertions(+), 31 deletions(-)

diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index b3ed69457fc..03dd3754940 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -9440,6 +9440,15 @@ ATExecDropColumn(List **wqueue, Relation rel, const char *colName,
 /*
  * Prepare to add a primary key on table, by adding not-null constraints
  * on all columns.
+ *
+ * An important aspect of this is pg_dump creation of primary keys on
+ * partitioned tables, which is done by first creating the primary key
+ * constraint on the partitioned table itself as not-recursive (so that the
+ * creation of the PK itself doesn't recurse to the partitions), then creating
+ * the corresponding indexes on the partitions, then doing ALTER INDEX ATTACH
+ * PARTITION.  This maximizes parallelism.  However, it means that we must
+ * ensure the creation of not-null constraints on the partitions even if asked
+ * not to recurse.
  */
 static void
 ATPrepAddPrimaryKey(List **wqueue, Relation rel, AlterTableCmd *cmd,
@@ -9447,42 +9456,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)
 	{
@@ -9528,12 +9508,50 @@ ATPrepAddPrimaryKey(List **wqueue, Relation rel, AlterTableCmd *cmd,
 			heap_freetuple(tuple);
 			continue;
 		}
+		else if (!recurse)
+		{
+			/*
+			 * If a not-null constraint for this column doesn't exist, and we
+			 * were asked not to recurse to children for the primary key, then
+			 * we must verify that the columns on children tables already have
+			 * a not-null constraint.  The reason for this, is that the
+			 * constraint is necessary so that our pg_dump strategy for
+			 * partitioned tables works, as explained above; and we don't want
+			 * such constraints to be created implicitly by the
+			 * makeNotNullConstraint() call below.  So here we check that a
+			 * not-null constraint exists on this column and raise an error if
+			 * not.  (We don't need to check on columns where a not-null
+			 * constraint exists on the parent, as we already verified that
+			 * it's not NO INHERIT.)
+			 */
+			if (!got_children)
+				children = find_inheritance_children(RelationGetRelid(rel),
+													 lockmode);
+
+			foreach_oid(childrelid, children)
+			{
+				HeapTuple	tup;
+				Form_pg_attribute attrForm;
+
+				tup = SearchSysCacheAttName(childrelid, strVal(column));
+				if (!tup)
+					elog(ERROR, "cache lookup failed for attribute %s of relation %u",
+						 strVal(column), childrelid);
+				attrForm = (Form_pg_attribute) GETSTRUCT(tup);
+				if (!attrForm->attnotnull)
+					ereport(ERROR,
+							errmsg("column \"%s\" of table \"%s\" is not marked NOT NULL",
+								   strVal(column), get_rel_name(childrelid)));
+				ReleaseSysCache(tup);
+			}
+		}
 
 		/* 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;
 
-- 
2.39.5