HEAD_v2-0001-Fix-data-replicated-twice-when-specifying-PUBLISH.patch

From e8d9cf1b5b3cbfc590a5f1853aacac905823bb71 Mon Sep 17 00:00:00 2001
From: wangw <wangw.fnst@fujitsu.com>
Date: Thu, 28 Apr 2022 13:37:19 +0800
Subject: [PATCH v2] Fix data replicated twice when specifying
 PUBLISH_VIA_PARTITION_ROOT option.

If there are two publications that publish the parent table and the child table
separately, and both specify the option PUBLISH_VIA_PARTITION_ROOT, when
subscribing to both publications using one subscription, the data is replicated
twice in inital copy. What we expect is to be copied only once.

To fix this, we extended the API of the function pg_get_publication_tables.
(Change the parameter type from text to any).
So, the function pg_get_publication_tables could also receive the array of
publications. And then, if we specify option viaroot, we could exclude the
partitioned table whose ancestor belongs to this publications array when
getting the table list.
---
 src/backend/catalog/pg_publication.c       | 169 ++++++++++++++++-----
 src/backend/commands/subscriptioncmds.c    |  17 ++-
 src/include/catalog/pg_proc.dat            |   6 +-
 src/test/regress/expected/rules.out        |   2 +-
 src/test/subscription/t/013_partition.pl   |  20 +--
 src/test/subscription/t/028_row_filter.pl  |   9 +-
 src/test/subscription/t/031_column_list.pl |   5 +-
 7 files changed, 164 insertions(+), 64 deletions(-)

