From 397d424cd64020bcb4aff219deefca7fc7b6f88d Mon Sep 17 00:00:00 2001
From: Alvaro Herrera <alvherre@alvh.no-ip.org>
Date: Tue, 9 Apr 2024 19:18:57 +0200
Subject: [PATCH] Correctly reset attnotnull when constraints dropped
 indirectly

---
 src/backend/catalog/pg_constraint.c       | 105 +++++++++++++++++++++-
 src/backend/commands/tablecmds.c          |  67 ++++++--------
 src/test/regress/expected/constraints.out |  11 +++
 src/test/regress/sql/constraints.sql      |   7 ++
 4 files changed, 150 insertions(+), 40 deletions(-)

diff --git a/src/backend/catalog/pg_constraint.c b/src/backend/catalog/pg_constraint.c
index 604280d322..17c7a2d0f3 100644
--- a/src/backend/catalog/pg_constraint.c
+++ b/src/backend/catalog/pg_constraint.c
@@ -19,6 +19,7 @@
 #include "access/htup_details.h"
 #include "access/sysattr.h"
 #include "access/table.h"
+#include "access/xact.h"
 #include "catalog/catalog.h"
 #include "catalog/dependency.h"
 #include "catalog/heap.h"
@@ -916,6 +917,8 @@ RemoveConstraintById(Oid conId)
 	Relation	conDesc;
 	HeapTuple	tup;
 	Form_pg_constraint con;
+	bool		dropping_pk = false;
+	List	   *unconstrained_cols = NIL;
 
 	conDesc = table_open(ConstraintRelationId, RowExclusiveLock);
 
@@ -940,7 +943,9 @@ RemoveConstraintById(Oid conId)
 		/*
 		 * We need to update the relchecks count if it is a check constraint
 		 * being dropped.  This update will force backends to rebuild relcache
-		 * entries when we commit.
+		 * entries when we commit.  For not-null and primary key constraints,
+		 * obtain the list of columns affected, so that we can reset their
+		 * attnotnull flags below.
 		 */
 		if (con->contype == CONSTRAINT_CHECK)
 		{
@@ -967,6 +972,36 @@ RemoveConstraintById(Oid conId)
 
 			table_close(pgrel, RowExclusiveLock);
 		}
+		else if (con->contype == CONSTRAINT_NOTNULL)
+		{
+			unconstrained_cols = list_make1_int(extractNotNullColumn(tup));
+		}
+		else if (con->contype == CONSTRAINT_PRIMARY)
+		{
+			Datum		adatum;
+			ArrayType  *arr;
+			int			numkeys;
+			bool		isNull;
+			int16	   *attnums;
+
+			dropping_pk = true;
+
+			adatum = heap_getattr(tup, Anum_pg_constraint_conkey,
+								  RelationGetDescr(conDesc), &isNull);
+			if (isNull)
+				elog(ERROR, "null conkey for constraint %u", con->oid);
+			arr = DatumGetArrayTypeP(adatum);	/* ensure not toasted */
+			numkeys = ARR_DIMS(arr)[0];
+			if (ARR_NDIM(arr) != 1 ||
+				numkeys < 0 ||
+				ARR_HASNULL(arr) ||
+				ARR_ELEMTYPE(arr) != INT2OID)
+				elog(ERROR, "conkey is not a 1-D smallint array");
+			attnums = (int16 *) ARR_DATA_PTR(arr);
+
+			for (int i = 0; i < numkeys; i++)
+				unconstrained_cols = lappend_int(unconstrained_cols, attnums[i]);
+		}
 
 		/* Keep lock on constraint's rel until end of xact */
 		table_close(rel, NoLock);
@@ -986,6 +1021,74 @@ RemoveConstraintById(Oid conId)
 	/* Fry the constraint itself */
 	CatalogTupleDelete(conDesc, &tup->t_self);
 
