From d9f7d81cf7137cb978e07d161ca0d01ebbfafca4 Mon Sep 17 00:00:00 2001
From: Ashutosh Bapat <ashutosh.bapat@enterprisedb.com>
Date: Mon, 18 Dec 2023 16:39:30 +0530
Subject: [PATCH 11/14] DROP IDENTITY on partitioned table recurses to its
 partitions

Since identity property is inherited by partition tables, dropping
identity of a column in partitioned table should result in dropping it
from all the partitions.

I tried to implement the DROP SEQUENCE as a separate step to be executed
at the end of ALTER TABLE command. But that does not work since the
sequence has a dependency on the column (whose identity is being
dropped) and not on the identity property itself. There is no step/ALTER
TABLE command to drop that dependecy.

Ashutosh Bapat
---
 src/backend/commands/tablecmds.c       | 68 +++++++++++++++++++++-----
 src/test/regress/expected/identity.out | 26 ++++++++++
 src/test/regress/sql/identity.sql      | 15 ++++++
 3 files changed, 96 insertions(+), 13 deletions(-)

diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index cf521a65df..79af2f4e85 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -456,7 +456,9 @@ static ObjectAddress ATExecAddIdentity(Relation rel, const char *colName,
 									   bool recurse, bool recursing);
 static ObjectAddress ATExecSetIdentity(Relation rel, const char *colName,
 									   Node *def, LOCKMODE lockmode);
