From d2b9eaf168d7a168671e4f164230eaa8d4b4965d 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/27] Support adding indentity column to a partitioned table

Adding an identity column to a partitioned table just means propagating
it 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       | 10 ++++++++--
 src/test/regress/expected/identity.out | 22 ++++++++++++++++++++++
 src/test/regress/sql/identity.sql      | 13 +++++++++++++
 3 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 5719df2b76..527a106c48 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -7091,11 +7091,17 @@ 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 are independent enough not to inherit the
+	 * identity column from parent hence can not recursively add  identity
+	 * column if the table has inheritance children.
+	 *
+	 * Partitions, on the other hand, are integral part of a partitioned table
+	 * and inherit indetity column. Hence propagate identity column down the
+	 * partition 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 222b3684da..93d9a60b03 100644
--- a/src/test/regress/expected/identity.out
+++ b/src/test/regress/expected/identity.out
@@ -565,6 +565,28 @@ SELECT tableoid::regclass, f1, f2, f3 FROM pitest1;
  pitest1_p2 | 08-04-2016 | from pitest1     |   4
 (5 rows)
 
+-- add identity column
+CREATE TABLE pitest2 (f1 date NOT NULL, f2 text) PARTITION BY RANGE (f1);
+CREATE TABLE pitest2_p1 PARTITION OF pitest2 FOR VALUES FROM ('2016-07-01') TO ('2016-08-01');
+CREATE TABLE pitest2_p2 PARTITION OF pitest2 FOR VALUES FROM ('2016-08-01') TO ('2016-09-01');
+INSERT into pitest2(f1, f2) VALUES ('2016-07-2', 'from pitest2');
+INSERT INTO pitest2 (f1, f2) VALUES ('2016-08-2', 'from pitest2');
+ALTER TABLE pitest2 ADD COLUMN f3 int GENERATED ALWAYS AS IDENTITY;
+INSERT into pitest2_p1 (f1, f2) VALUES ('2016-07-3', 'from pitest2_p1');
+INSERT INTO pitest2_p2 (f1, f2) VALUES ('2016-08-3', 'from pitest2_p2');
+INSERT into pitest2(f1, f2) VALUES ('2016-07-4', 'from pitest2');
+INSERT INTO pitest2 (f1, f2) VALUES ('2016-08-4', 'from pitest2');
+SELECT tableoid::regclass, f1, f2, f3 FROM pitest2;
+  tableoid  |     f1     |       f2        | f3 
+------------+------------+-----------------+----
+ pitest2_p1 | 07-02-2016 | from pitest2    |  1
+ pitest2_p1 | 07-03-2016 | from pitest2_p1 |  3
+ pitest2_p1 | 07-04-2016 | from pitest2    |  5
+ pitest2_p2 | 08-02-2016 | from pitest2    |  2
+ pitest2_p2 | 08-03-2016 | from pitest2_p2 |  4
+ pitest2_p2 | 08-04-2016 | from pitest2    |  6
+(6 rows)
+
 -- 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 (
diff --git a/src/test/regress/sql/identity.sql b/src/test/regress/sql/identity.sql
index be29cbf119..ba655002c0 100644
--- a/src/test/regress/sql/identity.sql
+++ b/src/test/regress/sql/identity.sql
@@ -349,6 +349,19 @@ INSERT INTO pitest1_p2 (f1, f2) VALUES ('2016-08-3', 'from pitest1_p2');
 INSERT INTO pitest1 (f1, f2) VALUES ('2016-08-4', 'from pitest1');
 SELECT tableoid::regclass, f1, f2, f3 FROM pitest1;
 
+-- add identity column
+CREATE TABLE pitest2 (f1 date NOT NULL, f2 text) PARTITION BY RANGE (f1);
+CREATE TABLE pitest2_p1 PARTITION OF pitest2 FOR VALUES FROM ('2016-07-01') TO ('2016-08-01');
+CREATE TABLE pitest2_p2 PARTITION OF pitest2 FOR VALUES FROM ('2016-08-01') TO ('2016-09-01');
+INSERT into pitest2(f1, f2) VALUES ('2016-07-2', 'from pitest2');
+INSERT INTO pitest2 (f1, f2) VALUES ('2016-08-2', 'from pitest2');
+ALTER TABLE pitest2 ADD COLUMN f3 int GENERATED ALWAYS AS IDENTITY;
+INSERT into pitest2_p1 (f1, f2) VALUES ('2016-07-3', 'from pitest2_p1');
+INSERT INTO pitest2_p2 (f1, f2) VALUES ('2016-08-3', 'from pitest2_p2');
+INSERT into pitest2(f1, f2) VALUES ('2016-07-4', 'from pitest2');
+INSERT INTO pitest2 (f1, f2) VALUES ('2016-08-4', 'from pitest2');
+SELECT tableoid::regclass, f1, f2, f3 FROM pitest2;
+
 -- 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

