From 349750385cd3df12c02654a91997f48f2abd9832 Mon Sep 17 00:00:00 2001
From: Shlok Kyal <shlok.kyal.oss@gmail.com>
Date: Wed, 29 Jan 2025 10:18:53 +0530
Subject: [PATCH v3] Restrict publishing of partitioned table with a foreign
 table as partition

Logical replication of foreign table is not supported and we throw an
error in this case. But when create a publication on a partitioned
table that has a foreign table as partition, the initial sync of such
table is successful. We should also throw an error in such cases.
With this patch we will throw an error when we try create a publication
on (or add to existing publication) a partitioned table with foreign
table as its partition. We will also throw an error when we try to
attach such table to existing published tables.
---
 src/backend/catalog/pg_publication.c      | 151 ++++++++++++++++++++++
 src/backend/commands/foreigncmds.c        |  40 ++++++
 src/backend/commands/publicationcmds.c    |   3 +
 src/backend/commands/tablecmds.c          |  44 +++++++
 src/include/catalog/pg_publication.h      |   8 ++
 src/test/regress/expected/publication.out |  39 ++++++
 src/test/regress/sql/publication.sql      |  37 ++++++
 7 files changed, 322 insertions(+)

diff --git a/src/backend/catalog/pg_publication.c b/src/backend/catalog/pg_publication.c
index 41ffd494c8..071b428bc0 100644
--- a/src/backend/catalog/pg_publication.c
+++ b/src/backend/catalog/pg_publication.c
@@ -55,6 +55,8 @@ typedef struct
 static void
 check_publication_add_relation(Relation targetrel)
 {
+	Oid			foreign_tbl_relid;
+
 	/* Must be a regular or partitioned table */
 	if (RelationGetForm(targetrel)->relkind != RELKIND_RELATION &&
 		RelationGetForm(targetrel)->relkind != RELKIND_PARTITIONED_TABLE)
@@ -64,6 +66,19 @@ check_publication_add_relation(Relation targetrel)
 						RelationGetRelationName(targetrel)),
 				 errdetail_relkind_not_supported(RelationGetForm(targetrel)->relkind)));
 
+	/*
+	 * Check if it is a partitioned table and any foreign table is its
+	 * partition
+	 */
+	if (check_partrel_has_foreign_table(RelationGetForm(targetrel), &foreign_tbl_relid))
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				 errmsg("cannot add relation \"%s\" to publication",
+						RelationGetRelationName(targetrel)),
+				 errdetail("foreign table \"%s\" is a partition of partitioned table \"%s\"",
+						   get_rel_name(foreign_tbl_relid),
+						   RelationGetRelationName(targetrel))));
+
 	/* Can't be system table */
 	if (IsCatalogRelation(targetrel))
 		ereport(ERROR,
@@ -695,6 +710,9 @@ publication_add_schema(Oid pubid, Oid schemaid, bool if_not_exists)
 
 	check_publication_add_schema(schemaid);
 
+	/* check if schema has any foreign table as partition table */
+	check_foreign_tables_in_schema(schemaid);
+
 	/* Form a tuple */
 	memset(values, 0, sizeof(values));
 	memset(nulls, false, sizeof(nulls));
@@ -1324,3 +1342,136 @@ pg_get_publication_tables(PG_FUNCTION_ARGS)
 
 	SRF_RETURN_DONE(funcctx);
 }
