From 215e8929bc8a61998047a8c060b4a0fd95c242f2 Mon Sep 17 00:00:00 2001
From: jian he <jian.universality@gmail.com>
Date: Tue, 31 Dec 2024 10:15:27 +0800
Subject: [PATCH v1 1/1] disallow attach partition if referencing another
 partitioned table

discussion: https://postgr.es/m/18741-e4ef6f7aa8a956cb@postgresql.org
---
 src/backend/commands/tablecmds.c          | 69 +++++++++++++++++++++--
 src/test/regress/expected/foreign_key.out | 20 +++++++
 src/test/regress/sql/foreign_key.sql      | 21 +++++++
 3 files changed, 104 insertions(+), 6 deletions(-)

diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 4937478262..314a1afd1d 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -11090,14 +11090,16 @@ CloneFkReferencing(List **wqueue, Relation parentRel, Relation partRel)
 	AttrMap    *attmap;
 	List	   *partFKs;
 	List	   *clone = NIL;
+	List		*fkpart	= NIL;
+	List		*fkparentRel = NIL;
 	ListCell   *cell;
 	Relation	trigrel;
 
 	/* obtain a list of constraints that we need to clone */
-	foreach(cell, RelationGetFKeyList(parentRel))
-	{
-		ForeignKeyCacheInfo *fk = lfirst(cell);
+	fkparentRel = RelationGetFKeyList(parentRel);
 
+	foreach_node(ForeignKeyCacheInfo, fk, fkparentRel)
+	{
 		/*
 		 * Refuse to attach a table as partition that this partitioned table
 		 * already has a foreign key to.  This isn't useful schema, which is
@@ -11118,6 +11120,60 @@ CloneFkReferencing(List **wqueue, Relation parentRel, Relation partRel)
 		clone = lappend_oid(clone, fk->conoid);
 	}
 
+	/*
+	 * we cannot attach a table as a partition if it is still referencing
+	 * another partitioned table. However, this restriction does not apply if
+	 * the referenced table is the partitioned table itself (i.e., `partRel`
+	 * referencing `parentRel`) or if both the table and the partitioned table
+	 * reference the same another partitioned table.
+	 *
+	 * so the following example should be OK.
+	 * CREATE TABLE pt2 (a int primary key) PARTITION BY RANGE (a);
+	 * CREATE TABLE p2 (a int not null, FOREIGN KEY (a) REFERENCES pt2);
+	 * ALTER TABLE pt2 ATTACH PARTITION p2 FOR VALUES FROM (0) TO (1);
+	 *
+	 * if the table is referencing non-partitioned table, that's OK.
+	*/
+	fkpart = RelationGetFKeyList(partRel);
+	foreach_node(ForeignKeyCacheInfo, fk_attrel, fkpart)
+	{
+		bool	skip = false;
+		foreach_node(ForeignKeyCacheInfo, fk, fkparentRel)
+		{
+			if (fk_attrel->confrelid == fk->confrelid)
+			{
+				skip = true;
+				break;
+			}
+		}
+
+		/*
+		 * if the attachee referencing a partitioned table and that partitioned table root
+		 * table is same as the attachee, then we don't error out.
+		*/
+		if (!skip && get_rel_relkind(fk_attrel->confrelid) == RELKIND_PARTITIONED_TABLE)
+		{
+			List	   *ancestors;
+			Oid			rootrelid;
+
+			ancestors = get_partition_ancestors(fk_attrel->confrelid);
+			if (ancestors == NIL)
+				rootrelid = fk_attrel->confrelid;
+			else
+				rootrelid = llast_oid(ancestors);
+
+			if (ancestors != NIL)
+				list_free(ancestors);
+
+			if (RelationGetRelid(parentRel) != rootrelid)
+				ereport(ERROR,
+						errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+						errmsg("table \"%s\" is still referenced by another partitioned table",
+								RelationGetRelationName(partRel)),
+						errdetail("table referenced by another partitioned table can not be attach as a partition."));
+		}
+	}
+
 	/*
 	 * Silently do nothing if there's nothing to do.  In particular, this
 	 * avoids throwing a spurious error for foreign tables.
@@ -19890,9 +19946,10 @@ DetachPartitionFinalize(Relation rel, Relation partRel, bool concurrent,
 		GetForeignKeyCheckTriggers(trigrel,
 								   fk->conoid, fk->confrelid, fk->conrelid,
 								   &insertTriggerOid, &updateTriggerOid);
-		Assert(OidIsValid(insertTriggerOid));
-		TriggerSetParentTrigger(trigrel, insertTriggerOid, InvalidOid,
-								RelationGetRelid(partRel));
+		// Assert(OidIsValid(insertTriggerOid));
+		if (OidIsValid(insertTriggerOid))
+			TriggerSetParentTrigger(trigrel, insertTriggerOid, InvalidOid,
+									RelationGetRelid(partRel));
 		Assert(OidIsValid(updateTriggerOid));
 		TriggerSetParentTrigger(trigrel, updateTriggerOid, InvalidOid,
 								RelationGetRelid(partRel));
diff --git a/src/test/regress/expected/foreign_key.out b/src/test/regress/expected/foreign_key.out
index 3f459f70ac..3399efe555 100644
--- a/src/test/regress/expected/foreign_key.out
+++ b/src/test/regress/expected/foreign_key.out
@@ -2946,6 +2946,26 @@ DETAIL:  drop cascades to table fkpart11.pk
 drop cascades to table fkpart11.fk_parted
 drop cascades to table fkpart11.fk_another
 drop cascades to function fkpart11.print_row()
+---BUG #18741 https://postgr.es/m/18741-e4ef6f7aa8a956cb@postgresql.org
+CREATE TABLE pk (a int PRIMARY KEY) PARTITION BY RANGE (a);
+ERROR:  relation "pk" already exists
+CREATE TABLE fk (a int) PARTITION BY RANGE (a);
+CREATE TABLE fk_1 (a int, FOREIGN KEY (a) REFERENCES pk);
+--should fail. fk_1 already referencing pk, but fk itself didn't
+ALTER TABLE fk ATTACH PARTITION fk_1 FOR VALUES FROM (0) TO (1);
+ERROR:  table "fk_1" is still referenced by another partitioned table
+DETAIL:  table referenced by another partitioned table can not be attach as a partition.
+DROP TABLE fk, fk_1, pk;
+CREATE TABLE pk (a int primary key) PARTITION BY RANGE (a);
+CREATE TABLE fk (a int not null, FOREIGN KEY (a) REFERENCES pk);
+ALTER TABLE pk ATTACH PARTITION fk FOR VALUES FROM (0) TO (1);
+DROP TABLE fk, pk;
+--should ok. since fk_1 referencing "pk" is non-partitioned table
+CREATE TABLE pk (a int primary key);
+CREATE TABLE fk (a int not null) partition by range(a);
+CREATE TABLE fk_1 (a int not null, FOREIGN KEY (a) REFERENCES pk) partition by range(a);
+ALTER TABLE fk ATTACH PARTITION fk_1 FOR VALUES FROM (0) TO (1);
+DROP TABLE fk, pk;
 -- When a table is attached as partition to a partitioned table that has
 -- a foreign key to another partitioned table, it acquires a clone of the
 -- FK.  Upon detach, this clone is not removed, but instead becomes an
diff --git a/src/test/regress/sql/foreign_key.sql b/src/test/regress/sql/foreign_key.sql
index 2e710e419c..d240b9d583 100644
--- a/src/test/regress/sql/foreign_key.sql
+++ b/src/test/regress/sql/foreign_key.sql
@@ -2087,6 +2087,27 @@ UPDATE fkpart11.pk SET a = 1 WHERE a = 2;
 
 DROP SCHEMA fkpart11 CASCADE;
 
+
+---BUG #18741 https://postgr.es/m/18741-e4ef6f7aa8a956cb@postgresql.org
+CREATE TABLE pk (a int PRIMARY KEY) PARTITION BY RANGE (a);
+CREATE TABLE fk (a int) PARTITION BY RANGE (a);
+CREATE TABLE fk_1 (a int, FOREIGN KEY (a) REFERENCES pk);
+--should fail. fk_1 already referencing pk, but fk itself didn't
+ALTER TABLE fk ATTACH PARTITION fk_1 FOR VALUES FROM (0) TO (1);
+DROP TABLE fk, fk_1, pk;
+
+CREATE TABLE pk (a int primary key) PARTITION BY RANGE (a);
+CREATE TABLE fk (a int not null, FOREIGN KEY (a) REFERENCES pk);
+ALTER TABLE pk ATTACH PARTITION fk FOR VALUES FROM (0) TO (1);
+DROP TABLE fk, pk;
+
+--should ok. since fk_1 referencing "pk" is non-partitioned table
+CREATE TABLE pk (a int primary key);
+CREATE TABLE fk (a int not null) partition by range(a);
+CREATE TABLE fk_1 (a int not null, FOREIGN KEY (a) REFERENCES pk) partition by range(a);
+ALTER TABLE fk ATTACH PARTITION fk_1 FOR VALUES FROM (0) TO (1);
+DROP TABLE fk, pk;
+
 -- When a table is attached as partition to a partitioned table that has
 -- a foreign key to another partitioned table, it acquires a clone of the
 -- FK.  Upon detach, this clone is not removed, but instead becomes an
-- 
2.34.1

