From 546b7a154db10edae4d63fbd3935be6b1342c1e5 Mon Sep 17 00:00:00 2001
From: Ashutosh Bapat <ashutosh.bapat@enterprisedb.com>
Date: Mon, 11 Dec 2023 16:41:09 +0530
Subject: [PATCH 09/14] Adding identity to partitioned table adds it to all
 partitions

... additionally do not allow adding identity to only partitioned table.

Ashutosh Bapat
---
 src/backend/commands/tablecmds.c       | 42 +++++++++++++++++++++++---
 src/test/regress/expected/identity.out | 25 +++++++++++++++
 src/test/regress/sql/identity.sql      | 17 +++++++++++
 3 files changed, 80 insertions(+), 4 deletions(-)

diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 7b94d7f434..cf521a65df 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -452,7 +452,8 @@ static ObjectAddress ATExecColumnDefault(Relation rel, const char *colName,
 static ObjectAddress ATExecCookedColumnDefault(Relation rel, AttrNumber attnum,
 											   Node *newDefault);
 static ObjectAddress ATExecAddIdentity(Relation rel, const char *colName,
-									   Node *def, LOCKMODE lockmode);
+									   Node *def, LOCKMODE lockmode,
+									   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);
@@ -4823,7 +4824,9 @@ ATPrepCmd(List **wqueue, Relation rel, AlterTableCmd *cmd,
 			break;
 		case AT_AddIdentity:
 			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_ADD_OTHERCONSTR;
 			break;
 		case AT_SetIdentity:
@@ -5222,7 +5225,8 @@ ATExecCmd(List **wqueue, AlteredTableInfo *tab,
 			cmd = ATParseTransformCmd(wqueue, tab, rel, cmd, false, lockmode,
 									  cur_pass, context);
 			Assert(cmd != NULL);
-			address = ATExecAddIdentity(rel, cmd->name, cmd->def, lockmode);
+			address = ATExecAddIdentity(rel, cmd->name, cmd->def, lockmode,
+										cmd->recurse, false);
 			break;
 		case AT_SetIdentity:
 			cmd = ATParseTransformCmd(wqueue, tab, rel, cmd, false, lockmode,
@@ -8100,7 +8104,7 @@ ATExecCookedColumnDefault(Relation rel, AttrNumber attnum,
  */
 static ObjectAddress
 ATExecAddIdentity(Relation rel, const char *colName,
-				  Node *def, LOCKMODE lockmode)
+				  Node *def, LOCKMODE lockmode, bool recurse, bool recursing)
 {
 	Relation	attrelation;
 	HeapTuple	tuple;
@@ -8108,6 +8112,14 @@ ATExecAddIdentity(Relation rel, const char *colName,
 	AttrNumber	attnum;
 	ObjectAddress address;
 	ColumnDef  *cdef = castNode(ColumnDef, def);
+	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 added to column \"%s\" across all the partitions of the relation \"%s\"",
+						colName, RelationGetRelationName(rel))));
 
 	attrelation = table_open(AttributeRelationId, RowExclusiveLock);
 
@@ -8162,6 +8174,28 @@ ATExecAddIdentity(Relation rel, const char *colName,
 
 	table_close(attrelation, RowExclusiveLock);
 
+	/*
+	 * Recurse to propagate the identity column to partitions. Identity is not
+	 * inherited in regular inheritance children.
+	 */
+	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);
+			ATExecAddIdentity(childrel, colName, def, lockmode, recurse, true);
+			table_close(childrel, NoLock);
+		}
+	}
+
 	return address;
 }
 
diff --git a/src/test/regress/expected/identity.out b/src/test/regress/expected/identity.out
index 9984f7ed19..96659d4399 100644
--- a/src/test/regress/expected/identity.out
+++ b/src/test/regress/expected/identity.out
@@ -595,6 +595,31 @@ SELECT tableoid::regclass, f1, f2, f3 FROM itest_parted;
  itest_p2 | 08-04-2016 | from itest_parted |  6
 (6 rows)
 
+DROP TABLE itest_parted;
+-- changing a regular column to identity column in a partitioned table
+CREATE TABLE itest_parted (f1 date NOT NULL, f2 text, f3 int) 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 VALUES ('2016-07-2', 'from itest_parted', 1);
+INSERT into itest_p1 VALUES ('2016-07-3', 'from itest_p1', 2);
+ALTER TABLE ONLY itest_parted -- ONLY causes the command to fail
+            ALTER COLUMN f3 SET NOT NULL,
+            ALTER COLUMN f3 ADD generated always as identity (start with 3);
+ERROR:  constraint must be added to child tables too
+HINT:  Do not specify the ONLY keyword.
+ALTER TABLE itest_parted
+            ALTER COLUMN f3 SET NOT NULL, -- avoids error
+            ALTER COLUMN f3 ADD generated always as identity (start with 3);
+INSERT into itest_parted(f1, f2) VALUES ('2016-07-4', 'from itest_parted');
+INSERT into itest_p1 (f1, f2) VALUES ('2016-07-5', 'from itest_p1');
+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 f23ccd2a86..31615914f5 100644
--- a/src/test/regress/sql/identity.sql
+++ b/src/test/regress/sql/identity.sql
@@ -369,6 +369,23 @@ SELECT tableoid::regclass, f1, f2, f3 FROM itest_parted;
 
 DROP TABLE itest_parted;
 
+-- changing a regular column to identity column in a partitioned table
+CREATE TABLE itest_parted (f1 date NOT NULL, f2 text, f3 int) 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 VALUES ('2016-07-2', 'from itest_parted', 1);
+INSERT into itest_p1 VALUES ('2016-07-3', 'from itest_p1', 2);
+ALTER TABLE ONLY itest_parted -- ONLY causes the command to fail
+            ALTER COLUMN f3 SET NOT NULL,
+            ALTER COLUMN f3 ADD generated always as identity (start with 3);
+ALTER TABLE itest_parted
+            ALTER COLUMN f3 SET NOT NULL, -- avoids error
+            ALTER COLUMN f3 ADD generated always as identity (start with 3);
+INSERT into itest_parted(f1, f2) VALUES ('2016-07-4', 'from itest_parted');
+INSERT into itest_p1 (f1, f2) VALUES ('2016-07-5', 'from itest_p1');
+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

