From e40489264b2f083808833a56c706e1762523ec6f Mon Sep 17 00:00:00 2001
From: Ashutosh Bapat <ashutosh.bapat@enterprisedb.com>
Date: Tue, 5 Dec 2023 17:08:11 +0530
Subject: [PATCH 02/14] A newly created partition inherits indentity property

The partitions of a partitioned table are integral part of the
partitioned table. A partition inherits identity columns from the
partitioned table. An indentity column of a partition shares the
identity space with the corresponding column of the partitioned table.
This is effected by sharing the same underlying sequence.  When
INSERTing directly into a partition, sequence associated with the
topmost partitioned table is used to calculate the value of the
corresponding identity column.

Children participating in a regular inheritance, however, may enjoy
identity properties and spaces different from their parent since they
are treated as independent objects.

Ashutosh Bapat
---
 src/backend/commands/tablecmds.c       |  8 ++++++++
 src/backend/rewrite/rewriteHandler.c   | 19 ++++++++++++++++++-
 src/test/regress/expected/identity.out | 16 +++++++++++++++-
 src/test/regress/sql/identity.sql      | 11 ++++++++++-
 4 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 7206da7c53..3e3cd2ed75 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -2852,6 +2852,14 @@ MergeAttributes(List *columns, const List *supers, char relpersistence,
 					def->is_not_null = true;
 				def->storage = attribute->attstorage;
 				def->generated = attribute->attgenerated;
+
+				/*
+				 * partitions inherit identity property from the parent but
+				 * inheritance child does not.
+				 */
+				if (is_partition)
+					def->identity = attribute->attidentity;
+
 				if (CompressionMethodIsValid(attribute->attcompression))
 					def->compression =
 						pstrdup(GetCompressionMethodName(attribute->attcompression));
diff --git a/src/backend/rewrite/rewriteHandler.c b/src/backend/rewrite/rewriteHandler.c
index 41a362310a..d5f1993665 100644
--- a/src/backend/rewrite/rewriteHandler.c
+++ b/src/backend/rewrite/rewriteHandler.c
@@ -25,6 +25,7 @@
 #include "access/table.h"
 #include "catalog/dependency.h"
 #include "catalog/pg_type.h"
+#include "catalog/partition.h"
 #include "commands/trigger.h"
 #include "executor/executor.h"
 #include "foreign/fdwapi.h"
@@ -1234,8 +1235,24 @@ build_column_default(Relation rel, int attrno)
 	if (att_tup->attidentity)
 	{
 		NextValueExpr *nve = makeNode(NextValueExpr);
+		Oid			reloid;
 
-		nve->seqid = getIdentitySequence(RelationGetRelid(rel), attrno, false);
+		/*
+		 * The identity sequence is associated with the topmost partitioned
+		 * table.
+		 */
+		if (rel->rd_rel->relispartition)
+		{
+			List	   *ancestors =
+				get_partition_ancestors(RelationGetRelid(rel));
+
+			reloid = llast_oid(ancestors);
+			list_free(ancestors);
+		}
+		else
+			reloid = RelationGetRelid(rel);
+
+		nve->seqid = getIdentitySequence(reloid, attrno, false);
 		nve->typeId = att_tup->atttypid;
 
 		return (Node *) nve;
diff --git a/src/test/regress/expected/identity.out b/src/test/regress/expected/identity.out
index 7c6e87e8a5..eff73add59 100644
--- a/src/test/regress/expected/identity.out
+++ b/src/test/regress/expected/identity.out
@@ -539,7 +539,21 @@ CREATE TYPE itest_type AS (f1 integer, f2 text, f3 bigint);
 CREATE TABLE itest12 OF itest_type (f1 WITH OPTIONS GENERATED ALWAYS AS IDENTITY); -- error
 ERROR:  identity columns are not supported on typed tables
 DROP TYPE itest_type CASCADE;
--- table partitions (currently not supported)
+-- table partitions
+-- a newly created partition inherits identity and shares the sequence
+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');
+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
+(2 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);
 CREATE TABLE itest_child PARTITION OF itest_parent (
     f3 WITH OPTIONS GENERATED ALWAYS AS IDENTITY
diff --git a/src/test/regress/sql/identity.sql b/src/test/regress/sql/identity.sql
index 9b8db2e4a3..ad2a4249b9 100644
--- a/src/test/regress/sql/identity.sql
+++ b/src/test/regress/sql/identity.sql
@@ -331,8 +331,17 @@ CREATE TABLE itest12 OF itest_type (f1 WITH OPTIONS GENERATED ALWAYS AS IDENTITY
 DROP TYPE itest_type CASCADE;
 
 
--- table partitions (currently not supported)
+-- table partitions
 
+-- a newly created partition inherits identity and shares the sequence
+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');
+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 (
     f3 WITH OPTIONS GENERATED ALWAYS AS IDENTITY
-- 
2.25.1

