From afb98ba63161f9e7b88aaff4b9d21d7ca64c8244 Mon Sep 17 00:00:00 2001
From: "tender.wang" <tender.wang@openpie.com>
Date: Wed, 27 Mar 2024 22:00:17 +0800
Subject: [PATCH v4] Fix pg_attribute attnotnull not reset when dropped PK
 indirectly.

Since b0e96f3119, we will reset pg_attribute attnotnull when drop
pk constraint and not difine NOT NULL. But in some cases, PK will
be dropped indirectly, which will not call tablecmds.c codes.

The result is \d will show not null, but we can't drop NOT NULL.
Moving some code from dropconstraint_internal to RemoveConstraintById
can fix this issue.

We should add SetAttrNotNull sub-command for some Alter Table. Because
those DDLs will check attnotnull is true during DDL.
---
 src/backend/catalog/pg_constraint.c       | 105 +++++++++++++++++++++-
 src/backend/commands/tablecmds.c          |  13 +--
 src/test/regress/expected/constraints.out |   6 ++
 src/test/regress/sql/constraints.sql      |   7 ++
 4 files changed, 125 insertions(+), 6 deletions(-)

diff --git a/src/backend/catalog/pg_constraint.c b/src/backend/catalog/pg_constraint.c
index 604280d322..f16f21985d 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"
@@ -906,7 +907,6 @@ RelationGetNotNullConstraints(Oid relid, bool cooked)
 	return notnulls;
 }
 
-
 /*
  * Delete a single constraint record.
  */
@@ -916,6 +916,7 @@ RemoveConstraintById(Oid conId)
 	Relation	conDesc;
 	HeapTuple	tup;
 	Form_pg_constraint con;
+	List	*unconstrained_cols = NIL;
 
 	conDesc = table_open(ConstraintRelationId, RowExclusiveLock);
 
@@ -967,6 +968,30 @@ RemoveConstraintById(Oid conId)
 
 			table_close(pgrel, RowExclusiveLock);
 		}
+		else if (con->contype == CONSTRAINT_PRIMARY)
+		{
+			Datum adatum;
+			ArrayType *arr;
+			int numkeys;
+			bool isNull;
+			int16 *attnums;
+
+			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 +1011,84 @@ 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 != NIL)
+	{
+		Relation	attrel;
+		Relation	rel;
+		Bitmapset  *ircols;
+		ListCell   *lc;
+
+		/* Make the above deletion visible */
+		CommandCounterIncrement();
+
+		attrel = table_open(AttributeRelationId, RowExclusiveLock);
+		rel = table_open(con->conrelid, AccessExclusiveLock);
+
+		ircols = RelationGetIndexAttrBitmap(rel, INDEX_ATTR_BITMAP_IDENTITY_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(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);
+
+			/* GENERATED AS IDENTITY column need not null */
+			if (attForm->attidentity)
+			{
+				heap_freetuple(atttup);
+				continue;
+			}
+
+			/* A column in the replica identity index need not null */
+			if (bms_is_member(lfirst_int(lc) - FirstLowInvalidHeapAttributeNumber, ircols))
+			{
+				heap_freetuple(atttup);
+				continue;
+			}
+
+			/*
+			 * 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 exists, and reset
+			 * attnotnull if none.
+			 */
+			contup = findNotNullConstraintAttnum(RelationGetRelid(rel), attnum);
+			if (HeapTupleIsValid(contup))
+			{
+				heap_freetuple(contup);
+				heap_freetuple(atttup);
+				continue;
+			}
+
+			/* Reset attnotnull */
+			if (attForm->attnotnull)
+			{
+				attForm->attnotnull = false;
+				CatalogTupleUpdate(attrel, &atttup->t_self, atttup);
+			}
+
+			heap_freetuple(atttup);
+		}
+		table_close(rel, AccessExclusiveLock);
+		table_close(attrel, RowExclusiveLock);
+	}
 	/* Clean up */
 	ReleaseSysCache(tup);
 	table_close(conDesc, RowExclusiveLock);
diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 259b4237a2..57fafb839a 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -147,6 +147,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-set column NOT NULL */
 	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 */
@@ -14629,12 +14630,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..00fab56bf8 100644
--- a/src/test/regress/expected/constraints.out
+++ b/src/test/regress/expected/constraints.out
@@ -867,6 +867,12 @@ select conname, contype, conkey from pg_constraint where conrelid = 'notnull_tbl
 (1 row)
 
 DROP TABLE notnull_tbl1;
+-- NOT NULL mask should be reset when drop PK
+CREATE TABLE notnull_t1(c0 int, c1 int);
+ALTER TABLE  notnull_t1 ADD CONSTRAINT Q PRIMARY KEY(c0, c1);
+ALTER TABLE  notnull_t1 DROP c1;
+ALTER TABLE  notnull_t1  ALTER c0 DROP NOT NULL;
+DROP TABLE notnull_t1;
 -- nope
 CREATE TABLE notnull_tbl2 (a INTEGER CONSTRAINT blah NOT NULL, b INTEGER CONSTRAINT blah NOT NULL);
 ERROR:  constraint "blah" for relation "notnull_tbl2" already exists
diff --git a/src/test/regress/sql/constraints.sql b/src/test/regress/sql/constraints.sql
index 87d685ae39..a146b67416 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;
 
+-- NOT NULL mask should be reset when drop PK
+CREATE TABLE notnull_t1(c0 int, c1 int);
+ALTER TABLE  notnull_t1 ADD CONSTRAINT Q PRIMARY KEY(c0, c1);
+ALTER TABLE  notnull_t1 DROP c1;
+ALTER TABLE  notnull_t1  ALTER c0 DROP NOT NULL;
+DROP TABLE notnull_t1;
+
 -- nope
 CREATE TABLE notnull_tbl2 (a INTEGER CONSTRAINT blah NOT NULL, b INTEGER CONSTRAINT blah NOT NULL);
 
-- 
2.25.1

