From 30c4b8a61a21e83b27251908f961395f183a7647 Mon Sep 17 00:00:00 2001
From: "houzj.fnst" <houzj.fnst@cn.fujitsu.com>
Date: Sat, 16 Oct 2021 12:32:13 +0800
Subject: [PATCH] Topup Hou

---
 src/backend/catalog/pg_publication.c      | 121 +++++++++++++++++++++++++++---
 src/backend/commands/publicationcmds.c    |   6 +-
 src/include/catalog/pg_publication.h      |   6 +-
 src/test/regress/expected/publication.out |  59 +++++++++++++++
 src/test/regress/sql/publication.sql      |  41 ++++++++++
 5 files changed, 217 insertions(+), 16 deletions(-)

diff --git a/src/backend/catalog/pg_publication.c b/src/backend/catalog/pg_publication.c
index 04be114..f6343f6 100644
--- a/src/backend/catalog/pg_publication.c
+++ b/src/backend/catalog/pg_publication.c
@@ -131,6 +131,59 @@ is_publishable_class(Oid relid, Form_pg_class reltuple)
 }
 
 /*
+ * Filter out the partitions whose parent tables was also specified in
+ * the publication.
+ */
+static List *
+filter_out_partitions(Oid pubid, List *relids,
+					  PublicationObjSpecType reltype)
+{
+	List	   *result = NIL;
+	List	   *published_relids = NIL;
+	List	   *schemaids = NIL;
+	ListCell   *lc;
+	ListCell   *lc2;
+
+	if (reltype == PUBLICATIONOBJ_REL_IN_SCHEMA)
+		published_relids = GetPublicationRelations(pubid,
+												   PUBLICATION_PART_ROOT);
+	else if (reltype == PUBLICATIONOBJ_TABLE)
+		schemaids = GetPublicationSchemas(pubid);
+
+	foreach(lc, relids)
+	{
+		bool		skip = false;
+		List	   *ancestors = NIL;
+		Oid			relid = lfirst_oid(lc);
+
+		if (get_rel_relispartition(relid))
+			ancestors = get_partition_ancestors(relid);
+
+		foreach(lc2, ancestors)
+		{
+			/*
+			 * Check if the parent table exists in the published table list or
+			 * the schema of the parent table was published in this
+			 * publication.
+			 */
+			if (list_member_oid(relids, lfirst_oid(lc2)) ||
+				list_member_oid(published_relids, lfirst_oid(lc2)) ||
+				(schemaids != NIL &&
+				 list_member_oid(schemaids, get_rel_namespace(lfirst_oid(lc2)))))
+			{
+				skip = true;
+				break;
+			}
+		}
+
+		if (!skip)
+			result = lappend_oid(result, relid);
+	}
+
+	return result;
+}
+
+/*
  * Another variant of this, taking a Relation.
  */
 bool
@@ -363,7 +416,8 @@ publication_add_schema(Oid pubid, Oid schemaid, bool if_not_exists)
 	 * publication_add_relation for why we need to consider all the
 	 * partitions.
 	 */
-	schemaRels = GetSchemaPublicationRelations(schemaid, PUBLICATION_PART_ALL);
+	schemaRels = GetSchemaPublicationRelations(schemaid,
+											   PUBLICATION_PART_ALL);
 	InvalidatePublicationRels(schemaRels);
 
 	return myself;
@@ -436,6 +490,34 @@ GetPublicationRelations(Oid pubid, PublicationPartOpt pub_partopt)
 }
 
 /*
+ * Gets list of relation explicitly mentioned in the publication.
+ *
+ * This should only be used FOR TABLE publications, the FOR ALL TABLES
+ * should use GetAllTablesPublicationRelations().
+ *
+ * If pubviaroot is true, we must exclude partitions in favor of including the
+ * root partitioned tables.
+ */
+List *
+GetPublicationRelationsExtended(Oid pubid, bool pubviaroot)
+{
+	List	   *result;
+
+	result = GetPublicationRelations(pubid, PUBLICATION_PART_ROOT);
+
+	/*
+	 * Filter out the partitions whose parent table is also in the same
+	 * publication.
+	 */
+	if (pubviaroot)
+		return filter_out_partitions(pubid, result, PUBLICATIONOBJ_TABLE);
+	else
+		return result;
+}
+
+
+
+/*
  * Gets list of publication oids for publications marked as FOR ALL TABLES.
  */
 List *
