From c8256786b5db282f5ad0cb20b8b75dbf399da35f Mon Sep 17 00:00:00 2001
From: Ashutosh Bapat <ashutosh.bapat@enterprisedb.com>
Date: Thu, 7 Dec 2023 11:20:07 +0530
Subject: [PATCH 06/14] Support adding indentity column to a partitioned table

All partitions in a partitioned table share the same identity space for an
identity column. Adding an identity column to a partitioned table just means
propagating the identity property along with the column down the partitioning
hierarchy. All the partitions share the same underlying sequence.

As for the implementation, we need to just disable the code prohibiting this
case for partitioned tables. The column definition is transformed only once.
The statements related to identity sequence are executed only once before
executing any subcommands. Thus only one sequence, associated with the identity
column, is created. The transformed column definition and related commands are
copied as they are for all the partitions. Hence default expressions of all the
partition use the same sequence when rewriting the table. Also the NOT NULL
constraints required by the identity column are propagated.

Ashutosh Bapat
---
 src/backend/commands/tablecmds.c       |  7 +++++--
 src/test/regress/expected/identity.out | 24 ++++++++++++++++++++++++
 src/test/regress/sql/identity.sql      | 16 ++++++++++++++++
 3 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index b298747e13..b471cb4d6b 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -7087,11 +7087,14 @@ ATExecAddColumn(List **wqueue, AlteredTableInfo *tab, Relation rel,
 	}
 
 	/*
-	 * Cannot add identity column if table has children, because identity does
-	 * not inherit.  (Adding column and identity separately will work.)
+	 * Regular inheritance children do not inherit the identity hence prohibit
+	 * adding identity column if the table has inheritance children. But all
+	 * the partitions in a partition hierarchy share identity space so
+	 * propagate it down the partitioning hierarchy.
 	 */
 	if (colDef->identity &&
 		recurse &&
+		rel->rd_rel->relkind != RELKIND_PARTITIONED_TABLE &&
 		find_inheritance_children(myrelid, NoLock) != NIL)
 		ereport(ERROR,
 				(errcode(ERRCODE_INVALID_TABLE_DEFINITION),
diff --git a/src/test/regress/expected/identity.out b/src/test/regress/expected/identity.out
index 6b9edb22c1..9984f7ed19 100644
--- a/src/test/regress/expected/identity.out
+++ b/src/test/regress/expected/identity.out
@@ -571,6 +571,30 @@ SELECT tableoid::regclass, f1, f2, f3 FROM itest_parted;
  itest_p2 | 08-04-2016 | from itest_parted |   4
 (5 rows)
 
+DROP TABLE itest_parted;
+-- adding an identity column to a partitioned table adds it to all the
+-- partitions
+CREATE TABLE itest_parted (f1 date NOT NULL, f2 text) PARTITION BY RANGE (f1);
+CREATE TABLE itest_p1 PARTITION OF itest_parted FOR VALUES FROM ('2016-07-01') TO ('2016-08-01');
+CREATE TABLE itest_p2 PARTITION OF itest_parted FOR VALUES FROM ('2016-08-01') TO ('2016-09-01');
+INSERT into itest_parted(f1, f2) VALUES ('2016-07-2', 'from itest_parted');
+INSERT INTO itest_parted (f1, f2) VALUES ('2016-08-2', 'from itest_parted');
+ALTER TABLE itest_parted ADD COLUMN f3 int generated always as identity;
+INSERT into itest_p1 (f1, f2) VALUES ('2016-07-3', 'from itest_p1');
+INSERT INTO itest_p2 (f1, f2) VALUES ('2016-08-3', 'from itest_p2');
+INSERT into itest_parted(f1, f2) VALUES ('2016-07-4', 'from itest_parted');
+INSERT INTO itest_parted (f1, f2) VALUES ('2016-08-4', 'from itest_parted');
+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     |  3
+ itest_p1 | 07-04-2016 | from itest_parted |  5
+ itest_p2 | 08-02-2016 | from itest_parted |  2
+ itest_p2 | 08-03-2016 | from itest_p2     |  4
+ itest_p2 | 08-04-2016 | from itest_parted |  6
+(6 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 a9babd1baa..f23ccd2a86 100644
--- a/src/test/regress/sql/identity.sql
+++ b/src/test/regress/sql/identity.sql
@@ -353,6 +353,22 @@ SELECT tableoid::regclass, f1, f2, f3 FROM itest_parted;
 
 DROP TABLE itest_parted;
 
+-- adding an identity column to a partitioned table adds it to all the
+-- partitions
+CREATE TABLE itest_parted (f1 date NOT NULL, f2 text) PARTITION BY RANGE (f1);
+CREATE TABLE itest_p1 PARTITION OF itest_parted FOR VALUES FROM ('2016-07-01') TO ('2016-08-01');
+CREATE TABLE itest_p2 PARTITION OF itest_parted FOR VALUES FROM ('2016-08-01') TO ('2016-09-01');
+INSERT into itest_parted(f1, f2) VALUES ('2016-07-2', 'from itest_parted');
+INSERT INTO itest_parted (f1, f2) VALUES ('2016-08-2', 'from itest_parted');
+ALTER TABLE itest_parted ADD COLUMN f3 int generated always as identity;
+INSERT into itest_p1 (f1, f2) VALUES ('2016-07-3', 'from itest_p1');
+INSERT INTO itest_p2 (f1, f2) VALUES ('2016-08-3', 'from itest_p2');
+INSERT into itest_parted(f1, f2) VALUES ('2016-07-4', 'from itest_parted');
+INSERT INTO itest_parted (f1, f2) VALUES ('2016-08-4', 'from itest_parted');
+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

