From 4a370f0d6586091979512a11797e8cb0b05affce Mon Sep 17 00:00:00 2001
From: "houzj.fnst" <houzj.fnst@fujitsu.com>
Date: Thu, 9 Dec 2021 09:13:50 +0800
Subject: [PATCH] do REPLICA IDENTITY validation when UPDATE or DELETE

For publish mode "delete" "update", validates that any columns referenced
in the filter expression must be part of REPLICA IDENTITY or Primary Key.

Move the row filter columns invalidation to CheckCmdReplicaIdentity, so that
the invalidation is executed only when actual UPDATE or DELETE executed on the
published relation. It's consistent with the existing check about replica
identity and can detect the change related to the row filter in time.

Cache the results of the validation for row filter columns in relcache to
reduce the cost of the validation. It's safe because every operation that
change the row filter and replica identity will invalidate the relcache.

---
 src/backend/catalog/pg_publication.c      | 104 +--------
 src/backend/executor/execReplication.c    |  36 ++-
 src/backend/utils/cache/relcache.c        | 262 +++++++++++++++++++---
 src/include/utils/rel.h                   |   7 +
 src/include/utils/relcache.h              |   1 +
 src/test/regress/expected/publication.out |  94 +++++---
 src/test/regress/sql/publication.sql      |  79 ++++---
 7 files changed, 388 insertions(+), 195 deletions(-)

diff --git a/src/backend/catalog/pg_publication.c b/src/backend/catalog/pg_publication.c
index 1630d2650c..3decb3935c 100644
--- a/src/backend/catalog/pg_publication.c
+++ b/src/backend/catalog/pg_publication.c
@@ -251,13 +251,6 @@ pg_relation_is_publishable(PG_FUNCTION_ARGS)
 	PG_RETURN_BOOL(result);
 }
 
-/* For rowfilter_walker. */
-typedef struct {
-	Relation	rel;
-	bool		check_replident; /* check if Var is bms_replident member? */
-	Bitmapset  *bms_replident; /* bitset of replica identity col indexes */
-} rf_context;
-
 /*
  * The row filter walker checks that the row filter expression is legal.
  *
@@ -291,15 +284,9 @@ typedef struct {
  * We don't allow anything other than immutable built-in functions because those
  * (not immutable ones) can access database and would lead to the problem (b)
  * mentioned in the previous paragraph.
- *
- * Rules: Replica Identity validation
- * -----------------------------------
- * If the flag context.check_replident is true then validate that every variable
- * referenced by the filter expression is a valid member of the allowed set of
- * replica identity columns (context.bms_replindent)
  */
 static bool
-rowfilter_walker(Node *node, rf_context *context)
+rowfilter_walker(Node *node, Relation relation)
 {
 	char *forbidden = NULL;
 	bool too_complex = false;
@@ -314,25 +301,6 @@ rowfilter_walker(Node *node, rf_context *context)
 		/* User-defined types not allowed. */
 		if (var->vartype >= FirstNormalObjectId)
 			forbidden = _("user-defined types are not allowed");
-
-		/* Optionally, do replica identify validation of the referenced column. */
-		if (context->check_replident)
-		{
-			Oid			relid = RelationGetRelid(context->rel);
-			AttrNumber	attnum = var->varattno;
-
-			if (!bms_is_member(attnum - FirstLowInvalidHeapAttributeNumber, context->bms_replident))
-			{
-				const char *colname = get_attname(relid, attnum, false);
-
-				ereport(ERROR,
-						(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
-						errmsg("cannot add relation \"%s\" to publication",
-							   RelationGetRelationName(context->rel)),
-						errdetail("Row filter column \"%s\" is not part of the REPLICA IDENTITY",
-								  colname)));
-			}
-		}
 	}
 	else if (IsA(node, Const) || IsA(node, BoolExpr) || IsA(node, NullIfExpr)
 			 || IsA(node, NullTest) || IsA(node, BooleanTest))
@@ -375,74 +343,18 @@ rowfilter_walker(Node *node, rf_context *context)
 	if (too_complex)
 		ereport(ERROR,
 				(errmsg("invalid publication WHERE expression for relation \"%s\"",
-						RelationGetRelationName(context->rel)),
+						RelationGetRelationName(relation)),
 				errhint("only simple expressions using columns, constants and immutable system functions are allowed")
 				));
 
 	if (forbidden)
 		ereport(ERROR,
 				(errmsg("invalid publication WHERE expression for relation \"%s\"",
-						RelationGetRelationName(context->rel)),
+						RelationGetRelationName(relation)),
 						errdetail("%s", forbidden)
 				));
 