+	/*
+	 * If this was a NOT NULL or the primary key, the constrained columns must
+	 * have had pg_attribute.attnotnull set.  See if we need to reset it, and
+	 * do so.
+	 */
+	if (unconstrained_cols)
+	{
+		Relation	tablerel;
+		Relation	attrel;
+		Bitmapset  *pkcols;
+		ListCell   *lc;
+
+		/* Make the above deletion visible */
+		CommandCounterIncrement();
+
+		tablerel = table_open(con->conrelid, NoLock);	/* already have lock */
+		attrel = table_open(AttributeRelationId, RowExclusiveLock);
+
+		/*
+		 * We want to test columns for their presence in the primary key, but
+		 * only if we're not dropping it.
+		 */
+		pkcols = dropping_pk ? NULL :
+			RelationGetIndexAttrBitmap(tablerel,
+									   INDEX_ATTR_BITMAP_PRIMARY_KEY);
+
+		foreach(lc, unconstrained_cols)
+		{
+			AttrNumber	attnum = lfirst_int(lc);
+			HeapTuple	atttup;
+			HeapTuple	contup;
+			Form_pg_attribute attForm;
+
+			/*
+			 * Obtain pg_attribute tuple and verify conditions on it.  We use
+			 * a copy we can scribble on.
+			 */
+			atttup = SearchSysCacheCopyAttNum(con->conrelid, attnum);
+			if (!HeapTupleIsValid(atttup))
+				elog(ERROR, "cache lookup failed for attribute %d of relation %u",
+					 attnum, con->conrelid);
+			attForm = (Form_pg_attribute) GETSTRUCT(atttup);
+
+			/*
+			 * Since the above deletion has been made visible, we can now
+			 * search for any remaining constraints on this column (or these
+			 * columns, in the case we're dropping a multicol primary key.)
+			 * Then, verify whether any further NOT NULL or primary key
+			 * exists, and reset attnotnull if none.
+			 */
+			contup = findNotNullConstraintAttnum(con->conrelid, attnum);
+			if (contup ||
+				bms_is_member(attnum - FirstLowInvalidHeapAttributeNumber,
+							  pkcols))
+				continue;
+
+			/* Reset attnotnull */
+			if (attForm->attnotnull)
+			{
+				attForm->attnotnull = false;
+				CatalogTupleUpdate(attrel, &atttup->t_self, atttup);
+			}
+		}
+
+		table_close(attrel, RowExclusiveLock);
+		table_close(tablerel, NoLock);
+	}
+
 	/* Clean up */
 	ReleaseSysCache(tup);
 	table_close(conDesc, RowExclusiveLock);
diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 865c6331c1..25d3613fe4 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -149,6 +149,7 @@ typedef enum AlterTablePass
 	AT_PASS_ALTER_TYPE,			/* ALTER COLUMN TYPE */
 	AT_PASS_ADD_COL,			/* ADD COLUMN */
 	AT_PASS_SET_EXPRESSION,		/* ALTER SET EXPRESSION */
+	AT_PASS_OLD_COL_ATTRS,		/* re-install attnotnull */
 	AT_PASS_OLD_INDEX,			/* re-add existing indexes */
 	AT_PASS_OLD_CONSTR,			/* re-add existing constraints */
 	/* We could support a RENAME COLUMN pass here, but not currently used */
