From bddf695e17becbc39bde4593f63a90c64dd97ae3 Mon Sep 17 00:00:00 2001
From: "houzj.fnst" <houzj.fnst@fujitsu.com>
Date: Thu, 16 Sep 2021 09:27:29 +0800
Subject: [PATCH] Invalidate all partitions for partitioned table in publication.

Updates/Deletes on a partition were allowed even without replica identity
after the parent table was added to a publiaction. This would later lead to
an error on subscribers.

The reason was that we were not invalidating the partition's relcache and the
publication information for partitions was not getting rebuilt. Similarly, we
were not invalidating the partitions' relcache after dropping a partitioned
table from publication which will prohibit Updates/Deletes on its partition
without replica identity even without any publication.

---
 src/backend/catalog/pg_publication.c      | 17 ++++++++++++++---
 src/backend/commands/publicationcmds.c    | 14 ++++++++++++--
 src/include/catalog/pg_publication.h      |  3 +++
 src/test/regress/expected/publication.out | 13 +++++++++++--
 src/test/regress/sql/publication.sql      | 11 +++++++++--
 5 files changed, 49 insertions(+), 9 deletions(-)

diff --git a/src/backend/catalog/pg_publication.c b/src/backend/catalog/pg_publication.c
index fa47f7052e..c6ff2d4acf 100644
--- a/src/backend/catalog/pg_publication.c
+++ b/src/backend/catalog/pg_publication.c
@@ -31,6 +31,7 @@
 #include "catalog/pg_publication.h"
 #include "catalog/pg_publication_rel.h"
 #include "catalog/pg_type.h"
+#include "commands/publicationcmds.h"
 #include "funcapi.h"
 #include "miscadmin.h"
 #include "utils/array.h"
@@ -140,7 +141,7 @@ pg_relation_is_publishable(PG_FUNCTION_ARGS)
  * Gets the relations based on the publication partition option for a specified
  * relation.
  */