-	return expression_tree_walker(node, rowfilter_walker, (void *)context);
-}
-
-/*
- * Check if the row-filter is valid according to the following rules:
- *
- * 1. Only certain simple node types are permitted in the expression. See
- * function rowfilter_walker for details.
- *
- * 2. If the publish operation contains "delete" or "update" then only columns
- * that are allowed by the REPLICA IDENTITY rules are permitted to be used in
- * the row-filter WHERE clause.
- */
-static void
-rowfilter_expr_checker(Publication *pub, Relation rel, Node *rfnode)
-{
-	rf_context	context = {0};
-
-	context.rel = rel;
-
-	/*
-	 * For "delete" or "update", check that filter cols are also valid replica
-	 * identity cols.
-	 */
-	if (pub->pubactions.pubdelete || pub->pubactions.pubupdate)
-	{
-		char replica_identity = rel->rd_rel->relreplident;
-
-		if (replica_identity == REPLICA_IDENTITY_FULL)
-		{
-			/*
-			 * FULL means all cols are in the REPLICA IDENTITY, so all cols are
-			 * allowed in the row-filter too.
-			 */
-		}
-		else
-		{
-			context.check_replident = true;
-
-			/*
-			 * Find what are the cols that are part of the REPLICA IDENTITY.
-			 * Note that REPLICA IDENTIY DEFAULT means primary key or nothing.
-			 */
-			if (replica_identity == REPLICA_IDENTITY_DEFAULT)
-				context.bms_replident = RelationGetIndexAttrBitmap(rel, INDEX_ATTR_BITMAP_PRIMARY_KEY);
-			else
-				context.bms_replident = RelationGetIndexAttrBitmap(rel, INDEX_ATTR_BITMAP_IDENTITY_KEY);
-		}
-	}
-
-	/*
-	 * Walk the parse-tree of this publication row filter expression and throw an
-	 * error if anything not permitted or unexpected is encountered.
-	 */
-	rowfilter_walker(rfnode, &context);
-
-	bms_free(context.bms_replident);
+	return expression_tree_walker(node, rowfilter_walker, (void *) relation);
 }
 
 List *
@@ -556,8 +468,12 @@ publication_add_relation(Oid pubid, PublicationRelInfo *pri,
 		/* Fix up collation information */
 		whereclause = GetTransformedWhereClause(pstate, pri, true);
 
-		/* Validate the row-filter. */
-		rowfilter_expr_checker(pub, targetrel, whereclause);
+		/*
+		 * Walk the parse-tree of this publication row filter expression and
+		 * throw an error if anything not permitted or unexpected is
+		 * encountered.
+		 */
+		rowfilter_walker(whereclause, targetrel);
 	}
 
 	/* Form a tuple. */