+
+/* Check if a partitioned table has a foreign partition*/
+bool
+check_partrel_has_foreign_table(Form_pg_class relform, Oid *foreign_tbl_relid)
+{
+	bool		has_foreign_tbl = false;
+
+	*foreign_tbl_relid = InvalidOid;
+	if (relform->relkind == RELKIND_PARTITIONED_TABLE)
+	{
+		List	   *relids = NIL;
+
+		relids = GetPubPartitionOptionRelations(relids, PUBLICATION_PART_LEAF,
+												relform->oid);
+
+		foreach_oid(relid, relids)
+		{
+			Relation	rel = table_open(relid, AccessShareLock);
+
+			if (RelationGetForm(rel)->relkind == RELKIND_FOREIGN_TABLE)
+			{
+				has_foreign_tbl = true;
+				*foreign_tbl_relid = relid;
+			}
+
+			table_close(rel, AccessShareLock);
+
+			if (has_foreign_tbl)
+				break;
+		}
+	}
+
+	return has_foreign_tbl;
+}
+
+/* Check if a schema has a partitioned table which has a foreign partition */
+void
+check_foreign_tables_in_schema(Oid schemaid)
+{
+	Relation	classRel;
+	ScanKeyData key[2];
+	TableScanDesc scan;
+	HeapTuple	tuple;
+
+	classRel = table_open(RelationRelationId, AccessShareLock);
+
+	ScanKeyInit(&key[0],
+				Anum_pg_class_relnamespace,
+				BTEqualStrategyNumber, F_OIDEQ,
+				schemaid);
+	ScanKeyInit(&key[1],
+				Anum_pg_class_relkind,
+				BTEqualStrategyNumber, F_CHAREQ,
+				CharGetDatum(RELKIND_PARTITIONED_TABLE));
+
+	scan = table_beginscan_catalog(classRel, 2, key);
+	while ((tuple = heap_getnext(scan, ForwardScanDirection)) != NULL)
+	{
+		Form_pg_class relForm = (Form_pg_class) GETSTRUCT(tuple);
+		Oid			foreign_tbl_relid;
+
+		if (check_partrel_has_foreign_table(relForm, &foreign_tbl_relid))
+		{
+			List	   *ancestors = get_partition_ancestors(relForm->oid);
+			Oid			parent_oid = relForm->oid;
+			char	   *parent_name;
+
+			foreach_oid(ancestor, ancestors)
+			{
+				Oid			ancestor_schemaid = get_rel_namespace(ancestor);
+
+				if (ancestor_schemaid == schemaid)
+					parent_oid = ancestor;
+			}
+
+			list_free(ancestors);
+			parent_name = get_rel_name(parent_oid);
+
+			ereport(ERROR,
+					(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+					 errmsg("cannot add relation \"%s\" to publication",
+							parent_name),
+					 errdetail("foreign table \"%s\" is a partition of partitioned table \"%s\"",
+							   get_rel_name(foreign_tbl_relid), parent_name)));
+		}
+	}
+
+	table_endscan(scan);
+	table_close(classRel, AccessShareLock);
+}
+
+/* Check if any foreign table is a partition table */
+void
+check_foreign_tables(void)
+{
+	Relation	classRel;
+	ScanKeyData key[1];
+	TableScanDesc scan;
+	HeapTuple	tuple;
+
+	classRel = table_open(RelationRelationId, AccessShareLock);
+
+	ScanKeyInit(&key[0],
+				Anum_pg_class_relkind,
+				BTEqualStrategyNumber, F_CHAREQ,
+				CharGetDatum(RELKIND_FOREIGN_TABLE));
+
+	scan = table_beginscan_catalog(classRel, 1, key);
+	while ((tuple = heap_getnext(scan, ForwardScanDirection)) != NULL)
+	{
+		Form_pg_class relForm = (Form_pg_class) GETSTRUCT(tuple);
+
+		if (relForm->relispartition)
+		{
+			Oid			parent_oid;
+			char	   *parent_name;
+			List	   *ancestors = get_partition_ancestors(relForm->oid);
+
+			parent_oid = llast_oid(ancestors);
+			parent_name = get_rel_name(parent_oid);
+
+			ereport(ERROR,
+					(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+					 errmsg("cannot add relation \"%s\" to publication",
+							parent_name),
+					 errdetail("foreign table \"%s\" is a partition of partitioned table \"%s\"",
+							   NameStr(relForm->relname), parent_name)));
+		}
+	}
+
+	table_endscan(scan);
+	table_close(classRel, AccessShareLock);
+}
diff --git a/src/backend/commands/foreigncmds.c b/src/backend/commands/foreigncmds.c
index c14e038d54..004edd9515 100644
--- a/src/backend/commands/foreigncmds.c
+++ b/src/backend/commands/foreigncmds.c
@@ -21,6 +21,7 @@
 #include "catalog/dependency.h"
 #include "catalog/indexing.h"
 #include "catalog/objectaccess.h"
+#include "catalog/partition.h"
 #include "catalog/pg_foreign_data_wrapper.h"
 #include "catalog/pg_foreign_server.h"
 #include "catalog/pg_foreign_table.h"