-static List *
+List *
 GetPubPartitionOptionRelations(List *result, PublicationPartOpt pub_partopt,
 							   Oid relid)
 {
@@ -189,6 +190,7 @@ publication_add_relation(Oid pubid, PublicationRelInfo *targetrel,
 	Publication *pub = GetPublication(pubid);
 	ObjectAddress myself,
 				referenced;
+	List	   *relids = NIL;
 
 	rel = table_open(PublicationRelRelationId, RowExclusiveLock);
 
@@ -244,8 +246,17 @@ publication_add_relation(Oid pubid, PublicationRelInfo *targetrel,
 	/* Close the table. */
 	table_close(rel, RowExclusiveLock);
 
-	/* Invalidate relcache so that publication info is rebuilt. */
-	CacheInvalidateRelcache(targetrel->relation);
+	/*
+	 * Invalidate relcache so that publication info is rebuilt.
+	 *
+	 * For partitioned table contained in the publication, we must
+	 * invalidate all partitions contained in the respective partition
+	 * trees, not just the one explicitly mentioned in the publication.
+	 */
+	relids = GetPubPartitionOptionRelations(relids, PUBLICATION_PART_ALL,
+											relid);
+
+	InvalidatePublicationRels(relids);
 
 	return myself;
 }
diff --git a/src/backend/commands/publicationcmds.c b/src/backend/commands/publicationcmds.c
index 945df49078..89972fa3cc 100644
--- a/src/backend/commands/publicationcmds.c
+++ b/src/backend/commands/publicationcmds.c
@@ -490,6 +490,7 @@ RemovePublicationRelById(Oid proid)
 	Relation	rel;
 	HeapTuple	tup;
 	Form_pg_publication_rel pubrel;
+	List	   *relids = NIL;
 
 	rel = table_open(PublicationRelRelationId, RowExclusiveLock);
 
@@ -501,8 +502,17 @@ RemovePublicationRelById(Oid proid)
 
 	pubrel = (Form_pg_publication_rel) GETSTRUCT(tup);
 
-	/* Invalidate relcache so that publication info is rebuilt. */
-	CacheInvalidateRelcacheByRelid(pubrel->prrelid);
+	/*
+	 * Invalidate relcache so that publication info is rebuilt.
+	 *
+	 * For partitioned table contained in the publication, we must
+	 * invalidate all partitions contained in the respective partition
+	 * trees, not just the one explicitly mentioned in the publication.
+	 */
+	relids = GetPubPartitionOptionRelations(relids, PUBLICATION_PART_ALL,
+											pubrel->prrelid);
+
+	InvalidatePublicationRels(relids);
 
 	CatalogTupleDelete(rel, &tup->t_self);
 
diff --git a/src/include/catalog/pg_publication.h b/src/include/catalog/pg_publication.h
index 561266aa3e..82f2536c65 100644
--- a/src/include/catalog/pg_publication.h
+++ b/src/include/catalog/pg_publication.h
@@ -111,6 +111,9 @@ typedef enum PublicationPartOpt
 extern List *GetPublicationRelations(Oid pubid, PublicationPartOpt pub_partopt);
 extern List *GetAllTablesPublications(void);
 extern List *GetAllTablesPublicationRelations(bool pubviaroot);
+extern List *GetPubPartitionOptionRelations(List *result,
+											PublicationPartOpt pub_partopt,
+											Oid relid);
 
 extern bool is_publishable_relation(Relation rel);
 extern ObjectAddress publication_add_relation(Oid pubid, PublicationRelInfo *targetrel,
diff --git a/src/test/regress/expected/publication.out b/src/test/regress/expected/publication.out
index cad1b374be..82bce9be09 100644
--- a/src/test/regress/expected/publication.out
+++ b/src/test/regress/expected/publication.out
@@ -126,10 +126,12 @@ CREATE PUBLICATION testpub_forparted;
 CREATE PUBLICATION testpub_forparted1;
 RESET client_min_messages;
 CREATE TABLE testpub_parted1 (LIKE testpub_parted);
+CREATE TABLE testpub_parted2 (LIKE testpub_parted);
 ALTER PUBLICATION testpub_forparted1 SET (publish='insert');
+ALTER TABLE testpub_parted ATTACH PARTITION testpub_parted1 FOR VALUES IN (1);
+ALTER TABLE testpub_parted ATTACH PARTITION testpub_parted2 FOR VALUES IN (2);
 -- works despite missing REPLICA IDENTITY, because updates are not replicated
 UPDATE testpub_parted1 SET a = 1;
-ALTER TABLE testpub_parted ATTACH PARTITION testpub_parted1 FOR VALUES IN (1);
 -- only parent is listed as being in publication, not the partition
 ALTER PUBLICATION testpub_forparted ADD TABLE testpub_parted;
 \dRp+ testpub_forparted
@@ -156,7 +158,14 @@ ALTER PUBLICATION testpub_forparted SET (publish_via_partition_root = true);
 Tables:
     "public.testpub_parted"
 
-DROP TABLE testpub_parted1;
+-- still fail, because parent's publication replicates updates
+UPDATE testpub_parted2 SET a = 2;
+ERROR:  cannot update table "testpub_parted2" because it does not have a replica identity and publishes updates
+HINT:  To enable updating the table, set REPLICA IDENTITY using ALTER TABLE.
+ALTER PUBLICATION testpub_forparted DROP TABLE testpub_parted;
+-- works again, because update is no longer replicated
+UPDATE testpub_parted2 SET a = 2;
+DROP TABLE testpub_parted1, testpub_parted2;
 DROP PUBLICATION testpub_forparted, testpub_forparted1;
 -- Test cache invalidation FOR ALL TABLES publication
 SET client_min_messages = 'ERROR';
diff --git a/src/test/regress/sql/publication.sql b/src/test/regress/sql/publication.sql
index 04b34ee299..e5745d575b 100644
--- a/src/test/regress/sql/publication.sql
+++ b/src/test/regress/sql/publication.sql
@@ -76,10 +76,12 @@ CREATE PUBLICATION testpub_forparted;
 CREATE PUBLICATION testpub_forparted1;
 RESET client_min_messages;
 CREATE TABLE testpub_parted1 (LIKE testpub_parted);
+CREATE TABLE testpub_parted2 (LIKE testpub_parted);
 ALTER PUBLICATION testpub_forparted1 SET (publish='insert');
+ALTER TABLE testpub_parted ATTACH PARTITION testpub_parted1 FOR VALUES IN (1);
+ALTER TABLE testpub_parted ATTACH PARTITION testpub_parted2 FOR VALUES IN (2);
 -- works despite missing REPLICA IDENTITY, because updates are not replicated
 UPDATE testpub_parted1 SET a = 1;
-ALTER TABLE testpub_parted ATTACH PARTITION testpub_parted1 FOR VALUES IN (1);
 -- only parent is listed as being in publication, not the partition
 ALTER PUBLICATION testpub_forparted ADD TABLE testpub_parted;
 \dRp+ testpub_forparted
@@ -90,7 +92,12 @@ ALTER TABLE testpub_parted DETACH PARTITION testpub_parted1;
 UPDATE testpub_parted1 SET a = 1;
 ALTER PUBLICATION testpub_forparted SET (publish_via_partition_root = true);
 \dRp+ testpub_forparted
-DROP TABLE testpub_parted1;
+-- still fail, because parent's publication replicates updates
+UPDATE testpub_parted2 SET a = 2;
+ALTER PUBLICATION testpub_forparted DROP TABLE testpub_parted;
+-- works again, because update is no longer replicated
+UPDATE testpub_parted2 SET a = 2;
+DROP TABLE testpub_parted1, testpub_parted2;
 DROP PUBLICATION testpub_forparted, testpub_forparted1;
 
 -- Test cache invalidation FOR ALL TABLES publication
-- 
2.18.4