diff --git a/src/backend/executor/execReplication.c b/src/backend/executor/execReplication.c
index 574d7d27fd..c069863d56 100644
--- a/src/backend/executor/execReplication.c
+++ b/src/backend/executor/execReplication.c
@@ -568,14 +568,46 @@ void
 CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 {
 	PublicationActions *pubactions;
+	AttrNumber			invalid_rfcolnum;
 
 	/* We only need to do checks for UPDATE and DELETE. */
 	if (cmd != CMD_UPDATE && cmd != CMD_DELETE)
 		return;
 
+	if (rel->rd_rel->relreplident == REPLICA_IDENTITY_FULL)
+		return;
+
+	invalid_rfcolnum = RelationGetInvalidRowFilterCol(rel);
+
+	/*
+	 * It is only safe to execute UPDATE/DELETE when all columns referenced in
+	 * the row filters from publications which the relation is in are valid,
+	 * which means all referenced columns are part of REPLICA IDENTITY, or the
+	 * table do not publish UPDATES or DELETES.
+	 */
+	if (invalid_rfcolnum)
+	{
+		const char *colname = get_attname(RelationGetRelid(rel),
+										  invalid_rfcolnum, false);
+
+		if (cmd == CMD_UPDATE)
+			ereport(ERROR,
+					(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+					 errmsg("cannot update table \"%s\"",
+							RelationGetRelationName(rel)),
+					 errdetail("Row filter column \"%s\" is not part of the REPLICA IDENTITY",
+							   colname)));
+		else if (cmd == CMD_DELETE)
+			ereport(ERROR,
+					(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+					 errmsg("cannot delete from table \"%s\"",
+							RelationGetRelationName(rel)),
+					 errdetail("Row filter column \"%s\" is not part of the REPLICA IDENTITY",
+							   colname)));
+	}
+
 	/* If relation has replica identity we are always good. */
-	if (rel->rd_rel->relreplident == REPLICA_IDENTITY_FULL ||
-		OidIsValid(RelationGetReplicaIndex(rel)))
+	if (OidIsValid(RelationGetReplicaIndex(rel)))
 		return;
 
 	/*
diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
index 105d8d4601..d0c1b3d316 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -56,6 +56,7 @@
 #include "catalog/pg_opclass.h"
 #include "catalog/pg_proc.h"
 #include "catalog/pg_publication.h"
+#include "catalog/pg_publication_rel.h"
 #include "catalog/pg_rewrite.h"
 #include "catalog/pg_shseclabel.h"
 #include "catalog/pg_statistic_ext.h"
@@ -71,6 +72,8 @@
 #include "nodes/makefuncs.h"
 #include "nodes/nodeFuncs.h"
 #include "optimizer/optimizer.h"
+#include "parser/parse_clause.h"
+#include "parser/parse_relation.h"
 #include "rewrite/rewriteDefine.h"
 #include "rewrite/rowsecurity.h"
 #include "storage/lmgr.h"
@@ -84,6 +87,7 @@
 #include "utils/memutils.h"
 #include "utils/relmapper.h"
 #include "utils/resowner_private.h"
+#include "utils/ruleutils.h"
 #include "utils/snapmgr.h"
 #include "utils/syscache.h"
 
@@ -5521,57 +5525,169 @@ RelationGetExclusionInfo(Relation indexRelation,
 	MemoryContextSwitchTo(oldcxt);
 }
 
+/* For invalid_rowfilter_column_walker. */
+typedef struct {
+	AttrNumber	invalid_rfcolnum; /* invalid column number */
+	Bitmapset  *bms_replident;	/* bitset of replica identity col indexes */
+	bool		pubviaroot;		/* true if we are validating the parent
+								 * relation's row filter */
+	Oid			relid;			/* relid of the relation */
+	Oid			parentid;		/* relid of the parent relation */
+} rf_context;
+
 /*
- * Get publication actions for the given relation.
+ * Check if any columns used in the row-filter WHERE clause are not part of
+ * REPLICA IDENTITY and save the invalid column number in
+ * rf_context::invalid_rfcolnum.
  */
-struct PublicationActions *
-GetRelationPublicationActions(Relation relation)
+static bool
+invalid_rowfilter_column_walker(Node *node, rf_context *context)
 {
-	List	   *puboids;
-	ListCell   *lc;
-	MemoryContext oldcxt;
-	Oid			schemaid;
-	PublicationActions *pubactions = palloc0(sizeof(PublicationActions));
+	if (node == NULL)
+		return false;
+
+	if (IsA(node, Var))
+	{
+		Var		   *var = (Var *) node;
+		AttrNumber	attnum = var->varattno;
+
+		/*
+		 * If pubviaroot is true, we need to convert the column number of
+		 * parent to the column number of child relation first.
+		 */
+		if (context->pubviaroot)
+		{
+			char *colname = get_attname(context->parentid, attnum, false);
+			attnum = get_attnum(context->relid, colname);
+		}
+
+		if (!bms_is_member(attnum - FirstLowInvalidHeapAttributeNumber,
+						   context->bms_replident))
+		{
+			context->invalid_rfcolnum = attnum;
+			return true;
+		}
+	}
+
+	return expression_tree_walker(node, invalid_rowfilter_column_walker,
+								  (void *) context);
+}
+
+/*
+ * Append to cur_puboids each member of add_puboids that isn't already in
+ * cur_puboids.
+ *
+ * Also update the top most parent relation's relid in the publication.
+ */
+static void
+concat_publication_oid(Oid relid,
+					   List **cur_puboids,
+					   List **toprelid_in_pub,
+					   const List *add_puboids)
+{
+	ListCell   *lc1,
+			   *lc2,
+			   *lc3;
+
+	foreach(lc1, add_puboids)
+	{
+		bool		is_member = false;
+
+		forboth(lc2, *cur_puboids, lc3, *toprelid_in_pub)
+		{
+			if (lfirst_oid(lc2) == lfirst_oid(lc1))
+			{
+				is_member = true;
+				lfirst_oid(lc3) = relid;
+			}
+		}
+
+		if (!is_member)
+		{
+			*cur_puboids = lappend_oid(*cur_puboids, lfirst_oid(lc1));
+			*toprelid_in_pub = lappend_oid(*toprelid_in_pub, relid);
+		}
+	}
+}
+
+/*
+ * Get the invalid row filter column number for the given relation.
+ *
+ * Traverse all the publications which the relation is in to get the
+ * publication actions. If the publication actions include UPDATE or DELETE,
+ * then validate that if all columns referenced in the row filter expression
+ * are part of REPLICA IDENTITY.
+ *
+ * If not all the row filter columns are part of REPLICA IDENTITY, return the
+ * invalid column number, otherwise InvalidAttrNumber.
+ */
+AttrNumber
+RelationGetInvalidRowFilterCol(Relation relation)
+{
+	List		   *puboids,
+				   *toprelid_in_pub;
+	ListCell	   *lc;
+	MemoryContext	oldcxt;
+	Oid				schemaid;
+	Oid				relid = RelationGetRelid(relation);
+	rf_context		context = { 0 };
+	PublicationActions pubactions = { 0 };
+	bool			rfcol_valid = true;
+	AttrNumber		invalid_rfcolnum = InvalidAttrNumber;
 
 	/*
 	 * If not publishable, it publishes no actions.  (pgoutput_change() will
 	 * ignore it.)
 	 */
-	if (!is_publishable_relation(relation))
-		return pubactions;
-
-	if (relation->rd_pubactions)
-		return memcpy(pubactions, relation->rd_pubactions,
-					  sizeof(PublicationActions));
+	if (!is_publishable_relation(relation) || relation->rd_rfcol_valid)
+		return invalid_rfcolnum;
 
 	/* Fetch the publication membership info. */
-	puboids = GetRelationPublications(RelationGetRelid(relation));
+	toprelid_in_pub = puboids = NIL;
+	concat_publication_oid(relid, &puboids, &toprelid_in_pub,
+						   GetRelationPublications(relid));
 	schemaid = RelationGetNamespace(relation);
-	puboids = list_concat_unique_oid(puboids, GetSchemaPublications(schemaid));
+	concat_publication_oid(relid, &puboids, &toprelid_in_pub,
+						   GetSchemaPublications(schemaid));
 
 	if (relation->rd_rel->relispartition)
 	{
 		/* Add publications that the ancestors are in too. */
-		List	   *ancestors = get_partition_ancestors(RelationGetRelid(relation));
+		List	   *ancestors = get_partition_ancestors(relid);
 		ListCell   *lc;
 
 		foreach(lc, ancestors)
 		{
 			Oid			ancestor = lfirst_oid(lc);
 
-			puboids = list_concat_unique_oid(puboids,
-											 GetRelationPublications(ancestor));
+			concat_publication_oid(ancestor, &puboids, &toprelid_in_pub,
+								   GetRelationPublications(ancestor));
 			schemaid = get_rel_namespace(ancestor);
-			puboids = list_concat_unique_oid(puboids,
-											 GetSchemaPublications(schemaid));
+			concat_publication_oid(ancestor, &puboids, &toprelid_in_pub,
+								   GetSchemaPublications(schemaid));
 		}
+
+		relid = llast_oid(ancestors);
 	}
-	puboids = list_concat_unique_oid(puboids, GetAllTablesPublications());
+	concat_publication_oid(relid, &puboids, &toprelid_in_pub,
+						   GetAllTablesPublications());
+
+	/*
+	 * Find what are the cols that are part of the REPLICA IDENTITY.
+	 * Note that REPLICA IDENTITY DEFAULT means primary key or nothing.
+	 */
+	if (relation->rd_rel->relreplident == REPLICA_IDENTITY_DEFAULT)
+		context.bms_replident = RelationGetIndexAttrBitmap(relation,
+														   INDEX_ATTR_BITMAP_PRIMARY_KEY);
+	else if (relation->rd_rel->relreplident == REPLICA_IDENTITY_INDEX)
+		context.bms_replident = RelationGetIndexAttrBitmap(relation,
+														   INDEX_ATTR_BITMAP_IDENTITY_KEY);
 
 	foreach(lc, puboids)
 	{
 		Oid			pubid = lfirst_oid(lc);
 		HeapTuple	tup;
+
 		Form_pg_publication pubform;
 
 		tup = SearchSysCache1(PUBLICATIONOID, ObjectIdGetDatum(pubid));
@@ -5581,35 +5697,116 @@ GetRelationPublicationActions(Relation relation)
 
 		pubform = (Form_pg_publication) GETSTRUCT(tup);
 
-		pubactions->pubinsert |= pubform->pubinsert;
-		pubactions->pubupdate |= pubform->pubupdate;
-		pubactions->pubdelete |= pubform->pubdelete;
-		pubactions->pubtruncate |= pubform->pubtruncate;
+		pubactions.pubinsert |= pubform->pubinsert;
+		pubactions.pubupdate |= pubform->pubupdate;
+		pubactions.pubdelete |= pubform->pubdelete;
+		pubactions.pubtruncate |= pubform->pubtruncate;
 
 		ReleaseSysCache(tup);
 
 		/*
-		 * If we know everything is replicated, there is no point to check for
-		 * other publications.
+		 * If the publication action include UPDATE and DELETE, validates
+		 * that any columns referenced in the filter expression are part of
+		 * REPLICA IDENTITY index.
+		 *
+		 * FULL means all cols are in the REPLICA IDENTITY, so all cols are
+		 * allowed in the row-filter and we can skip the validation.
+		 *
+		 * If we already found the column in row filter which is not part
+		 * of REPLICA IDENTITY index, skip the validation too.
 		 */
-		if (pubactions->pubinsert && pubactions->pubupdate &&
-			pubactions->pubdelete && pubactions->pubtruncate)
+		if ((pubform->pubupdate || pubform->pubdelete) &&
+			relation->rd_rel->relreplident != REPLICA_IDENTITY_FULL &&
+			rfcol_valid)
+		{
+			HeapTuple	rftuple;
+
+			if (pubform->pubviaroot)
+				relid = list_nth_oid(toprelid_in_pub,
+									 foreach_current_index(lc));
+			else
+				relid = RelationGetRelid(relation);
+
+			rftuple = SearchSysCache2(PUBLICATIONRELMAP,
+									  ObjectIdGetDatum(relid),
+									  ObjectIdGetDatum(pubid));
+
+			if (HeapTupleIsValid(rftuple))
+			{
+				Datum		rfdatum;
+				bool		rfisnull;
+				Node	   *rfnode;
+
+				context.pubviaroot = pubform->pubviaroot;
+				context.parentid = relid;
+				context.relid = RelationGetRelid(relation);
+
+				rfdatum = SysCacheGetAttr(PUBLICATIONRELMAP, rftuple,
+										  Anum_pg_publication_rel_prqual,
+										  &rfisnull);
+
+				if (!rfisnull)
+				{
+					rfnode = stringToNode(TextDatumGetCString(rfdatum));
+					rfcol_valid = !invalid_rowfilter_column_walker(rfnode,
+																   &context);
+					invalid_rfcolnum = context.invalid_rfcolnum;
+					pfree(rfnode);
+				}
+
+				ReleaseSysCache(rftuple);
+			}
+		}
+
+		/*
+		 * If we know everything is replicated and some columns are not part of
+		 * replica identity, there is no point to check for other publications.
+		 */
+		if (pubactions.pubinsert && pubactions.pubupdate &&
+			pubactions.pubdelete && pubactions.pubtruncate &&
+			!rfcol_valid)
 			break;
 	}
 
+	bms_free(context.bms_replident);
+
 	if (relation->rd_pubactions)
 	{
 		pfree(relation->rd_pubactions);
 		relation->rd_pubactions = NULL;
 	}
 
+	relation->rd_rfcol_valid = rfcol_valid;
+
 	/* Now save copy of the actions in the relcache entry. */
 	oldcxt = MemoryContextSwitchTo(CacheMemoryContext);
 	relation->rd_pubactions = palloc(sizeof(PublicationActions));
-	memcpy(relation->rd_pubactions, pubactions, sizeof(PublicationActions));
+	memcpy(relation->rd_pubactions, &pubactions, sizeof(PublicationActions));
 	MemoryContextSwitchTo(oldcxt);
 
-	return pubactions;
+	return invalid_rfcolnum;
+}
+
+/*
+ * Get publication actions for the given relation.
+ */
+struct PublicationActions *
+GetRelationPublicationActions(Relation relation)
+{
+	PublicationActions *pubactions = palloc0(sizeof(PublicationActions));
+
+	/*
+	 * If not publishable, it publishes no actions.  (pgoutput_change() will
+	 * ignore it.)
+	 */
+	if (!is_publishable_relation(relation))
+		return pubactions;
+
+	if (!relation->rd_pubactions)
+		(void) RelationGetInvalidRowFilterCol(relation);
+
+	return memcpy(pubactions, relation->rd_pubactions,
+				  sizeof(PublicationActions));
 }
 
 /*
@@ -6163,6 +6360,7 @@ load_relcache_init_file(bool shared)
 		rel->rd_idattr = NULL;
 		rel->rd_hotblockingattr = NULL;
 		rel->rd_pubactions = NULL;
+		rel->rd_rfcol_valid = false;
 		rel->rd_statvalid = false;
 		rel->rd_statlist = NIL;
 		rel->rd_fkeyvalid = false;
diff --git a/src/include/utils/rel.h b/src/include/utils/rel.h
index 31281279cf..27cec813c0 100644
--- a/src/include/utils/rel.h
+++ b/src/include/utils/rel.h
@@ -163,6 +163,13 @@ typedef struct RelationData
 
 	PublicationActions *rd_pubactions;	/* publication actions */
 
+	/*
+	 * true if the columns referenced in row filters from all the publications
+	 * the relation is in are part of replica identity, or the publication
+	 * actions do not include UPDATE and DELETE.
+	 */
+	bool		rd_rfcol_valid;
+
 	/*
 	 * rd_options is set whenever rd_rel is loaded into the relcache entry.
 	 * Note that you can NOT look into rd_rel for this data.  NULL means "use
diff --git a/src/include/utils/relcache.h b/src/include/utils/relcache.h
index 82316bba54..25c759f289 100644
--- a/src/include/utils/relcache.h
+++ b/src/include/utils/relcache.h
@@ -76,6 +76,7 @@ extern void RelationInitIndexAccessInfo(Relation relation);
 /* caller must include pg_publication.h */
 struct PublicationActions;
 extern struct PublicationActions *GetRelationPublicationActions(Relation relation);
+extern AttrNumber RelationGetInvalidRowFilterCol(Relation relation);
 
 extern void RelationInitTableAccessMethod(Relation relation);
 
diff --git a/src/test/regress/expected/publication.out b/src/test/regress/expected/publication.out
index a7729758af..2ca45d8f2c 100644
--- a/src/test/regress/expected/publication.out
+++ b/src/test/regress/expected/publication.out
@@ -416,58 +416,61 @@ DROP FUNCTION testpub_rf_func99();
 -- More row filter tests for validating column references
 CREATE TABLE rf_tbl_abcd_nopk(a int, b int, c int, d int);
 CREATE TABLE rf_tbl_abcd_pk(a int, b int, c int, d int, PRIMARY KEY(a,b));
+create table rf_tbl_abcd_part_pk (a int primary key, b int) partition by RANGE (a);
+create table rf_tbl_abcd_part_pk_1 (b int, a int primary key);
+alter table rf_tbl_abcd_part_pk attach partition rf_tbl_abcd_part_pk_1 FOR VALUES FROM (1) TO (10);
 -- Case 1. REPLICA IDENTITY DEFAULT (means use primary key or nothing)
 -- 1a. REPLICA IDENTITY is DEFAULT and table has a PK.
--- ok - "a" is a PK col
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_pk WHERE (a > 99);
 RESET client_min_messages;
-DROP PUBLICATION testpub6;
+-- ok - "a" is a PK col
+UPDATE rf_tbl_abcd_pk SET a = 1;
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_pk WHERE (b > 99);
 -- ok - "b" is a PK col
-SET client_min_messages = 'ERROR';
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_pk WHERE (b > 99);
-RESET client_min_messages;
-DROP PUBLICATION testpub6;
+UPDATE rf_tbl_abcd_pk SET a = 1;
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_pk WHERE (c > 99);
 -- fail - "c" is not part of the PK
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_pk WHERE (c > 99);
-ERROR:  cannot add relation "rf_tbl_abcd_pk" to publication
+UPDATE rf_tbl_abcd_pk SET a = 1;
+ERROR:  cannot update table "rf_tbl_abcd_pk"
 DETAIL:  Row filter column "c" is not part of the REPLICA IDENTITY
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_pk WHERE (d > 99);
 -- fail - "d" is not part of the PK
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_pk WHERE (d > 99);
-ERROR:  cannot add relation "rf_tbl_abcd_pk" to publication
+UPDATE rf_tbl_abcd_pk SET a = 1;
+ERROR:  cannot update table "rf_tbl_abcd_pk"
 DETAIL:  Row filter column "d" is not part of the REPLICA IDENTITY
 -- 1b. REPLICA IDENTITY is DEFAULT and table has no PK
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_nopk WHERE (a > 99);
 -- fail - "a" is not part of REPLICA IDENTITY
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_nopk WHERE (a > 99);
-ERROR:  cannot add relation "rf_tbl_abcd_nopk" to publication
+UPDATE rf_tbl_abcd_nopk SET a = 1;
+ERROR:  cannot update table "rf_tbl_abcd_nopk"
 DETAIL:  Row filter column "a" is not part of the REPLICA IDENTITY
 -- Case 2. REPLICA IDENTITY FULL
 ALTER TABLE rf_tbl_abcd_pk REPLICA IDENTITY FULL;
 ALTER TABLE rf_tbl_abcd_nopk REPLICA IDENTITY FULL;
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_pk WHERE (c > 99);
 -- ok - "c" is in REPLICA IDENTITY now even though not in PK
-SET client_min_messages = 'ERROR';
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_pk WHERE (c > 99);
-DROP PUBLICATION testpub6;
-RESET client_min_messages;
+UPDATE rf_tbl_abcd_pk SET a = 1;
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_nopk WHERE (a > 99);
 -- ok - "a" is in REPLICA IDENTITY now
-SET client_min_messages = 'ERROR';
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_nopk WHERE (a > 99);
-DROP PUBLICATION testpub6;
-RESET client_min_messages;
+UPDATE rf_tbl_abcd_nopk SET a = 1;
 -- Case 3. REPLICA IDENTITY NOTHING
 ALTER TABLE rf_tbl_abcd_pk REPLICA IDENTITY NOTHING;
 ALTER TABLE rf_tbl_abcd_nopk REPLICA IDENTITY NOTHING;
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_pk WHERE (a > 99);
 -- fail - "a" is in PK but it is not part of REPLICA IDENTITY NOTHING
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_pk WHERE (a > 99);
-ERROR:  cannot add relation "rf_tbl_abcd_pk" to publication
+UPDATE rf_tbl_abcd_pk SET a = 1;
+ERROR:  cannot update table "rf_tbl_abcd_pk"
 DETAIL:  Row filter column "a" is not part of the REPLICA IDENTITY
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_pk WHERE (c > 99);
 -- fail - "c" is not in PK and not in REPLICA IDENTITY NOTHING
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_pk WHERE (c > 99);
-ERROR:  cannot add relation "rf_tbl_abcd_pk" to publication
+UPDATE rf_tbl_abcd_pk SET a = 1;
+ERROR:  cannot update table "rf_tbl_abcd_pk"
 DETAIL:  Row filter column "c" is not part of the REPLICA IDENTITY
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_nopk WHERE (a > 99);
 -- fail - "a" is not in REPLICA IDENTITY NOTHING
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_nopk WHERE (a > 99);
-ERROR:  cannot add relation "rf_tbl_abcd_nopk" to publication
+UPDATE rf_tbl_abcd_nopk SET a = 1;
+ERROR:  cannot update table "rf_tbl_abcd_nopk"
 DETAIL:  Row filter column "a" is not part of the REPLICA IDENTITY
 -- Case 4. REPLICA IDENTITY INDEX
 ALTER TABLE rf_tbl_abcd_pk ALTER COLUMN c SET NOT NULL;
@@ -476,26 +479,43 @@ ALTER TABLE rf_tbl_abcd_pk REPLICA IDENTITY USING INDEX idx_abcd_pk_c;
 ALTER TABLE rf_tbl_abcd_nopk ALTER COLUMN c SET NOT NULL;
 CREATE UNIQUE INDEX idx_abcd_nopk_c ON rf_tbl_abcd_nopk(c);
 ALTER TABLE rf_tbl_abcd_nopk REPLICA IDENTITY USING INDEX idx_abcd_nopk_c;
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_pk WHERE (a > 99);
 -- fail - "a" is in PK but it is not part of REPLICA IDENTITY INDEX
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_pk WHERE (a > 99);
-ERROR:  cannot add relation "rf_tbl_abcd_pk" to publication
+UPDATE rf_tbl_abcd_pk SET a = 1;
+ERROR:  cannot update table "rf_tbl_abcd_pk"
 DETAIL:  Row filter column "a" is not part of the REPLICA IDENTITY
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_pk WHERE (c > 99);
 -- ok - "c" is not in PK but it is part of REPLICA IDENTITY INDEX
-SET client_min_messages = 'ERROR';
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_pk WHERE (c > 99);
-DROP PUBLICATION testpub6;
-RESET client_min_messages;
+UPDATE rf_tbl_abcd_pk SET a = 1;
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_nopk WHERE (a > 99);
 -- fail - "a" is not in REPLICA IDENTITY INDEX
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_nopk WHERE (a > 99);
-ERROR:  cannot add relation "rf_tbl_abcd_nopk" to publication
+UPDATE rf_tbl_abcd_nopk SET a = 1;
+ERROR:  cannot update table "rf_tbl_abcd_nopk"
 DETAIL:  Row filter column "a" is not part of the REPLICA IDENTITY
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_nopk WHERE (c > 99);
 -- ok - "c" is part of REPLICA IDENTITY INDEX
-SET client_min_messages = 'ERROR';
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_nopk WHERE (c > 99);
+UPDATE rf_tbl_abcd_nopk SET a = 1;
+-- Tests for partitioned table
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_part_pk WHERE (a > 99);
+ALTER PUBLICATION testpub6 SET (PUBLISH_VIA_PARTITION_ROOT=0);
+-- ok - partition does not have row filer
+UPDATE rf_tbl_abcd_part_pk SET a = 1;
+ALTER PUBLICATION testpub6 SET (PUBLISH_VIA_PARTITION_ROOT=1);
+-- ok - "a" is a OK col
+UPDATE rf_tbl_abcd_part_pk SET a = 1;
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_part_pk WHERE (b > 99);
+ALTER PUBLICATION testpub6 SET (PUBLISH_VIA_PARTITION_ROOT=0);
+-- ok - partition does not have row filer
+UPDATE rf_tbl_abcd_part_pk SET a = 1;
+ALTER PUBLICATION testpub6 SET (PUBLISH_VIA_PARTITION_ROOT=1);
+-- fail - "b" is not in REPLICA IDENTITY INDEX
+UPDATE rf_tbl_abcd_part_pk SET a = 1;
+ERROR:  cannot update table "rf_tbl_abcd_part_pk_1"
+DETAIL:  Row filter column "b" is not part of the REPLICA IDENTITY
 DROP PUBLICATION testpub6;
-RESET client_min_messages;
 DROP TABLE rf_tbl_abcd_pk;
 DROP TABLE rf_tbl_abcd_nopk;
+DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
 -- 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 e8242c95ee..4278bee6fd 100644
--- a/src/test/regress/sql/publication.sql
+++ b/src/test/regress/sql/publication.sql
@@ -235,50 +235,53 @@ DROP FUNCTION testpub_rf_func99();
 -- More row filter tests for validating column references
 CREATE TABLE rf_tbl_abcd_nopk(a int, b int, c int, d int);
 CREATE TABLE rf_tbl_abcd_pk(a int, b int, c int, d int, PRIMARY KEY(a,b));
+create table rf_tbl_abcd_part_pk (a int primary key, b int) partition by RANGE (a);
+create table rf_tbl_abcd_part_pk_1 (b int, a int primary key);
+alter table rf_tbl_abcd_part_pk attach partition rf_tbl_abcd_part_pk_1 FOR VALUES FROM (1) TO (10);
 
 -- Case 1. REPLICA IDENTITY DEFAULT (means use primary key or nothing)
 -- 1a. REPLICA IDENTITY is DEFAULT and table has a PK.
--- ok - "a" is a PK col
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_pk WHERE (a > 99);
 RESET client_min_messages;
-DROP PUBLICATION testpub6;
+-- ok - "a" is a PK col
+UPDATE rf_tbl_abcd_pk SET a = 1;
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_pk WHERE (b > 99);
 -- ok - "b" is a PK col
-SET client_min_messages = 'ERROR';
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_pk WHERE (b > 99);
-RESET client_min_messages;
-DROP PUBLICATION testpub6;
+UPDATE rf_tbl_abcd_pk SET a = 1;
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_pk WHERE (c > 99);
 -- fail - "c" is not part of the PK
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_pk WHERE (c > 99);
+UPDATE rf_tbl_abcd_pk SET a = 1;
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_pk WHERE (d > 99);
 -- fail - "d" is not part of the PK
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_pk WHERE (d > 99);
+UPDATE rf_tbl_abcd_pk SET a = 1;
 -- 1b. REPLICA IDENTITY is DEFAULT and table has no PK
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_nopk WHERE (a > 99);
 -- fail - "a" is not part of REPLICA IDENTITY
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_nopk WHERE (a > 99);
+UPDATE rf_tbl_abcd_nopk SET a = 1;
 
 -- Case 2. REPLICA IDENTITY FULL
 ALTER TABLE rf_tbl_abcd_pk REPLICA IDENTITY FULL;
 ALTER TABLE rf_tbl_abcd_nopk REPLICA IDENTITY FULL;
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_pk WHERE (c > 99);
 -- ok - "c" is in REPLICA IDENTITY now even though not in PK
-SET client_min_messages = 'ERROR';
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_pk WHERE (c > 99);
-DROP PUBLICATION testpub6;
-RESET client_min_messages;
+UPDATE rf_tbl_abcd_pk SET a = 1;
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_nopk WHERE (a > 99);
 -- ok - "a" is in REPLICA IDENTITY now
-SET client_min_messages = 'ERROR';
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_nopk WHERE (a > 99);
-DROP PUBLICATION testpub6;
-RESET client_min_messages;
+UPDATE rf_tbl_abcd_nopk SET a = 1;
 
 -- Case 3. REPLICA IDENTITY NOTHING
 ALTER TABLE rf_tbl_abcd_pk REPLICA IDENTITY NOTHING;
 ALTER TABLE rf_tbl_abcd_nopk REPLICA IDENTITY NOTHING;
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_pk WHERE (a > 99);
 -- fail - "a" is in PK but it is not part of REPLICA IDENTITY NOTHING
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_pk WHERE (a > 99);
+UPDATE rf_tbl_abcd_pk SET a = 1;
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_pk WHERE (c > 99);
 -- fail - "c" is not in PK and not in REPLICA IDENTITY NOTHING
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_pk WHERE (c > 99);
+UPDATE rf_tbl_abcd_pk SET a = 1;
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_nopk WHERE (a > 99);
 -- fail - "a" is not in REPLICA IDENTITY NOTHING
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_nopk WHERE (a > 99);
+UPDATE rf_tbl_abcd_nopk SET a = 1;
 
 -- Case 4. REPLICA IDENTITY INDEX
 ALTER TABLE rf_tbl_abcd_pk ALTER COLUMN c SET NOT NULL;
@@ -287,23 +290,39 @@ ALTER TABLE rf_tbl_abcd_pk REPLICA IDENTITY USING INDEX idx_abcd_pk_c;
 ALTER TABLE rf_tbl_abcd_nopk ALTER COLUMN c SET NOT NULL;
 CREATE UNIQUE INDEX idx_abcd_nopk_c ON rf_tbl_abcd_nopk(c);
 ALTER TABLE rf_tbl_abcd_nopk REPLICA IDENTITY USING INDEX idx_abcd_nopk_c;
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_pk WHERE (a > 99);
 -- fail - "a" is in PK but it is not part of REPLICA IDENTITY INDEX
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_pk WHERE (a > 99);
+UPDATE rf_tbl_abcd_pk SET a = 1;
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_pk WHERE (c > 99);
 -- ok - "c" is not in PK but it is part of REPLICA IDENTITY INDEX
-SET client_min_messages = 'ERROR';
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_pk WHERE (c > 99);
-DROP PUBLICATION testpub6;
-RESET client_min_messages;
+UPDATE rf_tbl_abcd_pk SET a = 1;
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_nopk WHERE (a > 99);
 -- fail - "a" is not in REPLICA IDENTITY INDEX
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_nopk WHERE (a > 99);
+UPDATE rf_tbl_abcd_nopk SET a = 1;
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_nopk WHERE (c > 99);
 -- ok - "c" is part of REPLICA IDENTITY INDEX
-SET client_min_messages = 'ERROR';
-CREATE PUBLICATION testpub6 FOR TABLE rf_tbl_abcd_nopk WHERE (c > 99);
-DROP PUBLICATION testpub6;
-RESET client_min_messages;
+UPDATE rf_tbl_abcd_nopk SET a = 1;
+
+-- Tests for partitioned table
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_part_pk WHERE (a > 99);
+ALTER PUBLICATION testpub6 SET (PUBLISH_VIA_PARTITION_ROOT=0);
+-- ok - partition does not have row filer
+UPDATE rf_tbl_abcd_part_pk SET a = 1;
+ALTER PUBLICATION testpub6 SET (PUBLISH_VIA_PARTITION_ROOT=1);
+-- ok - "a" is a OK col
+UPDATE rf_tbl_abcd_part_pk SET a = 1;
+ALTER PUBLICATION testpub6 SET TABLE rf_tbl_abcd_part_pk WHERE (b > 99);
+ALTER PUBLICATION testpub6 SET (PUBLISH_VIA_PARTITION_ROOT=0);
+-- ok - partition does not have row filer
+UPDATE rf_tbl_abcd_part_pk SET a = 1;
+ALTER PUBLICATION testpub6 SET (PUBLISH_VIA_PARTITION_ROOT=1);
+-- fail - "b" is not in REPLICA IDENTITY INDEX
+UPDATE rf_tbl_abcd_part_pk SET a = 1;
 
+DROP PUBLICATION testpub6;
 DROP TABLE rf_tbl_abcd_pk;
 DROP TABLE rf_tbl_abcd_nopk;
+DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
 
 -- Test cache invalidation FOR ALL TABLES publication
-- 
2.18.4