@@ -1423,6 +1424,45 @@ CreateForeignTable(CreateForeignTableStmt *stmt, Oid relid)
 
 	ftrel = table_open(ForeignTableRelationId, RowExclusiveLock);
 
+	/*
+	 * Check if it is a foreign partition and the partitioned table is
+	 * published
+	 */
+	if (stmt->base.partbound != NULL)
+	{
+		RangeVar   *root = castNode(RangeVar, lfirst(list_head(stmt->base.inhRelations)));
+		Relation	rootrel = table_openrv(root, AccessShareLock);
+
+		if (RelationGetForm(rootrel)->relkind == RELKIND_PARTITIONED_TABLE)
+		{
+			Oid			schemaid = RelationGetNamespace(rootrel);
+			List	   *puboids = GetRelationPublications(rootrel->rd_id);
+			List	   *ancestors;
+
+			puboids = list_concat_unique_oid(puboids, GetAllTablesPublications());
+			puboids = list_concat_unique_oid(puboids, GetSchemaPublications(schemaid));
+			ancestors = get_partition_ancestors(rootrel->rd_id);
+
+			foreach_oid(ancestor, ancestors)
+			{
+				puboids = list_concat_unique_oid(puboids,
+												 GetRelationPublications(ancestor));
+				schemaid = get_rel_namespace(ancestor);
+				puboids = list_concat_unique_oid(puboids,
+												 GetSchemaPublications(schemaid));
+			}
+			list_free(ancestors);
+
+			if (puboids)
+				ereport(ERROR,
+						(errcode(ERRCODE_WRONG_OBJECT_TYPE),
+						 errmsg("cannot create foreign partition \"%s\" as partitioned table \"%s\" is published",
+								get_rel_name(relid),
+								RelationGetRelationName(rootrel))));
+		}
+		table_close(rootrel, AccessShareLock);
+	}
+
 	/*
 	 * For now the owner cannot be specified on create. Use effective user ID.
 	 */
diff --git a/src/backend/commands/publicationcmds.c b/src/backend/commands/publicationcmds.c
index 951ffabb65..7258d6c33d 100644
--- a/src/backend/commands/publicationcmds.c
+++ b/src/backend/commands/publicationcmds.c
@@ -855,6 +855,9 @@ CreatePublication(ParseState *pstate, CreatePublicationStmt *stmt)
 	/* Associate objects with the publication. */
 	if (stmt->for_all_tables)
 	{
+		/* Check if any foreign table is a part of partitioned table */
+		check_foreign_tables();
+
 		/* Invalidate relcache so that publication info is rebuilt. */
 		CacheInvalidateRelcacheAll();
 	}
diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index d617c4bc63..82e07509a8 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -19222,6 +19222,7 @@ ATExecAttachPartition(List **wqueue, Relation rel, PartitionCmd *cmd,
 	Oid			defaultPartOid;
 	List	   *partBoundConstraint;
 	ParseState *pstate = make_parsestate(NULL);
+	Oid			foreign_tbl_relid;
 
 	pstate->p_sourcetext = context->queryString;
 
@@ -19347,6 +19348,49 @@ ATExecAttachPartition(List **wqueue, Relation rel, PartitionCmd *cmd,
 				(errcode(ERRCODE_WRONG_OBJECT_TYPE),
 				 errmsg("cannot attach temporary relation of another session as partition")));
 
+	/*
+	 * Check if attachrel is a foreign table or a partitioned table with
+	 * foreign partition and rel is not published.
+	 */
+	if (attachrel->rd_rel->relkind == RELKIND_FOREIGN_TABLE ||
+		check_partrel_has_foreign_table(RelationGetForm(attachrel), &foreign_tbl_relid))
+	{
+		Oid			schemaid = RelationGetNamespace(rel);
+		List	   *puboids = GetRelationPublications(rel->rd_id);
+		List	   *ancestors;
+
+		puboids = list_concat_unique_oid(puboids, GetAllTablesPublications());
+		puboids = list_concat_unique_oid(puboids, GetSchemaPublications(schemaid));
+		ancestors = get_partition_ancestors(rel->rd_id);
+
+		foreach_oid(ancestor, ancestors)
+		{
+			puboids = list_concat_unique_oid(puboids,
+											 GetRelationPublications(ancestor));
+			schemaid = get_rel_namespace(ancestor);
+			puboids = list_concat_unique_oid(puboids,
+											 GetSchemaPublications(schemaid));
+		}
+		list_free(ancestors);
+
+		if (puboids)
+		{
+			if (attachrel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE)
+				ereport(ERROR,
+						(errcode(ERRCODE_WRONG_OBJECT_TYPE),
+						 errmsg("cannot attach a partitioned table with a foreign partition to a published table"),
+						 errdetail("foreign table \"%s\" is a partition of partitioned table \"%s\"",
+								   get_rel_name(foreign_tbl_relid),
+								   RelationGetRelationName(attachrel))));
+			else
+				ereport(ERROR,
+						(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+						 errmsg("cannot attach foreign table \"%s\" to a published table",
+								get_rel_name(attachrel->rd_id))));
+
+		}
+	}
+
 	/*
 	 * Check if attachrel has any identity columns or any columns that aren't
 	 * in the parent.
diff --git a/src/include/catalog/pg_publication.h b/src/include/catalog/pg_publication.h
index 48c7d1a861..ba51f4a721 100644
--- a/src/include/catalog/pg_publication.h
+++ b/src/include/catalog/pg_publication.h
@@ -19,6 +19,7 @@
 
 #include "catalog/genbki.h"
 #include "catalog/objectaddress.h"
+#include "catalog/pg_class.h"
 #include "catalog/pg_publication_d.h"	/* IWYU pragma: export */
 
 /* ----------------
@@ -191,4 +192,11 @@ extern Bitmapset *pub_collist_to_bitmapset(Bitmapset *columns, Datum pubcols,
 extern Bitmapset *pub_form_cols_map(Relation relation,
 									PublishGencolsType include_gencols_type);
 
+extern bool check_partrel_has_foreign_table(Form_pg_class relform,
+											Oid *foreign_tbl_name);
+
+extern void check_foreign_tables_in_schema(Oid schemaid);
+
+extern void check_foreign_tables(void);
+
 #endif							/* PG_PUBLICATION_H */
diff --git a/src/test/regress/expected/publication.out b/src/test/regress/expected/publication.out
index bc3898fbe5..b973f41e94 100644
--- a/src/test/regress/expected/publication.out
+++ b/src/test/regress/expected/publication.out
@@ -1885,6 +1885,45 @@ DROP PUBLICATION pub1;
 DROP PUBLICATION pub2;
 DROP TABLE gencols;
 RESET client_min_messages;
+-- ======================================================
+-- Test when foreign table is a partition of a partitioned table on which
+-- publication is created
+SET client_min_messages = 'ERROR';
+CREATE FOREIGN DATA WRAPPER test_fdw;
+CREATE SERVER fdw_server FOREIGN DATA WRAPPER test_fdw;
+CREATE SCHEMA sch3;
+CREATE TABLE sch3.tmain(id int) PARTITION BY RANGE(id);
+CREATE TABLE sch3.part1 PARTITION OF sch3.tmain FOR VALUES FROM (0) TO (5);
+CREATE TABLE sch3.part2(id int) PARTITION BY RANGE(id);
+CREATE FOREIGN TABLE sch3.part2_1 PARTITION OF sch3.part2 FOR VALUES FROM (5) TO (10) SERVER fdw_server;
+ALTER TABLE sch3.tmain ATTACH PARTITION sch3.part2 FOR VALUES FROM (5) TO (10);
+CREATE PUBLICATION pub1 FOR TABLE sch3.tmain;
+ERROR:  cannot add relation "tmain" to publication
+DETAIL:  foreign table "part2_1" is a partition of partitioned table "tmain"
+CREATE PUBLICATION pub1 FOR TABLES IN SCHEMA sch3;
+ERROR:  cannot add relation "tmain" to publication
+DETAIL:  foreign table "part2_1" is a partition of partitioned table "tmain"
+CREATE PUBLICATION pub1 FOR ALL TABLES;
+ERROR:  cannot add relation "tmain" to publication
+DETAIL:  foreign table "part2_1" is a partition of partitioned table "tmain"
+-- Test when a partitioned table with foreign table as a partition is attached
+-- to partitioned table which is already published
+ALTER TABLE sch3.tmain DETACH PARTITION sch3.part2;
+CREATE PUBLICATION pub1 FOR TABLE sch3.tmain;
+ALTER TABLE sch3.tmain ATTACH PARTITION sch3.part2 FOR VALUES FROM (5) TO (10);
+ERROR:  cannot attach a partitioned table with a foreign partition to a published table
+DETAIL:  foreign table "part2_1" is a partition of partitioned table "part2"
+-- Can't create foreign partition of published table
+CREATE FOREIGN TABLE sch3.part3_1 PARTITION OF sch3.tmain FOR VALUES FROM (10) TO (15) SERVER fdw_server;
+ERROR:  cannot create foreign partition "part3_1" as partitioned table "tmain" is published
+-- Can't attach foreign partition to published table
+CREATE FOREIGN TABLE sch3.part3_2(id int) SERVER fdw_server;
+ALTER TABLE sch3.tmain ATTACH PARTITION sch3.part3_2 FOR VALUES FROM (15) TO (20);
+ERROR:  cannot attach foreign table "part3_2" to a published table
+DROP PUBLICATION pub1;
+DROP SCHEMA sch3 CASCADE;
+DROP SERVER fdw_server;
+DROP FOREIGN DATA WRAPPER test_fdw;
 RESET SESSION AUTHORIZATION;
 DROP ROLE regress_publication_user, regress_publication_user2;
 DROP ROLE regress_publication_user_dummy;
diff --git a/src/test/regress/sql/publication.sql b/src/test/regress/sql/publication.sql
index 47f0329c24..b9811033dc 100644
--- a/src/test/regress/sql/publication.sql
+++ b/src/test/regress/sql/publication.sql
@@ -1186,6 +1186,43 @@ DROP PUBLICATION pub2;
 DROP TABLE gencols;
 
 RESET client_min_messages;
+-- ======================================================
+
+-- Test when foreign table is a partition of a partitioned table on which
+-- publication is created
+SET client_min_messages = 'ERROR';
+CREATE FOREIGN DATA WRAPPER test_fdw;
+CREATE SERVER fdw_server FOREIGN DATA WRAPPER test_fdw;
+
+CREATE SCHEMA sch3;
+CREATE TABLE sch3.tmain(id int) PARTITION BY RANGE(id);
+CREATE TABLE sch3.part1 PARTITION OF sch3.tmain FOR VALUES FROM (0) TO (5);
+CREATE TABLE sch3.part2(id int) PARTITION BY RANGE(id);
+CREATE FOREIGN TABLE sch3.part2_1 PARTITION OF sch3.part2 FOR VALUES FROM (5) TO (10) SERVER fdw_server;
+ALTER TABLE sch3.tmain ATTACH PARTITION sch3.part2 FOR VALUES FROM (5) TO (10);
+
+CREATE PUBLICATION pub1 FOR TABLE sch3.tmain;
+CREATE PUBLICATION pub1 FOR TABLES IN SCHEMA sch3;
+CREATE PUBLICATION pub1 FOR ALL TABLES;
+
+-- Test when a partitioned table with foreign table as a partition is attached
+-- to partitioned table which is already published
+ALTER TABLE sch3.tmain DETACH PARTITION sch3.part2;
+CREATE PUBLICATION pub1 FOR TABLE sch3.tmain;
+ALTER TABLE sch3.tmain ATTACH PARTITION sch3.part2 FOR VALUES FROM (5) TO (10);
+
+-- Can't create foreign partition of published table
+CREATE FOREIGN TABLE sch3.part3_1 PARTITION OF sch3.tmain FOR VALUES FROM (10) TO (15) SERVER fdw_server;
+
+-- Can't attach foreign partition to published table
+CREATE FOREIGN TABLE sch3.part3_2(id int) SERVER fdw_server;
+ALTER TABLE sch3.tmain ATTACH PARTITION sch3.part3_2 FOR VALUES FROM (15) TO (20);
+
+DROP PUBLICATION pub1;
+DROP SCHEMA sch3 CASCADE;
+DROP SERVER fdw_server;
+DROP FOREIGN DATA WRAPPER test_fdw;
+
 RESET SESSION AUTHORIZATION;
 DROP ROLE regress_publication_user, regress_publication_user2;
 DROP ROLE regress_publication_user_dummy;
-- 
2.34.1