-static ObjectAddress ATExecDropIdentity(Relation rel, const char *colName, bool missing_ok, LOCKMODE lockmode);
+static ObjectAddress ATExecDropIdentity(Relation rel, const char *colName,
+										bool missing_ok, LOCKMODE lockmode,
+										bool recurse, bool recursing);
 static void ATPrepDropExpression(Relation rel, AlterTableCmd *cmd, bool recurse, bool recursing, LOCKMODE lockmode);
 static ObjectAddress ATExecDropExpression(Relation rel, const char *colName, bool missing_ok, LOCKMODE lockmode);
 static ObjectAddress ATExecSetStatistics(Relation rel, const char *colName, int16 colNum,
@@ -4837,7 +4839,9 @@ ATPrepCmd(List **wqueue, Relation rel, AlterTableCmd *cmd,
 			break;
 		case AT_DropIdentity:
 			ATSimplePermissions(cmd->subtype, rel, ATT_TABLE | ATT_VIEW | ATT_FOREIGN_TABLE);
-			/* This command never recurses */
+			/* Set up recursion for phase 2; no other prep needed */
+			if (recurse)
+				cmd->recurse = true;
 			pass = AT_PASS_DROP;
 			break;
 		case AT_DropNotNull:	/* ALTER COLUMN DROP NOT NULL */
@@ -5235,7 +5239,8 @@ ATExecCmd(List **wqueue, AlteredTableInfo *tab,
 			address = ATExecSetIdentity(rel, cmd->name, cmd->def, lockmode);
 			break;
 		case AT_DropIdentity:
-			address = ATExecDropIdentity(rel, cmd->name, cmd->missing_ok, lockmode);
+			address = ATExecDropIdentity(rel, cmd->name, cmd->missing_ok,
+										 lockmode, cmd->recurse, false);
 			break;
 		case AT_DropNotNull:	/* ALTER COLUMN DROP NOT NULL */
 			address = ATExecDropNotNull(rel, cmd->name, cmd->recurse, lockmode);
@@ -8287,7 +8292,9 @@ ATExecSetIdentity(Relation rel, const char *colName, Node *def, LOCKMODE lockmod
  * Return the address of the affected column.
  */
 static ObjectAddress
-ATExecDropIdentity(Relation rel, const char *colName, bool missing_ok, LOCKMODE lockmode)
+ATExecDropIdentity(Relation rel, const char *colName,
+				   bool missing_ok, LOCKMODE lockmode,
+				   bool recurse, bool recursing)
 {
 	HeapTuple	tuple;
 	Form_pg_attribute attTup;
@@ -8296,6 +8303,14 @@ ATExecDropIdentity(Relation rel, const char *colName, bool missing_ok, LOCKMODE
 	ObjectAddress address;
 	Oid			seqid;
 	ObjectAddress seqaddress;
+	bool		ispartitioned;
+
+	ispartitioned = (rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE);
+	if (ispartitioned && !recurse)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_TABLE_DEFINITION),
+				 errmsg("identity can only be dropped from column \"%s\" across all the partitions of the relation \"%s\"",
+						colName, RelationGetRelationName(rel))));
 
 	attrelation = table_open(AttributeRelationId, RowExclusiveLock);
 	tuple = SearchSysCacheCopyAttName(RelationGetRelid(rel), colName);
@@ -8344,15 +8359,42 @@ ATExecDropIdentity(Relation rel, const char *colName, bool missing_ok, LOCKMODE
 
 	table_close(attrelation, RowExclusiveLock);
 
-	/* drop the internal sequence */
-	seqid = getIdentitySequence(RelationGetRelid(rel), attnum, false);
-	deleteDependencyRecordsForClass(RelationRelationId, seqid,
-									RelationRelationId, DEPENDENCY_INTERNAL);
-	CommandCounterIncrement();
-	seqaddress.classId = RelationRelationId;
-	seqaddress.objectId = seqid;
-	seqaddress.objectSubId = 0;
-	performDeletion(&seqaddress, DROP_RESTRICT, PERFORM_DELETION_INTERNAL);
+	/*
+	 * Recurse to drop the identity from column in partitions. Identity is not
+	 * inherited in regular inheritance children so ignore them.
+	 */
+	if (recurse && ispartitioned)
+	{
+		List	   *children;
+		ListCell   *lc;
+
+		children = find_inheritance_children(RelationGetRelid(rel),
+											 lockmode);
+
+		foreach(lc, children)
+		{
+			Relation	childrel;
+
+			childrel = table_open(lfirst_oid(lc), NoLock);
+			ATExecDropIdentity(childrel, colName, false, lockmode, recurse,
+							   true);
+			table_close(childrel, NoLock);
+		}
+	}
+
+	if (!recursing)
+	{
+		/* drop the internal sequence */
+		seqid = getIdentitySequence(RelationGetRelid(rel), attnum, false);
+		deleteDependencyRecordsForClass(RelationRelationId, seqid,
+										RelationRelationId,
+										DEPENDENCY_INTERNAL);
+		CommandCounterIncrement();
+		seqaddress.classId = RelationRelationId;
+		seqaddress.objectId = seqid;
+		seqaddress.objectSubId = 0;
+		performDeletion(&seqaddress, DROP_RESTRICT, PERFORM_DELETION_INTERNAL);
+	}
 
 	return address;
 }
diff --git a/src/test/regress/expected/identity.out b/src/test/regress/expected/identity.out
index 96659d4399..cb783a3f6c 100644
--- a/src/test/regress/expected/identity.out
+++ b/src/test/regress/expected/identity.out
@@ -620,6 +620,32 @@ SELECT tableoid::regclass, f1, f2, f3 FROM itest_parted;
  itest_p1 | 07-05-2016 | from itest_p1     |  4
 (4 rows)
 
+DROP TABLE itest_parted;
+-- changing an identity column to a non-identity column in a partitioned table
+CREATE TABLE itest_parted (f1 date NOT NULL, f2 text, f3 bigint generated always as identity) PARTITION BY RANGE (f1);
+CREATE TABLE itest_p1 PARTITION OF itest_parted FOR VALUES FROM ('2016-07-01') TO ('2016-08-01');
+INSERT into itest_parted(f1, f2) VALUES ('2016-07-2', 'from itest_parted');
+INSERT into itest_p1 (f1, f2) VALUES ('2016-07-3', 'from itest_p1');
+ALTER TABLE ONLY itest_parted ALTER COLUMN f3 DROP IDENTITY; -- fails
+ERROR:  identity can only be dropped from column "f3" across all the partitions of the relation "itest_parted"
+ALTER TABLE itest_parted ALTER COLUMN f3 DROP IDENTITY;
+INSERT into itest_parted(f1, f2) VALUES ('2016-07-4', 'from itest_parted'); -- fails
+ERROR:  null value in column "f3" of relation "itest_p1" violates not-null constraint
+DETAIL:  Failing row contains (07-04-2016, from itest_parted, null).
+INSERT into itest_p1 (f1, f2) VALUES ('2016-07-5', 'from itest_p1'); -- fails
+ERROR:  null value in column "f3" of relation "itest_p1" violates not-null constraint
+DETAIL:  Failing row contains (07-05-2016, from itest_p1, null).
+INSERT into itest_parted(f1, f2, f3) VALUES ('2016-07-4', 'from itest_parted', 3);
+INSERT into itest_p1 (f1, f2, f3) VALUES ('2016-07-5', 'from itest_p1', 4);
+SELECT tableoid::regclass, f1, f2, f3 FROM itest_parted;
+ tableoid |     f1     |        f2         | f3 
+----------+------------+-------------------+----
+ itest_p1 | 07-02-2016 | from itest_parted |  1
+ itest_p1 | 07-03-2016 | from itest_p1     |  2
+ itest_p1 | 07-04-2016 | from itest_parted |  3
+ itest_p1 | 07-05-2016 | from itest_p1     |  4
+(4 rows)
+
 DROP TABLE itest_parted;
 -- partition with identity column of its own is not allowed
 CREATE TABLE itest_parent (f1 date NOT NULL, f2 text, f3 bigint) PARTITION BY RANGE (f1);
diff --git a/src/test/regress/sql/identity.sql b/src/test/regress/sql/identity.sql
index 31615914f5..fa823c2a30 100644
--- a/src/test/regress/sql/identity.sql
+++ b/src/test/regress/sql/identity.sql
@@ -386,6 +386,21 @@ SELECT tableoid::regclass, f1, f2, f3 FROM itest_parted;
 
 DROP TABLE itest_parted;
 
+-- changing an identity column to a non-identity column in a partitioned table
+CREATE TABLE itest_parted (f1 date NOT NULL, f2 text, f3 bigint generated always as identity) PARTITION BY RANGE (f1);
+CREATE TABLE itest_p1 PARTITION OF itest_parted FOR VALUES FROM ('2016-07-01') TO ('2016-08-01');
+INSERT into itest_parted(f1, f2) VALUES ('2016-07-2', 'from itest_parted');
+INSERT into itest_p1 (f1, f2) VALUES ('2016-07-3', 'from itest_p1');
+ALTER TABLE ONLY itest_parted ALTER COLUMN f3 DROP IDENTITY; -- fails
+ALTER TABLE itest_parted ALTER COLUMN f3 DROP IDENTITY;
+INSERT into itest_parted(f1, f2) VALUES ('2016-07-4', 'from itest_parted'); -- fails
+INSERT into itest_p1 (f1, f2) VALUES ('2016-07-5', 'from itest_p1'); -- fails
+INSERT into itest_parted(f1, f2, f3) VALUES ('2016-07-4', 'from itest_parted', 3);
+INSERT into itest_p1 (f1, f2, f3) VALUES ('2016-07-5', 'from itest_p1', 4);
+SELECT tableoid::regclass, f1, f2, f3 FROM itest_parted;
+
+DROP TABLE itest_parted;
+
 -- partition with identity column of its own is not allowed
 CREATE TABLE itest_parent (f1 date NOT NULL, f2 text, f3 bigint) PARTITION BY RANGE (f1);
 CREATE TABLE itest_child PARTITION OF itest_parent (
-- 
2.25.1