@@ -12933,10 +12934,10 @@ dropconstraint_internal(Relation rel, HeapTuple constraintTup, DropBehavior beha
 	List	   *children;
 	ListCell   *child;
 	bool		is_no_inherit_constraint = false;
-	bool		dropping_pk = false;
 	char	   *constrName;
 	List	   *unconstrained_cols = NIL;
-	char	   *colname;
+	char	   *colname = NULL;
+	bool		dropping_pk = false;
 
 	if (list_member_oid(*readyRels, RelationGetRelid(rel)))
 		return InvalidObjectAddress;
@@ -12968,10 +12969,12 @@ dropconstraint_internal(Relation rel, HeapTuple constraintTup, DropBehavior beha
 	 */
 	if (con->contype == CONSTRAINT_NOTNULL)
 	{
-		AttrNumber	colnum = extractNotNullColumn(constraintTup);
+		AttrNumber	colnum;
 
-		if (colnum != InvalidAttrNumber)
-			unconstrained_cols = list_make1_int(colnum);
+		colnum = extractNotNullColumn(constraintTup);
+		unconstrained_cols = list_make1_int(colnum);
+		colname = NameStr(TupleDescAttr(RelationGetDescr(rel),
+										colnum - 1)->attname);
 	}
 	else if (con->contype == CONSTRAINT_PRIMARY)
 	{
@@ -13027,9 +13030,8 @@ dropconstraint_internal(Relation rel, HeapTuple constraintTup, DropBehavior beha
 	performDeletion(&conobj, behavior, 0);
 
 	/*
-	 * If this was a NOT NULL or the primary key, the constrained columns must
-	 * have had pg_attribute.attnotnull set.  See if we need to reset it, and
-	 * do so.
+	 * If this was a NOT NULL or the primary key, verify that we still have
+	 * constraints to support GENERATED AS IDENTITY or the replica identity.
 	 */
 	if (unconstrained_cols)
 	{
@@ -13038,7 +13040,7 @@ dropconstraint_internal(Relation rel, HeapTuple constraintTup, DropBehavior beha
 		Bitmapset  *ircols;
 		ListCell   *lc;
 
-		/* Make the above deletion visible */
+		/* Make implicit attnotnull changes visible */
 		CommandCounterIncrement();
 
 		attrel = table_open(AttributeRelationId, RowExclusiveLock);
@@ -13058,27 +13060,26 @@ dropconstraint_internal(Relation rel, HeapTuple constraintTup, DropBehavior beha
 			HeapTuple	atttup;
 			HeapTuple	contup;
 			Form_pg_attribute attForm;
+			char		attidentity;
 
 			/*
-			 * Obtain pg_attribute tuple and verify conditions on it.  We use
-			 * a copy we can scribble on.
+			 * Obtain pg_attribute tuple and verify conditions on it.
 			 */
-			atttup = SearchSysCacheCopyAttNum(RelationGetRelid(rel), attnum);
+			atttup = SearchSysCacheAttNum(RelationGetRelid(rel), attnum);
 			if (!HeapTupleIsValid(atttup))
 				elog(ERROR, "cache lookup failed for attribute %d of relation %u",
 					 attnum, RelationGetRelid(rel));
 			attForm = (Form_pg_attribute) GETSTRUCT(atttup);
+			attidentity = attForm->attidentity;
+			ReleaseSysCache(atttup);
 
 			/*
 			 * Since the above deletion has been made visible, we can now
 			 * search for any remaining constraints on this column (or these
 			 * columns, in the case we're dropping a multicol primary key.)
 			 * Then, verify whether any further NOT NULL or primary key
-			 * exists, and reset attnotnull if none.
-			 *
-			 * However, if this is a generated identity column, abort the
-			 * whole thing with a specific error message, because the
-			 * constraint is required in that case.
+			 * exists: if none and this is a generated identity column or the
+			 * replica identity, abort the whole thing.
 			 */
 			contup = findNotNullConstraintAttnum(RelationGetRelid(rel), attnum);
 			if (contup ||
@@ -13090,7 +13091,7 @@ dropconstraint_internal(Relation rel, HeapTuple constraintTup, DropBehavior beha
 			 * It's not valid to drop the not-null constraint for a GENERATED
 			 * AS IDENTITY column.
 			 */
-			if (attForm->attidentity)
+			if (attidentity != '\0')
 				ereport(ERROR,
 						errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
 						errmsg("column \"%s\" of relation \"%s\" is an identity column",
@@ -13107,13 +13108,6 @@ dropconstraint_internal(Relation rel, HeapTuple constraintTup, DropBehavior beha
 						errcode(ERRCODE_INVALID_TABLE_DEFINITION),
 						errmsg("column \"%s\" is in index used as replica identity",
 							   get_attname(RelationGetRelid(rel), lfirst_int(lc), false)));
-
-			/* Reset attnotnull */
-			if (attForm->attnotnull)
-			{
-				attForm->attnotnull = false;
-				CatalogTupleUpdate(attrel, &atttup->t_self, atttup);
-			}
 		}
 		table_close(attrel, RowExclusiveLock);
 	}
@@ -13152,13 +13146,6 @@ dropconstraint_internal(Relation rel, HeapTuple constraintTup, DropBehavior beha
 				 errmsg("cannot remove constraint from only the partitioned table when partitions exist"),
 				 errhint("Do not specify the ONLY keyword.")));
 
-	/* For not-null constraints we recurse by column name */
-	if (con->contype == CONSTRAINT_NOTNULL)
-		colname = NameStr(TupleDescAttr(RelationGetDescr(rel),
-										linitial_int(unconstrained_cols) - 1)->attname);
-	else
-		colname = NULL;			/* keep compiler quiet */
-
 	foreach(child, children)
 	{
 		Oid			childrelid = lfirst_oid(child);
@@ -13174,8 +13161,8 @@ dropconstraint_internal(Relation rel, HeapTuple constraintTup, DropBehavior beha
 		CheckTableNotInUse(childrel, "ALTER TABLE");
 
 		/*
-		 * We search for not-null constraint by column number, and other
-		 * constraints by name.
+		 * We search for not-null constraints by column name, and others by
+		 * constraint name.
 		 */
 		if (con->contype == CONSTRAINT_NOTNULL)
 		{
@@ -14656,12 +14643,14 @@ ATPostAlterTypeParse(Oid oldId, Oid oldRelId, Oid refRelId, char *cmd,
 				else if (cmd->subtype == AT_SetAttNotNull)
 				{
 					/*
-					 * The parser will create AT_AttSetNotNull subcommands for
-					 * columns of PRIMARY KEY indexes/constraints, but we need
-					 * not do anything with them here, because the columns'
-					 * NOT NULL marks will already have been propagated into
-					 * the new table definition.
+					 * When doing alter primary key column type, Removing old
+					 * PK constraint as sub-command will reset attnotnull to
+					 * false in RemoveConstraintById(). We must set attnotnull
+					 * back to true before doing re-add the index sub-command.
 					 */
+					cmd->subtype = AT_SetAttNotNull;
+					tab->subcmds[AT_PASS_OLD_COL_ATTRS] =
+						lappend(tab->subcmds[AT_PASS_OLD_COL_ATTRS], cmd);
 				}
 				else
 					elog(ERROR, "unexpected statement subtype: %d",
diff --git a/src/test/regress/expected/constraints.out b/src/test/regress/expected/constraints.out
index d9d8408e86..ce8ba73a87 100644
--- a/src/test/regress/expected/constraints.out
+++ b/src/test/regress/expected/constraints.out
@@ -866,6 +866,17 @@ select conname, contype, conkey from pg_constraint where conrelid = 'notnull_tbl
  foobar  | n       | {1}
 (1 row)
 
+DROP TABLE notnull_tbl1;
+-- make sure attnotnull is reset correctly when a PK is dropped indirectly
+CREATE TABLE notnull_tbl1 (c0 int, c1 int);
+ALTER TABLE  notnull_tbl1 ADD CONSTRAINT Q PRIMARY KEY(c0, c1);
+ALTER TABLE  notnull_tbl1 DROP c1;
+\d+ notnull_tbl1
+                               Table "public.notnull_tbl1"
+ Column |  Type   | Collation | Nullable | Default | Storage | Stats target | Description 
+--------+---------+-----------+----------+---------+---------+--------------+-------------
+ c0     | integer |           |          |         | plain   |              | 
+
 DROP TABLE notnull_tbl1;
 -- nope
 CREATE TABLE notnull_tbl2 (a INTEGER CONSTRAINT blah NOT NULL, b INTEGER CONSTRAINT blah NOT NULL);
diff --git a/src/test/regress/sql/constraints.sql b/src/test/regress/sql/constraints.sql
index 87d685ae39..b25107da46 100644
--- a/src/test/regress/sql/constraints.sql
+++ b/src/test/regress/sql/constraints.sql
@@ -599,6 +599,13 @@ ALTER TABLE notnull_tbl1 ADD CONSTRAINT foobar NOT NULL a;
 select conname, contype, conkey from pg_constraint where conrelid = 'notnull_tbl1'::regclass;
 DROP TABLE notnull_tbl1;
 
+-- make sure attnotnull is reset correctly when a PK is dropped indirectly
+CREATE TABLE notnull_tbl1 (c0 int, c1 int);
+ALTER TABLE  notnull_tbl1 ADD CONSTRAINT Q PRIMARY KEY(c0, c1);
+ALTER TABLE  notnull_tbl1 DROP c1;
+\d+ notnull_tbl1
+DROP TABLE notnull_tbl1;
+
 -- nope
 CREATE TABLE notnull_tbl2 (a INTEGER CONSTRAINT blah NOT NULL, b INTEGER CONSTRAINT blah NOT NULL);
 
-- 
2.39.2