@@ -608,7 +690,8 @@ GetSchemaPublications(Oid schemaid)
  * relkind tables, so two rounds of scan are required.
  */
 List *
-GetSchemaPublicationRelations(Oid schemaid, PublicationPartOpt pub_partopt)
+GetSchemaPublicationRelations(Oid schemaid,
+							  PublicationPartOpt pub_partopt)
 {
 	Relation	classRel;
 	ScanKeyData key[3];
@@ -638,8 +721,7 @@ GetSchemaPublicationRelations(Oid schemaid, PublicationPartOpt pub_partopt)
 		Form_pg_class relForm = (Form_pg_class) GETSTRUCT(tuple);
 		Oid			relid = relForm->oid;
 
-		if (is_publishable_class(relid, relForm) &&
-			!(relForm->relispartition && pub_partopt == PUBLICATION_PART_ROOT))
+		if (is_publishable_class(relid, relForm))
 			result = lappend_oid(result, relid);
 	}
 
@@ -685,12 +767,16 @@ GetSchemaPublicationRelations(Oid schemaid, PublicationPartOpt pub_partopt)
 /*
  * Gets the list of all relations published by FOR ALL TABLES IN SCHEMA
  * publication.
+ *
+ * If pubviaroot is true, we must exclude partitions in favor of including the
+ * root partitioned tables.
  */
 List *
-GetAllSchemaPublicationRelations(Oid puboid, PublicationPartOpt pub_partopt)
+GetAllSchemaPublicationRelations(Oid pubid, PublicationPartOpt pub_partopt,
+								 bool pubviaroot)
 {
 	List	   *result = NIL;
-	List	   *pubschemalist = GetPublicationSchemas(puboid);
+	List	   *pubschemalist = GetPublicationSchemas(pubid);
 	ListCell   *cell;
 
 	foreach(cell, pubschemalist)
@@ -702,7 +788,14 @@ GetAllSchemaPublicationRelations(Oid puboid, PublicationPartOpt pub_partopt)
 		result = list_concat(result, schemaRels);
 	}
 
-	return result;
+	/*
+	 * Filter out the partitions whose parent table is also in the same
+	 * publication.
+	 */
+	if (pubviaroot)
+		return filter_out_partitions(pubid, result, PUBLICATIONOBJ_REL_IN_SCHEMA);
+	else
+		return result;
 }
 
 /*
@@ -840,14 +933,18 @@ pg_get_publication_tables(PG_FUNCTION_ARGS)
 			List	   *relids,
 					   *schemarelids;
 
-			relids = GetPublicationRelations(publication->oid,
-											 publication->pubviaroot ?
-											 PUBLICATION_PART_ROOT :
-											 PUBLICATION_PART_LEAF);
+			if (publication->pubviaroot)
+				relids = GetPublicationRelationsExtended(publication->oid,
+														 publication->pubviaroot);
+			else
+				relids = GetPublicationRelations(publication->oid,
+												 PUBLICATION_PART_LEAF);
+
 			schemarelids = GetAllSchemaPublicationRelations(publication->oid,
 															publication->pubviaroot ?
 															PUBLICATION_PART_ROOT :
-															PUBLICATION_PART_LEAF);
+															PUBLICATION_PART_LEAF,
+															publication->pubviaroot);
 
 			tables = list_concat_unique_oid(relids, schemarelids);
 		}
diff --git a/src/backend/commands/publicationcmds.c b/src/backend/commands/publicationcmds.c
index 55079a6..d867c4e 100644
--- a/src/backend/commands/publicationcmds.c
+++ b/src/backend/commands/publicationcmds.c
@@ -452,7 +452,8 @@ AlterPublicationOptions(ParseState *pstate, AlterPublicationStmt *stmt,
 		relids = GetPublicationRelations(pubform->oid,
 										 PUBLICATION_PART_ALL);
 		schemarelids = GetAllSchemaPublicationRelations(pubform->oid,
-														PUBLICATION_PART_ALL);
+														PUBLICATION_PART_ALL,
+														false);
 		relids = list_concat_unique_oid(relids, schemarelids);
 
 		InvalidatePublicationRels(relids);
@@ -830,7 +831,8 @@ RemovePublicationSchemaById(Oid psoid)
 
 	/*
 	 * Invalidate relcache so that publication info is rebuilt. See
-	 * RemovePublicationRelById for why we need to consider all the partitions.
+	 * RemovePublicationRelById for why we need to consider all the
+	 * partitions.
 	 */
 	schemaRels = GetSchemaPublicationRelations(pubsch->pnnspid,
 											   PUBLICATION_PART_ALL);