diff --git a/src/backend/catalog/pg_publication.c b/src/backend/catalog/pg_publication.c
index 2631558ff1..ada93ac8db 100644
--- a/src/backend/catalog/pg_publication.c
+++ b/src/backend/catalog/pg_publication.c
@@ -1077,20 +1077,55 @@ get_publication_name(Oid pubid, bool missing_ok)
 }
 
 /*
- * Returns Oids of tables in a publication.
+ * Returns Oids of tables in one or more publications.
  */
 Datum
 pg_get_publication_tables(PG_FUNCTION_ARGS)
 {
 	FuncCallContext *funcctx;
-	char	   *pubname = text_to_cstring(PG_GETARG_TEXT_PP(0));
-	Publication *publication;
-	List	   *tables;
+	Oid			typeid;
+	ArrayType  *arr;
+	Datum	   *elems;
+	bool	   *nulls;
+	int			nelems;
+	char	   *pubname;
+	List	   *tables_viaroot = NIL,
+			   *tables = NIL,
+			   *current_table = NIL;
+
+	typeid = get_fn_expr_argtype(fcinfo->flinfo, 0);
+
+	switch (typeid)
+	{
+		case TEXTARRAYOID:
+			/* get Oids of tables in publications array */
+			arr = PG_GETARG_ARRAYTYPE_P(0);
+			deconstruct_array(arr, TEXTOID, -1, false, TYPALIGN_INT,
+							  &elems, &nulls, &nelems);
+			break;
+		case TEXTOID:
+		case NAMEOID:
+		case UNKNOWNOID:
+			/* get Oids of tables in a publication */
+			elems = (Datum *) palloc(sizeof(Datum));
+			if (typeid == TEXTOID)
+				elems[0] = PG_GETARG_DATUM(0);
+			else
+				elems[0] = CStringGetTextDatum(PG_GETARG_POINTER(0));
+			nelems = 1;
+			break;
+		default:
+			ereport(ERROR,
+					(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+					 errmsg("do not support this argument type")));
+	}
 
 	/* stuff done only on the first call of the function */
 	if (SRF_IS_FIRSTCALL())
 	{
 		MemoryContext oldcontext;
+		Publication *publication;
+		int	i;
 
 		/* create a function context for cross-call persistence */
 		funcctx = SRF_FIRSTCALL_INIT();
@@ -1098,44 +1133,108 @@ pg_get_publication_tables(PG_FUNCTION_ARGS)
 		/* switch to memory context appropriate for multiple function calls */
 		oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
 
-		publication = GetPublicationByName(pubname, false);
-
-		/*
-		 * Publications support partitioned tables, although all changes are
-		 * replicated using leaf partition identity and schema, so we only
-		 * need those.
-		 */
-		if (publication->alltables)
-		{
-			tables = GetAllTablesPublicationRelations(publication->pubviaroot);
-		}
-		else
+		for (i = 0; i < nelems; i++)
 		{
-			List	   *relids,
-					   *schemarelids;
-
-			relids = GetPublicationRelations(publication->oid,
-											 publication->pubviaroot ?
-											 PUBLICATION_PART_ROOT :
-											 PUBLICATION_PART_LEAF);
-			schemarelids = GetAllSchemaPublicationRelations(publication->oid,
-															publication->pubviaroot ?
-															PUBLICATION_PART_ROOT :
-															PUBLICATION_PART_LEAF);
-			tables = list_concat_unique_oid(relids, schemarelids);
+			publication = GetPublicationByName(TextDatumGetCString(elems[i]), true);
+
+			if (!publication)
+				continue;
+
+			/*
+			 * Publications support partitioned tables, although all changes
+			 * are replicated using leaf partition identity and schema, so we
+			 * only need those.
+			 */
+			if (publication->alltables)
+			{
+				current_table = GetAllTablesPublicationRelations(publication->pubviaroot);
+			}
+			else
+			{
+				List	   *relids,
+						   *schemarelids;
+
+				relids = GetPublicationRelations(publication->oid,
+												 publication->pubviaroot ?
+												 PUBLICATION_PART_ROOT :
+												 PUBLICATION_PART_LEAF);
+				schemarelids = GetAllSchemaPublicationRelations(publication->oid,
+																publication->pubviaroot ?
+																PUBLICATION_PART_ROOT :
+																PUBLICATION_PART_LEAF);
+				current_table = list_concat_unique_oid(relids, schemarelids);
+			}
 
 			/*
-			 * If the publication publishes partition changes via their
-			 * respective root partitioned tables, we must exclude partitions
-			 * in favor of including the root partitioned tables. Otherwise,
-			 * the function could return both the child and parent tables
-			 * which could cause data of the child table to be
-			 * double-published on the subscriber side.
+			 * Store the published tables separately according to pubviaroot.
+			 * It is convenient for removing tables that do not need to be
+			 * published later.
 			 */
 			if (publication->pubviaroot)
-				tables = filter_partitions(tables);
+				tables_viaroot = list_concat_unique_oid(tables_viaroot, current_table);
+			else
+				tables = list_concat_unique_oid(tables, current_table);
+
+			/* Reset current_table */
+			list_free(current_table);
+			current_table = NIL;
 		}
 
+		/*
+		 * If a partition table is published in a publication with viaroot,
+		 * and its parent or child table is published in another publication
+		 * without viaroot. Then we need to move the parent or child table
+		 * from tables to tables_viaroot.
+		 *
+		 * If all publication(s)'s viaroot are the same, then skip this part.
+		 */
+		if (tables_viaroot && tables)
+		{
+			ListCell   *lc;
+			ListCell   *lc2;
+			List	   *change_tables = NIL;
+
+			foreach(lc, tables_viaroot)
+			{
+				Oid			relid_viaroot = lfirst_oid(lc);
+				Oid			ancestor_viaroot;
+
+				/*
+				 * If the relid is not a partition table, the relid may be the
+				 * root partitioned table of the relation in the tables, so
+				 * check itself.
+				 */
+				if (get_rel_relispartition(relid_viaroot))
+					ancestor_viaroot = llast_oid(get_partition_ancestors(relid_viaroot));
+				else
+					ancestor_viaroot = relid_viaroot;
+
+				foreach(lc2, tables)
+				{
+					Oid			relid = lfirst_oid(lc2);
+					Oid			ancestor;
+
+					if (get_rel_relispartition(relid))
+						ancestor = llast_oid(get_partition_ancestors(relid));
+					else
+						ancestor = relid;
+
+					if (ancestor_viaroot == ancestor)
+					{
+						tables = foreach_delete_current(tables, lc2);
+						change_tables = list_append_unique_oid(change_tables,
+															   relid);
+					}
+				}
+			}
+
+			if (change_tables)
+				tables_viaroot = list_concat_unique_oid(tables_viaroot, change_tables);
+		}
+
+		if (tables_viaroot)
+			tables = list_concat_unique_oid(tables, filter_partitions(tables_viaroot));
+
 		funcctx->user_fctx = (void *) tables;
 
 		MemoryContextSwitchTo(oldcontext);
diff --git a/src/backend/commands/subscriptioncmds.c b/src/backend/commands/subscriptioncmds.c
index b94236f74d..851eabf3f1 100644
--- a/src/backend/commands/subscriptioncmds.c
+++ b/src/backend/commands/subscriptioncmds.c
@@ -1759,17 +1759,22 @@ static List *
 fetch_table_list(WalReceiverConn *wrconn, List *publications)
 {
 	WalRcvExecResult *res;
-	StringInfoData cmd;
+	StringInfoData cmd,
+				pub_names;
 	TupleTableSlot *slot;
 	Oid			tableRow[2] = {TEXTOID, TEXTOID};
 	List	   *tablelist = NIL;
 
+	initStringInfo(&pub_names);
+	get_publications_str(publications, &pub_names, true);
+
 	initStringInfo(&cmd);
-	appendStringInfoString(&cmd, "SELECT DISTINCT t.schemaname, t.tablename\n"
-						   "  FROM pg_catalog.pg_publication_tables t\n"
-						   " WHERE t.pubname IN (");
-	get_publications_str(publications, &cmd, true);
-	appendStringInfoChar(&cmd, ')');
+	appendStringInfo(&cmd, "SELECT DISTINCT n.nspname, c.relname"
+					 " FROM pg_class c JOIN pg_namespace n"
+					 "   ON n.oid = c.relnamespace,"
+					 " LATERAL pg_get_publication_tables(array[ %s ]) gst"
+					 " WHERE c.oid = gst.relid;",
+					 pub_names.data);
 
 	res = walrcv_exec(wrconn, cmd.data, 2, tableRow);
 	pfree(cmd.data);
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index 6d378ff785..2cd3712365 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -11670,10 +11670,10 @@
   prosrc => 'pg_show_replication_origin_status' },
 
 # publications
-{ oid => '6119', descr => 'get OIDs of tables in a publication',
+{ oid => '6119', descr => 'get OIDs of tables in one or more publications',
   proname => 'pg_get_publication_tables', prorows => '1000', proretset => 't',
-  provolatile => 's', prorettype => 'oid', proargtypes => 'text',
-  proallargtypes => '{text,oid}', proargmodes => '{i,o}',
+  provolatile => 's', prorettype => 'oid', proargtypes => 'any',
+  proallargtypes => '{any,oid}', proargmodes => '{i,o}',
   proargnames => '{pubname,relid}', prosrc => 'pg_get_publication_tables' },
 { oid => '6121',
   descr => 'returns whether a relation can be part of a publication',
diff --git a/src/test/regress/expected/rules.out b/src/test/regress/expected/rules.out
index 21effe8315..35f1f4241d 100644
--- a/src/test/regress/expected/rules.out
+++ b/src/test/regress/expected/rules.out
@@ -1439,7 +1439,7 @@ pg_publication_tables| SELECT p.pubname,
     n.nspname AS schemaname,
     c.relname AS tablename
    FROM pg_publication p,
-    LATERAL pg_get_publication_tables((p.pubname)::text) gpt(relid),
+    LATERAL pg_get_publication_tables(p.pubname) gpt(relid),
     (pg_class c
      JOIN pg_namespace n ON ((n.oid = c.relnamespace)))
   WHERE (c.oid = gpt.relid);
diff --git a/src/test/subscription/t/013_partition.pl b/src/test/subscription/t/013_partition.pl
index 66e63e755e..6313d425dc 100644
--- a/src/test/subscription/t/013_partition.pl
+++ b/src/test/subscription/t/013_partition.pl
@@ -476,14 +476,14 @@ $node_subscriber2->safe_psql('postgres',
 	"CREATE TABLE tab3_1 (a int PRIMARY KEY, c text DEFAULT 'sub2_tab3_1', b text)"
 );
 
-# Note: We create two separate tables, not a partitioned one, so that we can
-# easily identity through which relation were the changes replicated.
+# Note: We only create one table for the partition table (tab4) here.
+# Because we specify option PUBLISH_VIA_PARTITION_ROOT (see pub_all and
+# pub_lower_level above), all data should be replicated to the partition table.
+# So we do not need to create table for the partitioned table.
 $node_subscriber2->safe_psql('postgres',
 	"CREATE TABLE tab4 (a int PRIMARY KEY)"
 );
-$node_subscriber2->safe_psql('postgres',
-	"CREATE TABLE tab4_1 (a int PRIMARY KEY)"
-);
+
 # Publication that sub2 points to now publishes via root, so must update
 # subscription target relations. We set the list of publications so that
 # the FOR ALL TABLES publication is second (the list order matters).
@@ -559,11 +559,6 @@ $result = $node_subscriber2->safe_psql('postgres',
 	"SELECT a FROM tab4 ORDER BY 1");
 is( $result, qq(0), 'inserts into tab4 replicated');
 
-$result = $node_subscriber2->safe_psql('postgres',
-	"SELECT a FROM tab4_1 ORDER BY 1");
-is( $result, qq(), 'inserts into tab4_1 replicated');
-
-
 # now switch the order of publications in the list, try again, the result
 # should be the same (no dependence on order of pulications)
 $node_subscriber2->safe_psql('postgres',
@@ -588,11 +583,6 @@ $result = $node_subscriber2->safe_psql('postgres',
 is( $result, qq(0
 1), 'inserts into tab4 replicated');
 
-$result = $node_subscriber2->safe_psql('postgres',
-	"SELECT a FROM tab4_1 ORDER BY 1");
-is( $result, qq(), 'inserts into tab4_1 replicated');
-
-
 # update (replicated as update)
 $node_publisher->safe_psql('postgres', "UPDATE tab1 SET a = 6 WHERE a = 5");
 $node_publisher->safe_psql('postgres', "UPDATE tab2 SET a = 6 WHERE a = 5");
diff --git a/src/test/subscription/t/028_row_filter.pl b/src/test/subscription/t/028_row_filter.pl
index 82c4eb6ef6..21800a03b6 100644
--- a/src/test/subscription/t/028_row_filter.pl
+++ b/src/test/subscription/t/028_row_filter.pl
@@ -394,6 +394,10 @@ $node_publisher->safe_psql('postgres',
 	"INSERT INTO tab_rowfilter_child(a, b) VALUES(0,'0'),(30,'30'),(40,'40')"
 );
 
+# insert data into partitioned table.
+$node_publisher->safe_psql('postgres',
+	"INSERT INTO tab_rowfilter_viaroot_part(a) VALUES(13), (17)");
+
 $node_subscriber->safe_psql('postgres',
 	"CREATE SUBSCRIPTION tap_sub CONNECTION '$publisher_connstr application_name=$appname' PUBLICATION tap_pub_1, tap_pub_2, tap_pub_3, tap_pub_4a, tap_pub_4b, tap_pub_5a, tap_pub_5b, tap_pub_toast, tap_pub_inherits, tap_pub_viaroot_2, tap_pub_viaroot_1"
 );
@@ -719,8 +723,9 @@ is($result, qq(t|1), 'check replicated rows to tab_rowfilter_toast');
 # - INSERT (16)        YES, 16 > 15
 $result =
   $node_subscriber->safe_psql('postgres',
-	"SELECT a FROM tab_rowfilter_viaroot_part");
-is( $result, qq(16),
+	"SELECT a FROM tab_rowfilter_viaroot_part ORDER BY 1");
+is( $result, qq(16
+17),
 	'check replicated rows to tab_rowfilter_viaroot_part'
 );
 
diff --git a/src/test/subscription/t/031_column_list.pl b/src/test/subscription/t/031_column_list.pl
index bdcf3e4a24..7c73913f4d 100644
--- a/src/test/subscription/t/031_column_list.pl
+++ b/src/test/subscription/t/031_column_list.pl
@@ -948,7 +948,8 @@ $node_publisher->safe_psql('postgres', qq(
 	CREATE TABLE test_root_1 PARTITION OF test_root FOR VALUES FROM (1) TO (10);
 	CREATE TABLE test_root_2 PARTITION OF test_root FOR VALUES FROM (10) TO (20);
 
-	CREATE PUBLICATION pub_root_true FOR TABLE test_root (a) WITH (publish_via_partition_root = true);
+	CREATE PUBLICATION pub_root_true_1 FOR TABLE test_root (a) WITH (publish_via_partition_root = true);
+	CREATE PUBLICATION pub_root_true_2 FOR TABLE test_root_1 (a, b) WITH (publish_via_partition_root = true);
 
 	-- initial data
 	INSERT INTO test_root VALUES (1, 2, 3);
@@ -956,7 +957,7 @@ $node_publisher->safe_psql('postgres', qq(
 ));
 
 $node_subscriber->safe_psql('postgres', qq(
-	CREATE SUBSCRIPTION sub1 CONNECTION '$publisher_connstr' PUBLICATION pub_root_true;
+	CREATE SUBSCRIPTION sub1 CONNECTION '$publisher_connstr' PUBLICATION pub_root_true_1, pub_root_true_2;
 ));
 
 wait_for_subscription_sync($node_subscriber);
-- 
2.18.4