diff --git a/src/include/catalog/pg_publication.h b/src/include/catalog/pg_publication.h
index a4c894e..9672649 100644
--- a/src/include/catalog/pg_publication.h
+++ b/src/include/catalog/pg_publication.h
@@ -104,6 +104,7 @@ typedef enum PublicationPartOpt
 } PublicationPartOpt;
 
 extern List *GetPublicationRelations(Oid pubid, PublicationPartOpt pub_partopt);
+extern List *GetPublicationRelationsExtended(Oid pubid, bool pubviaroot);
 extern List *GetAllTablesPublications(void);
 extern List *GetAllTablesPublicationRelations(bool pubviaroot);
 extern List *GetPublicationSchemas(Oid pubid);
@@ -111,10 +112,11 @@ extern List *GetSchemaPublications(Oid schemaid);
 extern List *GetSchemaPublicationRelations(Oid schemaid,
 										   PublicationPartOpt pub_partopt);
 extern List *GetAllSchemaPublicationRelations(Oid puboid,
-											  PublicationPartOpt pub_partopt);
+											  PublicationPartOpt pub_partopt,
+											  bool pubviaroot);
 
 extern bool is_publishable_relation(Relation rel);
-extern List *GetPubPartitionOptionRelations(List* result,
+extern List *GetPubPartitionOptionRelations(List *result,
 											PublicationPartOpt pub_partopt,
 											Oid relid);
 extern ObjectAddress publication_add_relation(Oid pubid, Relation targetrel,
diff --git a/src/test/regress/expected/publication.out b/src/test/regress/expected/publication.out
index 459371b..d34c1ff 100644
--- a/src/test/regress/expected/publication.out
+++ b/src/test/regress/expected/publication.out
@@ -772,6 +772,65 @@ DETAIL:  drop cascades to table pub_testpart1.parent1
 drop cascades to table pub_testpart1.child_parent2
 DROP SCHEMA pub_testpart2 CASCADE;
 NOTICE:  drop cascades to table pub_testpart2.parent2
+-- pg_publication_tables
+SET client_min_messages = 'ERROR';
+CREATE SCHEMA sch1;
+CREATE SCHEMA sch2;
+CREATE TABLE sch1.tbl1 (a int) PARTITION BY RANGE(a);
+CREATE TABLE sch2.tbl1_part1 PARTITION OF sch1.tbl1 FOR VALUES FROM (1) to (10);
+CREATE PUBLICATION pub FOR ALL TABLES IN SCHEMA sch2 WITH (PUBLISH_VIA_PARTITION_ROOT=1);
+SELECT * FROM pg_publication_tables;
+ pubname | schemaname | tablename  
+---------+------------+------------
+ pub     | sch2       | tbl1_part1
+(1 row)
+
+DROP PUBLICATION pub;
+CREATE PUBLICATION pub FOR TABLE sch2.tbl1_part1 WITH (PUBLISH_VIA_PARTITION_ROOT=1);
+SELECT * FROM pg_publication_tables;
+ pubname | schemaname | tablename  
+---------+------------+------------
+ pub     | sch2       | tbl1_part1
+(1 row)
+
+DROP PUBLICATION pub;
+CREATE PUBLICATION pub FOR ALL TABLES IN SCHEMA sch2 WITH (PUBLISH_VIA_PARTITION_ROOT=0);
+SELECT * FROM pg_publication_tables;
+ pubname | schemaname | tablename  
+---------+------------+------------
+ pub     | sch2       | tbl1_part1
+(1 row)
+
+DROP PUBLICATION pub;
+CREATE PUBLICATION pub FOR TABLE sch2.tbl1_part1 WITH (PUBLISH_VIA_PARTITION_ROOT=0);
+SELECT * FROM pg_publication_tables;
+ pubname | schemaname | tablename  
+---------+------------+------------
+ pub     | sch2       | tbl1_part1
+(1 row)
+
+DROP PUBLICATION pub;
+DROP TABLE sch2.tbl1_part1;
+DROP TABLE sch1.tbl1;
+CREATE TABLE sch1.tbl1 (a int) PARTITION BY RANGE(a);
+CREATE TABLE sch1.tbl1_part1 PARTITION OF sch1.tbl1 FOR VALUES FROM (1) to (10);
+CREATE TABLE sch1.tbl1_part2 PARTITION OF sch1.tbl1 FOR VALUES FROM (10) to (20);
+CREATE TABLE sch1.tbl1_part3 (a int) PARTITION BY RANGE(a);
+ALTER TABLE sch1.tbl1 ATTACH PARTITION sch1.tbl1_part3 FOR VALUES FROM (20) to (30);
+CREATE PUBLICATION pub FOR ALL TABLES IN SCHEMA sch1 WITH (PUBLISH_VIA_PARTITION_ROOT=1);
+SELECT * FROM pg_publication_tables;
+ pubname | schemaname | tablename 
+---------+------------+-----------
+ pub     | sch1       | tbl1
+(1 row)
+
+RESET client_min_messages;
+DROP PUBLICATION pub;
+DROP TABLE sch2.tbl1_part1;
+ERROR:  table "tbl1_part1" does not exist
+DROP TABLE sch1.tbl1;
+DROP SCHEMA sch1 cascade;
+DROP SCHEMA sch2 cascade;
 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 3ef2f68..5273117 100644
--- a/src/test/regress/sql/publication.sql
+++ b/src/test/regress/sql/publication.sql
@@ -428,6 +428,47 @@ DROP SCHEMA pub_test2 CASCADE;
 DROP SCHEMA pub_testpart1 CASCADE;
 DROP SCHEMA pub_testpart2 CASCADE;
 
+-- pg_publication_tables
+SET client_min_messages = 'ERROR';
+CREATE SCHEMA sch1;
+CREATE SCHEMA sch2;
+CREATE TABLE sch1.tbl1 (a int) PARTITION BY RANGE(a);
+CREATE TABLE sch2.tbl1_part1 PARTITION OF sch1.tbl1 FOR VALUES FROM (1) to (10);
+CREATE PUBLICATION pub FOR ALL TABLES IN SCHEMA sch2 WITH (PUBLISH_VIA_PARTITION_ROOT=1);
+SELECT * FROM pg_publication_tables;
+
+DROP PUBLICATION pub;
+CREATE PUBLICATION pub FOR TABLE sch2.tbl1_part1 WITH (PUBLISH_VIA_PARTITION_ROOT=1);
+SELECT * FROM pg_publication_tables;
+
+DROP PUBLICATION pub;
+CREATE PUBLICATION pub FOR ALL TABLES IN SCHEMA sch2 WITH (PUBLISH_VIA_PARTITION_ROOT=0);
+SELECT * FROM pg_publication_tables;
+
+DROP PUBLICATION pub;
+CREATE PUBLICATION pub FOR TABLE sch2.tbl1_part1 WITH (PUBLISH_VIA_PARTITION_ROOT=0);
+SELECT * FROM pg_publication_tables;
+
+DROP PUBLICATION pub;
+DROP TABLE sch2.tbl1_part1;
+DROP TABLE sch1.tbl1;
+
+CREATE TABLE sch1.tbl1 (a int) PARTITION BY RANGE(a);
+CREATE TABLE sch1.tbl1_part1 PARTITION OF sch1.tbl1 FOR VALUES FROM (1) to (10);
+CREATE TABLE sch1.tbl1_part2 PARTITION OF sch1.tbl1 FOR VALUES FROM (10) to (20);
+CREATE TABLE sch1.tbl1_part3 (a int) PARTITION BY RANGE(a);
+ALTER TABLE sch1.tbl1 ATTACH PARTITION sch1.tbl1_part3 FOR VALUES FROM (20) to (30);
+CREATE PUBLICATION pub FOR ALL TABLES IN SCHEMA sch1 WITH (PUBLISH_VIA_PARTITION_ROOT=1);
+SELECT * FROM pg_publication_tables;
+
+RESET client_min_messages;
+DROP PUBLICATION pub;
+DROP TABLE sch2.tbl1_part1;
+DROP TABLE sch1.tbl1;
+DROP SCHEMA sch1 cascade;
+DROP SCHEMA sch2 cascade;
+
+
 RESET SESSION AUTHORIZATION;
 DROP ROLE regress_publication_user, regress_publication_user2;
 DROP ROLE regress_publication_user_dummy;
-- 
2.7.2.windows.1

