Disallow UPDATE/DELETE on table with unpublished generated column as REPLICA IDENTITY

From 867fda5bb3f7e0ef2a8a1c5729fb7eca4d310832 Mon Sep 17 00:00:00 2001
From: Shlok Kyal <shlok.kyal.oss@gmail.com>
Date: Tue, 5 Nov 2024 11:11:25 +0530
Subject: [PATCH v2] Disallow UPDATE/DELETE on table with generated column as
 REPLICA IDENTITY

UPDATE/DELETE on table having unpublished generated column as REPLICA
IDENTITY is allowed. UPDATE/DELETE on such tables should not be allowed
---
 src/backend/commands/publicationcmds.c    | 103 ++++++++++++++++++++++
 src/backend/executor/execReplication.c    |  12 +++
 src/backend/utils/cache/relcache.c        |  16 ++++
 src/include/catalog/pg_publication.h      |   6 ++
 src/include/commands/publicationcmds.h    |   2 +
 src/test/regress/expected/publication.out |  24 +++++
 src/test/regress/sql/publication.sql      |  24 +++++
 7 files changed, 187 insertions(+)

diff --git a/src/backend/commands/publicationcmds.c b/src/backend/commands/publicationcmds.c
index 0129db18c6..fb132b96cd 100644
--- a/src/backend/commands/publicationcmds.c
+++ b/src/backend/commands/publicationcmds.c
@@ -444,6 +444,109 @@ pub_collist_contains_invalid_column(Oid pubid, Relation relation, List *ancestor
 	return result;
 }
 
+/*
+ * Check if REPLICA IDENTITY consists of any unpublished generated column.
+ *
+ * Returns true if any replica identity column is an unpublished generated column.
+ */
+bool
+replident_has_unpublished_gen_col(Oid pubid, Relation relation, List *ancestors,
+								  bool pubviaroot)
+{
+	HeapTuple	tuple;
+	Oid			relid = RelationGetRelid(relation);
+	Oid			publish_as_relid = RelationGetRelid(relation);
+	bool		result = false;
+	bool		isnull;
+
+	/*
+	 * For a partition, if pubviaroot is true, find the topmost ancestor that
+	 * is published via this publication as we need to use its column list for
+	 * the changes.
+	 *
+	 * Note that even though the column list used is for an ancestor, the
+	 * REPLICA IDENTITY used will be for the actual child table.
+	 */
+	if (pubviaroot && relation->rd_rel->relispartition)
+	{
+		publish_as_relid = GetTopMostAncestorInPublication(pubid, ancestors, NULL);
+
+		if (!OidIsValid(publish_as_relid))
+			publish_as_relid = relid;
+	}
+
+	tuple = SearchSysCache2(PUBLICATIONRELMAP,
+							ObjectIdGetDatum(publish_as_relid),
+							ObjectIdGetDatum(pubid));
+
+	if (!HeapTupleIsValid(tuple))
+		return false;
+
+	(void) SysCacheGetAttr(PUBLICATIONRELMAP, tuple,
+						   Anum_pg_publication_rel_prattrs,
+						   &isnull);
+
+	if(isnull)
+	{
+		int			x;
+		Bitmapset  *idattrs = NULL;
+
+		idattrs = RelationGetIndexAttrBitmap(relation,
+											 INDEX_ATTR_BITMAP_IDENTITY_KEY);
+
+		x = -1;
+
+		/*
+		 * Check if any REPLICA IDENTITY column is an generated column.
+		 */
+		while ((x = bms_next_member(idattrs, x)) >= 0)
+		{
+			AttrNumber	attnum = (x + FirstLowInvalidHeapAttributeNumber);
+			char attgenerated = get_attgenerated(relid, attnum);
+
+			/*
+			 * If pubviaroot is true, we are validating the column list of the
+			 * parent table, but the bitmap contains the replica identity
+			 * information of the child table. The parent/child attnums may
+			 * not match, so translate them to the parent - get the attname
+			 * from the child, and look it up in the parent.
+			 */
+			if (pubviaroot)
+			{
+				/* attribute name in the child table */
+				char	   *colname = get_attname(relid, attnum, false);
+
+				/*
+				 * Determine the attnum for the attribute name in parent (we
+				 * are using the column list defined on the parent).
+				 */
+				attnum = get_attnum(publish_as_relid, colname);
+				attgenerated = get_attgenerated(publish_as_relid, attnum);
+			}
+
+			/*
+			 * Check if the column is a generated column.
+			 *
+			 * 'publish_generated_columns = false' and no column list is
+			 * specified for publication. So if the column is a generated
+			 * column, this implies that the REPLICA IDENTITY consists an
+			 * unpublished generated column.
+			 */
+			if (attgenerated == ATTRIBUTE_GENERATED_STORED)
+			{
+				result = true;
+				break;
+			}
+		}
+
+		bms_free(idattrs);
+	}
+
+	ReleaseSysCache(tuple);
+
+	return result;
+}
+
 /* check_functions_in_node callback */
 static bool
 contain_mutable_or_user_functions_checker(Oid func_id, void *context)
diff --git a/src/backend/executor/execReplication.c b/src/backend/executor/execReplication.c
index 54025c9f15..1ba69320ea 100644
--- a/src/backend/executor/execReplication.c
+++ b/src/backend/executor/execReplication.c
@@ -809,6 +809,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot update table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_UPDATE && !pubdesc.replident_has_valid_gen_cols)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot update table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("REPLICA IDENTITY consists of an unpublished generated column.")));
 	else if (cmd == CMD_DELETE && !pubdesc.rf_valid_for_delete)
 		ereport(ERROR,
 				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
@@ -821,6 +827,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot delete from table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_DELETE && !pubdesc.replident_has_valid_gen_cols)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot delete from table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("REPLICA IDENTITY consists of an unpublished generated column.")));
 
 	/* If relation has replica identity we are always good. */
 	if (OidIsValid(RelationGetReplicaIndex(rel)))
diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
index 5bbb654a5d..eea5a9fb52 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -5712,6 +5712,7 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		pubdesc->rf_valid_for_delete = true;
 		pubdesc->cols_valid_for_update = true;
 		pubdesc->cols_valid_for_delete = true;
+		pubdesc->replident_has_valid_gen_cols = true;
 		return;
 	}
 
@@ -5726,6 +5727,7 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 	pubdesc->rf_valid_for_delete = true;
 	pubdesc->cols_valid_for_update = true;
 	pubdesc->cols_valid_for_delete = true;
+	pubdesc->replident_has_valid_gen_cols = true;
 
 	/* Fetch the publication membership info. */
 	puboids = GetRelationPublications(relid);
@@ -5803,6 +5805,20 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 				pubdesc->cols_valid_for_delete = false;
 		}
 
+		/*
+		 * Check if all columns which are part of the REPLICA IDENTITY is
+		 * published.
+		 *
+		 * If the publication is FOR ALL TABLES we can skip the validation.
+		 */
+		if (!pubform->puballtables && !pubform->pubgencols &&
+			(pubform->pubupdate || pubform->pubdelete) &&
+			replident_has_unpublished_gen_col(pubid, relation, ancestors,
+											  pubform->pubviaroot))
+		{
+			pubdesc->replident_has_valid_gen_cols = false;
+		}
+
 		ReleaseSysCache(tup);
 
 		/*
diff --git a/src/include/catalog/pg_publication.h b/src/include/catalog/pg_publication.h
index 9a83a72d6b..5eeeced91d 100644
--- a/src/include/catalog/pg_publication.h
+++ b/src/include/catalog/pg_publication.h
@@ -98,6 +98,12 @@ typedef struct PublicationDesc
 	 */
 	bool		cols_valid_for_update;
 	bool		cols_valid_for_delete;
+
+	/*
+	 * true if all generated columns which are part of replica identity are
+	 * published.
+	 */
+	bool		replident_has_valid_gen_cols;
 } PublicationDesc;
 
 typedef struct Publication
diff --git a/src/include/commands/publicationcmds.h b/src/include/commands/publicationcmds.h
index 5487c571f6..0a862e93a9 100644
--- a/src/include/commands/publicationcmds.h
+++ b/src/include/commands/publicationcmds.h
@@ -35,5 +35,7 @@ extern bool pub_rf_contains_invalid_column(Oid pubid, Relation relation,
 										   List *ancestors, bool pubviaroot);
 extern bool pub_collist_contains_invalid_column(Oid pubid, Relation relation,
 												List *ancestors, bool pubviaroot);
+extern bool replident_has_unpublished_gen_col(Oid pubid, Relation relation,
+												List *ancestors, bool pubviaroot);
 
 #endif							/* PUBLICATIONCMDS_H */
diff --git a/src/test/regress/expected/publication.out b/src/test/regress/expected/publication.out
index a8949ffc2c..d0dfbd969f 100644
--- a/src/test/regress/expected/publication.out
+++ b/src/test/regress/expected/publication.out
@@ -670,6 +670,30 @@ DROP TABLE rf_tbl_abcd_pk;
 DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+-- error - generated column "b" is not published but part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  REPLICA IDENTITY consists of an unpublished generated column.
+DELETE FROM testpub_gencol WHERE a = 100;
+ERROR:  cannot delete from table "testpub_gencol"
+DETAIL:  REPLICA IDENTITY consists of an unpublished generated column.
+DROP PUBLICATION pub_gencol;
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DELETE FROM testpub_gencol WHERE a = 100;
+DROP PUBLICATION pub_gencol;
+DROP INDEX testpub_gencol_idx;
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/regress/sql/publication.sql b/src/test/regress/sql/publication.sql
index 48e68bcca2..d0a31b6c29 100644
--- a/src/test/regress/sql/publication.sql
+++ b/src/test/regress/sql/publication.sql
@@ -396,6 +396,30 @@ DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
 
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+
+-- error - generated column "b" is not published but part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DELETE FROM testpub_gencol WHERE a = 100;
+DROP PUBLICATION pub_gencol;
+
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DELETE FROM testpub_gencol WHERE a = 100;
+DROP PUBLICATION pub_gencol;
+
+DROP INDEX testpub_gencol_idx;
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
+
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
-- 
2.34.1

From 0292ed3510a676dfd6f5f77360008de30af16475 Mon Sep 17 00:00:00 2001
From: Shlok Kyal <shlok.kyal.oss@gmail.com>
Date: Tue, 5 Nov 2024 11:11:25 +0530
Subject: [PATCH v3] Disallow UPDATE/DELETE on table with generated column as
 REPLICA IDENTITY

UPDATE/DELETE on table having unpublished generated column as REPLICA
IDENTITY is allowed. UPDATE/DELETE on such tables should not be allowed
---
 src/backend/commands/publicationcmds.c    | 117 ++++++++++++++++++++++
 src/backend/executor/execReplication.c    |  12 +++
 src/backend/utils/cache/relcache.c        |  14 +++
 src/include/catalog/pg_publication.h      |   6 ++
 src/include/commands/publicationcmds.h    |   2 +
 src/test/regress/expected/publication.out |  32 ++++++
 src/test/regress/sql/publication.sql      |  29 ++++++
 src/test/subscription/t/100_bugs.pl       |  16 +--
 8 files changed, 215 insertions(+), 13 deletions(-)

diff --git a/src/backend/commands/publicationcmds.c b/src/backend/commands/publicationcmds.c
index 0129db18c6..21ed6344ee 100644
--- a/src/backend/commands/publicationcmds.c
+++ b/src/backend/commands/publicationcmds.c
@@ -444,6 +444,123 @@ pub_collist_contains_invalid_column(Oid pubid, Relation relation, List *ancestor
 	return result;
 }
 
+/*
+ * Check if REPLICA IDENTITY consists of any unpublished generated column.
+ *
+ * Returns true if any replica identity column is an unpublished generated column.
+ */
+bool
+replident_has_unpublished_gen_col(Oid pubid, Relation relation, List *ancestors,
+								  bool pubviaroot, bool puballtables)
+{
+	HeapTuple	tuple;
+	Oid			relid = RelationGetRelid(relation);
+	Oid			publish_as_relid = RelationGetRelid(relation);
+	bool		result = false;
+	bool		isnull;
+
+	/*
+	 * For a partition, if pubviaroot is true, find the topmost ancestor that
+	 * is published via this publication as we need to use its column list for
+	 * the changes.
+	 *
+	 * Note that even though the column list used is for an ancestor, the
+	 * REPLICA IDENTITY used will be for the actual child table.
+	 */
+	if (pubviaroot && relation->rd_rel->relispartition)
+	{
+		publish_as_relid = GetTopMostAncestorInPublication(pubid, ancestors, NULL);
+
+		if (!OidIsValid(publish_as_relid))
+			publish_as_relid = relid;
+	}
+
+	if (!puballtables)
+	{
+		tuple = SearchSysCache2(PUBLICATIONRELMAP,
+								ObjectIdGetDatum(publish_as_relid),
+								ObjectIdGetDatum(pubid));
+
+		if (!HeapTupleIsValid(tuple))
+			return false;
+
+		(void) SysCacheGetAttr(PUBLICATIONRELMAP, tuple,
+							   Anum_pg_publication_rel_prattrs,
+							   &isnull);
+
+		ReleaseSysCache(tuple);
+	}
+
+	if(puballtables || isnull)
+	{
+		int			x;
+		Bitmapset  *idattrs = NULL;
+
+		/*
+		 * REPLICA IDENTITY can be FULL only if there is no column list for
+		 * publication. If REPLICA IDENTITY is set as FULL and relation has a
+		 * generated column we should error out.
+		 */
+		if(relation->rd_rel->relreplident == REPLICA_IDENTITY_FULL &&
+		   relation->rd_att->constr && relation->rd_att->constr->has_generated_stored)
+			result = true;
+
+		idattrs = RelationGetIndexAttrBitmap(relation,
+											 INDEX_ATTR_BITMAP_IDENTITY_KEY);
+
+		x = -1;
+
+		/*
+		 * Check if any REPLICA IDENTITY column is an generated column.
+		 */
+		while ((x = bms_next_member(idattrs, x)) >= 0)
+		{
+			AttrNumber	attnum = (x + FirstLowInvalidHeapAttributeNumber);
+			char attgenerated;
+
+			/*
+			 * If pubviaroot is true, we are validating the column list of the
+			 * parent table, but the bitmap contains the replica identity
+			 * information of the child table. The parent/child attnums may
+			 * not match, so translate them to the parent - get the attname
+			 * from the child, and look it up in the parent.
+			 */
+			if (pubviaroot)
+			{
+				/* attribute name in the child table */
+				char	   *colname = get_attname(relid, attnum, false);
+
+				/*
+				 * Determine the attnum for the attribute name in parent (we
+				 * are using the column list defined on the parent).
+				 */
+				attnum = get_attnum(publish_as_relid, colname);
+				attgenerated = get_attgenerated(publish_as_relid, attnum);
+			}
+			else
+				attgenerated = get_attgenerated(relid, attnum);
+
+			/*
+			 * Check if the column is a generated column.
+			 *
+			 * 'publish_generated_columns = false' and no column list is
+			 * specified for publication. So if the column is a generated
+			 * column, this implies that the REPLICA IDENTITY consists an
+			 * unpublished generated column.
+			 */
+			if (attgenerated == ATTRIBUTE_GENERATED_STORED)
+			{
+				result = true;
+				break;
+			}
+		}
+
+		bms_free(idattrs);
+	}
+
+	return result;
+}
+
 /* check_functions_in_node callback */
 static bool
 contain_mutable_or_user_functions_checker(Oid func_id, void *context)
diff --git a/src/backend/executor/execReplication.c b/src/backend/executor/execReplication.c
index 54025c9f15..d3517db8b3 100644
--- a/src/backend/executor/execReplication.c
+++ b/src/backend/executor/execReplication.c
@@ -809,6 +809,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot update table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_UPDATE && !pubdesc.replident_has_valid_gen_cols)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot update table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("replica identity consists of an unpublished generated column.")));
 	else if (cmd == CMD_DELETE && !pubdesc.rf_valid_for_delete)
 		ereport(ERROR,
 				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
@@ -821,6 +827,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot delete from table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_DELETE && !pubdesc.replident_has_valid_gen_cols)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot delete from table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("replica identity consists of an unpublished generated column.")));
 
 	/* If relation has replica identity we are always good. */
 	if (OidIsValid(RelationGetReplicaIndex(rel)))
diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
index 342467fd18..87e1c12bac 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -5736,6 +5736,7 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		pubdesc->rf_valid_for_delete = true;
 		pubdesc->cols_valid_for_update = true;
 		pubdesc->cols_valid_for_delete = true;
+		pubdesc->replident_has_valid_gen_cols = true;
 		return;
 	}
 
@@ -5750,6 +5751,7 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 	pubdesc->rf_valid_for_delete = true;
 	pubdesc->cols_valid_for_update = true;
 	pubdesc->cols_valid_for_delete = true;
+	pubdesc->replident_has_valid_gen_cols = true;
 
 	/* Fetch the publication membership info. */
 	puboids = GetRelationPublications(relid);
@@ -5827,6 +5829,18 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 				pubdesc->cols_valid_for_delete = false;
 		}
 
+		/*
+		 * Check if all columns which are part of the REPLICA IDENTITY is
+		 * published.
+		 */
+		if (!pubform->pubgencols &&
+			(pubform->pubupdate || pubform->pubdelete) &&
+			replident_has_unpublished_gen_col(pubid, relation, ancestors,
+											  pubform->pubviaroot, pubform->puballtables))
+		{
+			pubdesc->replident_has_valid_gen_cols = false;
+		}
+
 		ReleaseSysCache(tup);
 
 		/*
diff --git a/src/include/catalog/pg_publication.h b/src/include/catalog/pg_publication.h
index 9a83a72d6b..5eeeced91d 100644
--- a/src/include/catalog/pg_publication.h
+++ b/src/include/catalog/pg_publication.h
@@ -98,6 +98,12 @@ typedef struct PublicationDesc
 	 */
 	bool		cols_valid_for_update;
 	bool		cols_valid_for_delete;
+
+	/*
+	 * true if all generated columns which are part of replica identity are
+	 * published.
+	 */
+	bool		replident_has_valid_gen_cols;
 } PublicationDesc;
 
 typedef struct Publication
diff --git a/src/include/commands/publicationcmds.h b/src/include/commands/publicationcmds.h
index 5487c571f6..23153afd8a 100644
--- a/src/include/commands/publicationcmds.h
+++ b/src/include/commands/publicationcmds.h
@@ -35,5 +35,7 @@ extern bool pub_rf_contains_invalid_column(Oid pubid, Relation relation,
 										   List *ancestors, bool pubviaroot);
 extern bool pub_collist_contains_invalid_column(Oid pubid, Relation relation,
 												List *ancestors, bool pubviaroot);
+extern bool replident_has_unpublished_gen_col(Oid pubid, Relation relation,
+												List *ancestors, bool pubviaroot, bool puballtables);
 
 #endif							/* PUBLICATIONCMDS_H */
diff --git a/src/test/regress/expected/publication.out b/src/test/regress/expected/publication.out
index 5de2d64d01..fe17dda5c9 100644
--- a/src/test/regress/expected/publication.out
+++ b/src/test/regress/expected/publication.out
@@ -672,6 +672,38 @@ DROP TABLE rf_tbl_abcd_pk;
 DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  replica identity consists of an unpublished generated column.
+DELETE FROM testpub_gencol WHERE a = 100;
+ERROR:  cannot delete from table "testpub_gencol"
+DETAIL:  replica identity consists of an unpublished generated column.
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  replica identity consists of an unpublished generated column.
+DELETE FROM testpub_gencol WHERE a = 100;
+ERROR:  cannot delete from table "testpub_gencol"
+DETAIL:  replica identity consists of an unpublished generated column.
+DROP PUBLICATION pub_gencol;
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DELETE FROM testpub_gencol WHERE a = 100;
+DROP PUBLICATION pub_gencol;
+DROP INDEX testpub_gencol_idx;
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/regress/sql/publication.sql b/src/test/regress/sql/publication.sql
index 48e68bcca2..edf71e5325 100644
--- a/src/test/regress/sql/publication.sql
+++ b/src/test/regress/sql/publication.sql
@@ -396,6 +396,35 @@ DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
 
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DELETE FROM testpub_gencol WHERE a = 100;
+
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DELETE FROM testpub_gencol WHERE a = 100;
+DROP PUBLICATION pub_gencol;
+
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DELETE FROM testpub_gencol WHERE a = 100;
+DROP PUBLICATION pub_gencol;
+
+DROP INDEX testpub_gencol_idx;
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
+
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/subscription/t/100_bugs.pl b/src/test/subscription/t/100_bugs.pl
index cb36ca7b16..794b928f50 100644
--- a/src/test/subscription/t/100_bugs.pl
+++ b/src/test/subscription/t/100_bugs.pl
@@ -377,8 +377,8 @@ $node_publisher->safe_psql('postgres', "DROP PUBLICATION tap_pub_sch");
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');
 
-# The bug was that when the REPLICA IDENTITY FULL is used with dropped or
-# generated columns, we fail to apply updates and deletes
+# The bug was that when the REPLICA IDENTITY FULL is used with dropped
+# we fail to apply updates and deletes
 $node_publisher->rotate_logfile();
 $node_publisher->start();
 
@@ -389,18 +389,14 @@ $node_publisher->safe_psql(
 	'postgres', qq(
 	CREATE TABLE dropped_cols (a int, b_drop int, c int);
 	ALTER TABLE dropped_cols REPLICA IDENTITY FULL;
-	CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
-	ALTER TABLE generated_cols REPLICA IDENTITY FULL;
-	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols, generated_cols;
+	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols;
 	-- some initial data
 	INSERT INTO dropped_cols VALUES (1, 1, 1);
-	INSERT INTO generated_cols (a, c) VALUES (1, 1);
 ));
 
 $node_subscriber->safe_psql(
 	'postgres', qq(
 	 CREATE TABLE dropped_cols (a int, b_drop int, c int);
-	 CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
 ));
 
 $publisher_connstr = $node_publisher->connstr . ' dbname=postgres';
@@ -421,7 +417,6 @@ $node_subscriber->safe_psql(
 $node_publisher->safe_psql(
 	'postgres', qq(
 		UPDATE dropped_cols SET a = 100;
-		UPDATE generated_cols SET a = 100;
 ));
 $node_publisher->wait_for_catchup('sub_dropped_cols');
 
@@ -430,11 +425,6 @@ is( $node_subscriber->safe_psql(
 	qq(1),
 	'replication with RI FULL and dropped columns');
 
-is( $node_subscriber->safe_psql(
-		'postgres', "SELECT count(*) FROM generated_cols WHERE a = 100"),
-	qq(1),
-	'replication with RI FULL and generated columns');
-
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');
 
-- 
2.34.1

From 1aaa1cf9379ae1961a05f417c7bf6553bf789cf6 Mon Sep 17 00:00:00 2001
From: Shlok Kyal <shlok.kyal.oss@gmail.com>
Date: Tue, 5 Nov 2024 11:11:25 +0530
Subject: [PATCH v4] Disallow UPDATE/DELETE on table with generated column as
 REPLICA IDENTITY

UPDATE/DELETE on table having unpublished generated column as REPLICA
IDENTITY is allowed. UPDATE/DELETE on such tables should not be allowed
---
 src/backend/commands/publicationcmds.c    | 102 ++++++++++++++++++++++
 src/backend/executor/execReplication.c    |  12 +++
 src/backend/utils/cache/relcache.c        |  12 +++
 src/include/catalog/pg_publication.h      |   6 ++
 src/include/commands/publicationcmds.h    |   2 +
 src/test/regress/expected/publication.out |  25 ++++++
 src/test/regress/sql/publication.sql      |  26 ++++++
 src/test/subscription/t/100_bugs.pl       |  16 +---
 8 files changed, 188 insertions(+), 13 deletions(-)

diff --git a/src/backend/commands/publicationcmds.c b/src/backend/commands/publicationcmds.c
index 0129db18c6..a095c61edc 100644
--- a/src/backend/commands/publicationcmds.c
+++ b/src/backend/commands/publicationcmds.c
@@ -444,6 +444,108 @@ pub_collist_contains_invalid_column(Oid pubid, Relation relation, List *ancestor
 	return result;
 }
 
+/*
+ * Check if REPLICA IDENTITY consists of any unpublished generated column.
+ *
+ * Returns true if any replica identity column is an unpublished generated column.
+ */
+bool
+replident_has_unpublished_gen_col(Oid pubid, Relation relation, List *ancestors,
+								  bool pubviaroot)
+{
+	Oid			relid = RelationGetRelid(relation);
+	Oid			publish_as_relid = RelationGetRelid(relation);
+	bool		result = false;
+	bool		found;
+	Publication *pub;
+
+	/*
+	 * For a partition, if pubviaroot is true, find the topmost ancestor that
+	 * is published via this publication as we need to use its column list for
+	 * the changes.
+	 *
+	 * Note that even though the column list used is for an ancestor, the
+	 * REPLICA IDENTITY used will be for the actual child table.
+	 */
+	if (pubviaroot && relation->rd_rel->relispartition)
+	{
+		publish_as_relid = GetTopMostAncestorInPublication(pubid, ancestors, NULL);
+
+		if (!OidIsValid(publish_as_relid))
+			publish_as_relid = relid;
+	}
+
+	pub = GetPublication(pubid);
+	found = check_and_fetch_column_list(pub, publish_as_relid, NULL, NULL);
+
+	if (!found)
+	{
+		int			x;
+		Bitmapset  *idattrs = NULL;
+
+		/*
+		 * REPLICA IDENTITY can be FULL only if there is no column list for
+		 * publication. If REPLICA IDENTITY is set as FULL and relation has a
+		 * generated column we should error out.
+		 */
+		if (relation->rd_rel->relreplident == REPLICA_IDENTITY_FULL &&
+			relation->rd_att->constr && relation->rd_att->constr->has_generated_stored)
+			result = true;
+
+		idattrs = RelationGetIndexAttrBitmap(relation,
+											 INDEX_ATTR_BITMAP_IDENTITY_KEY);
+
+		x = -1;
+
+		/* Check if any REPLICA IDENTITY column is an generated column. */
+		while ((x = bms_next_member(idattrs, x)) >= 0)
+		{
+			AttrNumber	attnum = (x + FirstLowInvalidHeapAttributeNumber);
+			char		attgenerated;
+
+			/*
+			 * If pubviaroot is true, we are validating the column list of the
+			 * parent table, but the bitmap contains the replica identity
+			 * information of the child table. The parent/child attnums may
+			 * not match, so translate them to the parent - get the attname
+			 * from the child, and look it up in the parent.
+			 */
+			if (pubviaroot && relation->rd_rel->relispartition)
+			{
+				/* attribute name in the child table */
+				char	   *colname = get_attname(relid, attnum, false);
+
+				/*
+				 * Determine the attnum for the attribute name in parent (we
+				 * are using the column list defined on the parent).
+				 */
+				attnum = get_attnum(publish_as_relid, colname);
+				attgenerated = get_attgenerated(publish_as_relid, attnum);
+			}
+			else
+				attgenerated = get_attgenerated(relid, attnum);
+
+			/*
+			 * Check if the column is a generated column.
+			 *
+			 * 'publish_generated_columns = false' and no column list is
+			 * specified for publication. So if the column is a generated
+			 * column, this implies that the REPLICA IDENTITY consists an
+			 * unpublished generated column.
+			 */
+			if (attgenerated == ATTRIBUTE_GENERATED_STORED)
+			{
+				result = true;
+				break;
+			}
+		}
+
+		bms_free(idattrs);
+	}
+
+	return result;
+}
+
 /* check_functions_in_node callback */
 static bool
 contain_mutable_or_user_functions_checker(Oid func_id, void *context)
diff --git a/src/backend/executor/execReplication.c b/src/backend/executor/execReplication.c
index 54025c9f15..d3517db8b3 100644
--- a/src/backend/executor/execReplication.c
+++ b/src/backend/executor/execReplication.c
@@ -809,6 +809,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot update table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_UPDATE && !pubdesc.replident_has_valid_gen_cols)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot update table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("replica identity consists of an unpublished generated column.")));
 	else if (cmd == CMD_DELETE && !pubdesc.rf_valid_for_delete)
 		ereport(ERROR,
 				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
@@ -821,6 +827,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot delete from table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_DELETE && !pubdesc.replident_has_valid_gen_cols)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot delete from table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("replica identity consists of an unpublished generated column.")));
 
 	/* If relation has replica identity we are always good. */
 	if (OidIsValid(RelationGetReplicaIndex(rel)))
diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
index 342467fd18..02d7fba160 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -5736,6 +5736,7 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		pubdesc->rf_valid_for_delete = true;
 		pubdesc->cols_valid_for_update = true;
 		pubdesc->cols_valid_for_delete = true;
+		pubdesc->replident_has_valid_gen_cols = true;
 		return;
 	}
 
@@ -5750,6 +5751,7 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 	pubdesc->rf_valid_for_delete = true;
 	pubdesc->cols_valid_for_update = true;
 	pubdesc->cols_valid_for_delete = true;
+	pubdesc->replident_has_valid_gen_cols = true;
 
 	/* Fetch the publication membership info. */
 	puboids = GetRelationPublications(relid);
@@ -5827,6 +5829,16 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 				pubdesc->cols_valid_for_delete = false;
 		}
 
+		/*
+		 * Check if all columns which are part of the REPLICA IDENTITY is
+		 * published.
+		 */
+		if (!pubform->pubgencols &&
+			(pubform->pubupdate || pubform->pubdelete) &&
+			replident_has_unpublished_gen_col(pubid, relation, ancestors,
+											  pubform->pubviaroot))
+			pubdesc->replident_has_valid_gen_cols = false;
+
 		ReleaseSysCache(tup);
 
 		/*
diff --git a/src/include/catalog/pg_publication.h b/src/include/catalog/pg_publication.h
index 9a83a72d6b..5eeeced91d 100644
--- a/src/include/catalog/pg_publication.h
+++ b/src/include/catalog/pg_publication.h
@@ -98,6 +98,12 @@ typedef struct PublicationDesc
 	 */
 	bool		cols_valid_for_update;
 	bool		cols_valid_for_delete;
+
+	/*
+	 * true if all generated columns which are part of replica identity are
+	 * published.
+	 */
+	bool		replident_has_valid_gen_cols;
 } PublicationDesc;
 
 typedef struct Publication
diff --git a/src/include/commands/publicationcmds.h b/src/include/commands/publicationcmds.h
index 5487c571f6..b18e576b77 100644
--- a/src/include/commands/publicationcmds.h
+++ b/src/include/commands/publicationcmds.h
@@ -35,5 +35,7 @@ extern bool pub_rf_contains_invalid_column(Oid pubid, Relation relation,
 										   List *ancestors, bool pubviaroot);
 extern bool pub_collist_contains_invalid_column(Oid pubid, Relation relation,
 												List *ancestors, bool pubviaroot);
+extern bool replident_has_unpublished_gen_col(Oid pubid, Relation relation,
+											  List *ancestors, bool pubviaroot);
 
 #endif							/* PUBLICATIONCMDS_H */
diff --git a/src/test/regress/expected/publication.out b/src/test/regress/expected/publication.out
index 5de2d64d01..62198ead23 100644
--- a/src/test/regress/expected/publication.out
+++ b/src/test/regress/expected/publication.out
@@ -672,6 +672,31 @@ DROP TABLE rf_tbl_abcd_pk;
 DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  replica identity consists of an unpublished generated column.
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  replica identity consists of an unpublished generated column.
+DROP PUBLICATION pub_gencol;
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+DROP INDEX testpub_gencol_idx;
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/regress/sql/publication.sql b/src/test/regress/sql/publication.sql
index 48e68bcca2..3b1c661440 100644
--- a/src/test/regress/sql/publication.sql
+++ b/src/test/regress/sql/publication.sql
@@ -396,6 +396,32 @@ DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
 
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+
+DROP INDEX testpub_gencol_idx;
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
+
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/subscription/t/100_bugs.pl b/src/test/subscription/t/100_bugs.pl
index cb36ca7b16..794b928f50 100644
--- a/src/test/subscription/t/100_bugs.pl
+++ b/src/test/subscription/t/100_bugs.pl
@@ -377,8 +377,8 @@ $node_publisher->safe_psql('postgres', "DROP PUBLICATION tap_pub_sch");
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');
 
-# The bug was that when the REPLICA IDENTITY FULL is used with dropped or
-# generated columns, we fail to apply updates and deletes
+# The bug was that when the REPLICA IDENTITY FULL is used with dropped
+# we fail to apply updates and deletes
 $node_publisher->rotate_logfile();
 $node_publisher->start();
 
@@ -389,18 +389,14 @@ $node_publisher->safe_psql(
 	'postgres', qq(
 	CREATE TABLE dropped_cols (a int, b_drop int, c int);
 	ALTER TABLE dropped_cols REPLICA IDENTITY FULL;
-	CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
-	ALTER TABLE generated_cols REPLICA IDENTITY FULL;
-	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols, generated_cols;
+	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols;
 	-- some initial data
 	INSERT INTO dropped_cols VALUES (1, 1, 1);
-	INSERT INTO generated_cols (a, c) VALUES (1, 1);
 ));
 
 $node_subscriber->safe_psql(
 	'postgres', qq(
 	 CREATE TABLE dropped_cols (a int, b_drop int, c int);
-	 CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
 ));
 
 $publisher_connstr = $node_publisher->connstr . ' dbname=postgres';
@@ -421,7 +417,6 @@ $node_subscriber->safe_psql(
 $node_publisher->safe_psql(
 	'postgres', qq(
 		UPDATE dropped_cols SET a = 100;
-		UPDATE generated_cols SET a = 100;
 ));
 $node_publisher->wait_for_catchup('sub_dropped_cols');
 
@@ -430,11 +425,6 @@ is( $node_subscriber->safe_psql(
 	qq(1),
 	'replication with RI FULL and dropped columns');
 
-is( $node_subscriber->safe_psql(
-		'postgres', "SELECT count(*) FROM generated_cols WHERE a = 100"),
-	qq(1),
-	'replication with RI FULL and generated columns');
-
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');
 
-- 
2.34.1

From 924891000ac8ff1f4a77dc9c2a0c46ab141e8f5c Mon Sep 17 00:00:00 2001
From: Shlok Kyal <shlok.kyal.oss@gmail.com>
Date: Tue, 5 Nov 2024 11:11:25 +0530
Subject: [PATCH v5] Disallow UPDATE/DELETE on table with generated column as
 REPLICA IDENTITY

UPDATE/DELETE on table having unpublished generated column as REPLICA
IDENTITY is allowed. UPDATE/DELETE on such tables should not be allowed
---
 src/backend/commands/publicationcmds.c    | 88 +++++++++++++++++++++++
 src/backend/executor/execReplication.c    | 12 ++++
 src/backend/utils/cache/relcache.c        | 12 ++++
 src/include/catalog/pg_publication.h      |  6 ++
 src/include/commands/publicationcmds.h    |  2 +
 src/test/regress/expected/publication.out | 25 +++++++
 src/test/regress/sql/publication.sql      | 26 +++++++
 src/test/subscription/t/100_bugs.pl       | 16 +----
 8 files changed, 174 insertions(+), 13 deletions(-)

diff --git a/src/backend/commands/publicationcmds.c b/src/backend/commands/publicationcmds.c
index 0129db18c6..c354094c7d 100644
--- a/src/backend/commands/publicationcmds.c
+++ b/src/backend/commands/publicationcmds.c
@@ -444,6 +444,94 @@ pub_collist_contains_invalid_column(Oid pubid, Relation relation, List *ancestor
 	return result;
 }
 
+/*
+ * Check if REPLICA IDENTITY consists of any unpublished generated column.
+ *
+ * Returns true if any replica identity column is an unpublished generated
+ * column.
+ */
+bool
+replident_has_unpublished_gen_col(Oid pubid, Relation relation, List *ancestors,
+								  bool pubviaroot)
+{
+	Oid			relid = RelationGetRelid(relation);
+	Oid			publish_as_relid;
+	bool		result = false;
+	bool		found;
+	Publication *pub;
+
+	/*
+	 * For a partition, if pubviaroot is true, find the topmost ancestor that
+	 * is published via this publication as we need to use its column list for
+	 * the changes.
+	 *
+	 * Note that even though the column list used is for an ancestor, the
+	 * REPLICA IDENTITY used will be for the actual child table.
+	 */
+	if (pubviaroot && relation->rd_rel->relispartition)
+	{
+		publish_as_relid = GetTopMostAncestorInPublication(pubid, ancestors, NULL);
+
+		if (!OidIsValid(publish_as_relid))
+			publish_as_relid = relid;
+	}
+	else
+		publish_as_relid = relid;
+
+	pub = GetPublication(pubid);
+	found = check_and_fetch_column_list(pub, publish_as_relid, NULL, NULL);
+
+	if (!found)
+	{
+		int			x;
+		Bitmapset  *idattrs = NULL;
+		Bitmapset  *columns = NULL;
+
+		/*
+		 * REPLICA IDENTITY can be FULL only if there is no column list for
+		 * publication. If REPLICA IDENTITY is set as FULL and relation has a
+		 * generated column we should error out.
+		 */
+		if (relation->rd_rel->relreplident == REPLICA_IDENTITY_FULL &&
+			relation->rd_att->constr &&
+			relation->rd_att->constr->has_generated_stored)
+			result = true;
+
+		idattrs = RelationGetIndexAttrBitmap(relation,
+											 INDEX_ATTR_BITMAP_IDENTITY_KEY);
+
+		/*
+		 * Bitmap of published columns when publish_generated_columns is
+		 * 'false' and no column list is specified.
+		 */
+		columns = pub_form_cols_map(relation, false);
+
+		/*
+		 * Attnums in the bitmap returned by RelationGetIndexAttrBitmap are
+		 * offset (to handle system columns the usual way), while column list
+		 * does not use offset, so we can't do bms_is_subset(). Instead, we
+		 * have to loop over the idattrs and check all of them are in the
+		 * list.
+		 */
+		x = -1;
+		while ((x = bms_next_member(idattrs, x)) >= 0)
+		{
+			AttrNumber	attnum = (x + FirstLowInvalidHeapAttributeNumber);
+
+			/* replica identity column, not covered by the column list */
+			if (!bms_is_member(attnum, columns))
+			{
+				result = true;
+				break;
+			}
+		}
+
+		bms_free(idattrs);
+	}
+
+	return result;
+}
+
 /* check_functions_in_node callback */
 static bool
 contain_mutable_or_user_functions_checker(Oid func_id, void *context)
diff --git a/src/backend/executor/execReplication.c b/src/backend/executor/execReplication.c
index 54025c9f15..d3517db8b3 100644
--- a/src/backend/executor/execReplication.c
+++ b/src/backend/executor/execReplication.c
@@ -809,6 +809,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot update table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_UPDATE && !pubdesc.replident_has_valid_gen_cols)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot update table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("replica identity consists of an unpublished generated column.")));
 	else if (cmd == CMD_DELETE && !pubdesc.rf_valid_for_delete)
 		ereport(ERROR,
 				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
@@ -821,6 +827,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot delete from table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_DELETE && !pubdesc.replident_has_valid_gen_cols)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot delete from table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("replica identity consists of an unpublished generated column.")));
 
 	/* If relation has replica identity we are always good. */
 	if (OidIsValid(RelationGetReplicaIndex(rel)))
diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
index 342467fd18..02d7fba160 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -5736,6 +5736,7 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		pubdesc->rf_valid_for_delete = true;
 		pubdesc->cols_valid_for_update = true;
 		pubdesc->cols_valid_for_delete = true;
+		pubdesc->replident_has_valid_gen_cols = true;
 		return;
 	}
 
@@ -5750,6 +5751,7 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 	pubdesc->rf_valid_for_delete = true;
 	pubdesc->cols_valid_for_update = true;
 	pubdesc->cols_valid_for_delete = true;
+	pubdesc->replident_has_valid_gen_cols = true;
 
 	/* Fetch the publication membership info. */
 	puboids = GetRelationPublications(relid);
@@ -5827,6 +5829,16 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 				pubdesc->cols_valid_for_delete = false;
 		}
 
+		/*
+		 * Check if all columns which are part of the REPLICA IDENTITY is
+		 * published.
+		 */
+		if (!pubform->pubgencols &&
+			(pubform->pubupdate || pubform->pubdelete) &&
+			replident_has_unpublished_gen_col(pubid, relation, ancestors,
+											  pubform->pubviaroot))
+			pubdesc->replident_has_valid_gen_cols = false;
+
 		ReleaseSysCache(tup);
 
 		/*
diff --git a/src/include/catalog/pg_publication.h b/src/include/catalog/pg_publication.h
index 9a83a72d6b..5eeeced91d 100644
--- a/src/include/catalog/pg_publication.h
+++ b/src/include/catalog/pg_publication.h
@@ -98,6 +98,12 @@ typedef struct PublicationDesc
 	 */
 	bool		cols_valid_for_update;
 	bool		cols_valid_for_delete;
+
+	/*
+	 * true if all generated columns which are part of replica identity are
+	 * published.
+	 */
+	bool		replident_has_valid_gen_cols;
 } PublicationDesc;
 
 typedef struct Publication
diff --git a/src/include/commands/publicationcmds.h b/src/include/commands/publicationcmds.h
index 5487c571f6..b18e576b77 100644
--- a/src/include/commands/publicationcmds.h
+++ b/src/include/commands/publicationcmds.h
@@ -35,5 +35,7 @@ extern bool pub_rf_contains_invalid_column(Oid pubid, Relation relation,
 										   List *ancestors, bool pubviaroot);
 extern bool pub_collist_contains_invalid_column(Oid pubid, Relation relation,
 												List *ancestors, bool pubviaroot);
+extern bool replident_has_unpublished_gen_col(Oid pubid, Relation relation,
+											  List *ancestors, bool pubviaroot);
 
 #endif							/* PUBLICATIONCMDS_H */
diff --git a/src/test/regress/expected/publication.out b/src/test/regress/expected/publication.out
index 5de2d64d01..62198ead23 100644
--- a/src/test/regress/expected/publication.out
+++ b/src/test/regress/expected/publication.out
@@ -672,6 +672,31 @@ DROP TABLE rf_tbl_abcd_pk;
 DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  replica identity consists of an unpublished generated column.
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  replica identity consists of an unpublished generated column.
+DROP PUBLICATION pub_gencol;
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+DROP INDEX testpub_gencol_idx;
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/regress/sql/publication.sql b/src/test/regress/sql/publication.sql
index 48e68bcca2..3b1c661440 100644
--- a/src/test/regress/sql/publication.sql
+++ b/src/test/regress/sql/publication.sql
@@ -396,6 +396,32 @@ DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
 
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+
+DROP INDEX testpub_gencol_idx;
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
+
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/subscription/t/100_bugs.pl b/src/test/subscription/t/100_bugs.pl
index cb36ca7b16..794b928f50 100644
--- a/src/test/subscription/t/100_bugs.pl
+++ b/src/test/subscription/t/100_bugs.pl
@@ -377,8 +377,8 @@ $node_publisher->safe_psql('postgres', "DROP PUBLICATION tap_pub_sch");
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');
 
-# The bug was that when the REPLICA IDENTITY FULL is used with dropped or
-# generated columns, we fail to apply updates and deletes
+# The bug was that when the REPLICA IDENTITY FULL is used with dropped
+# we fail to apply updates and deletes
 $node_publisher->rotate_logfile();
 $node_publisher->start();
 
@@ -389,18 +389,14 @@ $node_publisher->safe_psql(
 	'postgres', qq(
 	CREATE TABLE dropped_cols (a int, b_drop int, c int);
 	ALTER TABLE dropped_cols REPLICA IDENTITY FULL;
-	CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
-	ALTER TABLE generated_cols REPLICA IDENTITY FULL;
-	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols, generated_cols;
+	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols;
 	-- some initial data
 	INSERT INTO dropped_cols VALUES (1, 1, 1);
-	INSERT INTO generated_cols (a, c) VALUES (1, 1);
 ));
 
 $node_subscriber->safe_psql(
 	'postgres', qq(
 	 CREATE TABLE dropped_cols (a int, b_drop int, c int);
-	 CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
 ));
 
 $publisher_connstr = $node_publisher->connstr . ' dbname=postgres';
@@ -421,7 +417,6 @@ $node_subscriber->safe_psql(
 $node_publisher->safe_psql(
 	'postgres', qq(
 		UPDATE dropped_cols SET a = 100;
-		UPDATE generated_cols SET a = 100;
 ));
 $node_publisher->wait_for_catchup('sub_dropped_cols');
 
@@ -430,11 +425,6 @@ is( $node_subscriber->safe_psql(
 	qq(1),
 	'replication with RI FULL and dropped columns');
 
-is( $node_subscriber->safe_psql(
-		'postgres', "SELECT count(*) FROM generated_cols WHERE a = 100"),
-	qq(1),
-	'replication with RI FULL and generated columns');
-
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');
 
-- 
2.34.1

From 6a59b708494ad8e669a22714f87f41fe08f5cb58 Mon Sep 17 00:00:00 2001
From: Shlok Kyal <shlok.kyal.oss@gmail.com>
Date: Tue, 5 Nov 2024 11:11:25 +0530
Subject: [PATCH v6] Disallow UPDATE/DELETE on table with generated column as
 REPLICA IDENTITY

UPDATE/DELETE on table having unpublished generated column as REPLICA
IDENTITY is allowed. UPDATE/DELETE on such tables should not be allowed
---
 src/backend/commands/publicationcmds.c    | 89 +++++++++++++++++++++++
 src/backend/executor/execReplication.c    | 12 +++
 src/backend/utils/cache/relcache.c        | 12 +++
 src/include/catalog/pg_publication.h      |  6 ++
 src/include/commands/publicationcmds.h    |  2 +
 src/test/regress/expected/publication.out | 25 +++++++
 src/test/regress/sql/publication.sql      | 26 +++++++
 src/test/subscription/t/100_bugs.pl       | 16 +---
 8 files changed, 175 insertions(+), 13 deletions(-)

diff --git a/src/backend/commands/publicationcmds.c b/src/backend/commands/publicationcmds.c
index 0129db18c6..05debf6d79 100644
--- a/src/backend/commands/publicationcmds.c
+++ b/src/backend/commands/publicationcmds.c
@@ -444,6 +444,95 @@ pub_collist_contains_invalid_column(Oid pubid, Relation relation, List *ancestor
 	return result;
 }
 
+/*
+ * Check if REPLICA IDENTITY consists of any unpublished generated column.
+ *
+ * Returns true if any replica identity column is an unpublished generated
+ * column.
+ */
+bool
+replident_has_unpublished_gen_col(Oid pubid, Relation relation, List *ancestors,
+								  bool pubviaroot)
+{
+	Oid			relid = RelationGetRelid(relation);
+	Oid			publish_as_relid;
+	bool		result = false;
+	bool		found;
+	Publication *pub;
+
+	/*
+	 * For a partition, if pubviaroot is true, find the topmost ancestor that
+	 * is published via this publication as we need to use its column list for
+	 * the changes.
+	 *
+	 * Note that even though the column list used is for an ancestor, the
+	 * REPLICA IDENTITY used will be for the actual child table.
+	 */
+	if (pubviaroot && relation->rd_rel->relispartition)
+	{
+		publish_as_relid = GetTopMostAncestorInPublication(pubid, ancestors, NULL);
+
+		if (!OidIsValid(publish_as_relid))
+			publish_as_relid = relid;
+	}
+	else
+		publish_as_relid = relid;
+
+	pub = GetPublication(pubid);
+	found = check_and_fetch_column_list(pub, publish_as_relid, NULL, NULL);
+
+	if (!found)
+	{
+		int			x;
+		Bitmapset  *idattrs = NULL;
+		Bitmapset  *columns = NULL;
+
+		/*
+		 * REPLICA IDENTITY can be FULL only if there is no column list for
+		 * publication. If REPLICA IDENTITY is set as FULL and relation has a
+		 * generated column we should error out.
+		 */
+		if (relation->rd_rel->relreplident == REPLICA_IDENTITY_FULL &&
+			relation->rd_att->constr &&
+			relation->rd_att->constr->has_generated_stored)
+			return true;
+
+		idattrs = RelationGetIndexAttrBitmap(relation,
+											 INDEX_ATTR_BITMAP_IDENTITY_KEY);
+
+		/*
+		 * Bitmap of published columns when publish_generated_columns is
+		 * 'false' and no column list is specified.
+		 */
+		columns = pub_form_cols_map(relation, false);
+
+		/*
+		 * Attnums in the bitmap returned by RelationGetIndexAttrBitmap are
+		 * offset (to handle system columns the usual way), while column list
+		 * does not use offset, so we can't do bms_is_subset(). Instead, we
+		 * have to loop over the idattrs and check all of them are in the
+		 * list.
+		 */
+		x = -1;
+		while ((x = bms_next_member(idattrs, x)) >= 0)
+		{
+			AttrNumber	attnum = (x + FirstLowInvalidHeapAttributeNumber);
+
+			/* replica identity column, not covered by the column list */
+			if (!bms_is_member(attnum, columns))
+			{
+				result = true;
+				break;
+			}
+		}
+
+		bms_free(idattrs);
+		bms_free(columns);
+	}
+
+	return result;
+}
+
 /* check_functions_in_node callback */
 static bool
 contain_mutable_or_user_functions_checker(Oid func_id, void *context)
diff --git a/src/backend/executor/execReplication.c b/src/backend/executor/execReplication.c
index 54025c9f15..6e42de104f 100644
--- a/src/backend/executor/execReplication.c
+++ b/src/backend/executor/execReplication.c
@@ -809,6 +809,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot update table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_UPDATE && !pubdesc.replident_has_valid_gen_cols)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot update table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("Replica identity consists of an unpublished generated column.")));
 	else if (cmd == CMD_DELETE && !pubdesc.rf_valid_for_delete)
 		ereport(ERROR,
 				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
@@ -821,6 +827,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot delete from table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_DELETE && !pubdesc.replident_has_valid_gen_cols)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot delete from table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("Replica identity consists of an unpublished generated column.")));
 
 	/* If relation has replica identity we are always good. */
 	if (OidIsValid(RelationGetReplicaIndex(rel)))
diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
index 342467fd18..02d7fba160 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -5736,6 +5736,7 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		pubdesc->rf_valid_for_delete = true;
 		pubdesc->cols_valid_for_update = true;
 		pubdesc->cols_valid_for_delete = true;
+		pubdesc->replident_has_valid_gen_cols = true;
 		return;
 	}
 
@@ -5750,6 +5751,7 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 	pubdesc->rf_valid_for_delete = true;
 	pubdesc->cols_valid_for_update = true;
 	pubdesc->cols_valid_for_delete = true;
+	pubdesc->replident_has_valid_gen_cols = true;
 
 	/* Fetch the publication membership info. */
 	puboids = GetRelationPublications(relid);
@@ -5827,6 +5829,16 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 				pubdesc->cols_valid_for_delete = false;
 		}
 
+		/*
+		 * Check if all columns which are part of the REPLICA IDENTITY is
+		 * published.
+		 */
+		if (!pubform->pubgencols &&
+			(pubform->pubupdate || pubform->pubdelete) &&
+			replident_has_unpublished_gen_col(pubid, relation, ancestors,
+											  pubform->pubviaroot))
+			pubdesc->replident_has_valid_gen_cols = false;
+
 		ReleaseSysCache(tup);
 
 		/*
diff --git a/src/include/catalog/pg_publication.h b/src/include/catalog/pg_publication.h
index 9a83a72d6b..5eeeced91d 100644
--- a/src/include/catalog/pg_publication.h
+++ b/src/include/catalog/pg_publication.h
@@ -98,6 +98,12 @@ typedef struct PublicationDesc
 	 */
 	bool		cols_valid_for_update;
 	bool		cols_valid_for_delete;
+
+	/*
+	 * true if all generated columns which are part of replica identity are
+	 * published.
+	 */
+	bool		replident_has_valid_gen_cols;
 } PublicationDesc;
 
 typedef struct Publication
diff --git a/src/include/commands/publicationcmds.h b/src/include/commands/publicationcmds.h
index 5487c571f6..b18e576b77 100644
--- a/src/include/commands/publicationcmds.h
+++ b/src/include/commands/publicationcmds.h
@@ -35,5 +35,7 @@ extern bool pub_rf_contains_invalid_column(Oid pubid, Relation relation,
 										   List *ancestors, bool pubviaroot);
 extern bool pub_collist_contains_invalid_column(Oid pubid, Relation relation,
 												List *ancestors, bool pubviaroot);
+extern bool replident_has_unpublished_gen_col(Oid pubid, Relation relation,
+											  List *ancestors, bool pubviaroot);
 
 #endif							/* PUBLICATIONCMDS_H */
diff --git a/src/test/regress/expected/publication.out b/src/test/regress/expected/publication.out
index 5de2d64d01..caf6e31554 100644
--- a/src/test/regress/expected/publication.out
+++ b/src/test/regress/expected/publication.out
@@ -672,6 +672,31 @@ DROP TABLE rf_tbl_abcd_pk;
 DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  Replica identity consists of an unpublished generated column.
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  Replica identity consists of an unpublished generated column.
+DROP PUBLICATION pub_gencol;
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+DROP INDEX testpub_gencol_idx;
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/regress/sql/publication.sql b/src/test/regress/sql/publication.sql
index 48e68bcca2..3b1c661440 100644
--- a/src/test/regress/sql/publication.sql
+++ b/src/test/regress/sql/publication.sql
@@ -396,6 +396,32 @@ DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
 
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+
+DROP INDEX testpub_gencol_idx;
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
+
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/subscription/t/100_bugs.pl b/src/test/subscription/t/100_bugs.pl
index cb36ca7b16..794b928f50 100644
--- a/src/test/subscription/t/100_bugs.pl
+++ b/src/test/subscription/t/100_bugs.pl
@@ -377,8 +377,8 @@ $node_publisher->safe_psql('postgres', "DROP PUBLICATION tap_pub_sch");
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');
 
-# The bug was that when the REPLICA IDENTITY FULL is used with dropped or
-# generated columns, we fail to apply updates and deletes
+# The bug was that when the REPLICA IDENTITY FULL is used with dropped
+# we fail to apply updates and deletes
 $node_publisher->rotate_logfile();
 $node_publisher->start();
 
@@ -389,18 +389,14 @@ $node_publisher->safe_psql(
 	'postgres', qq(
 	CREATE TABLE dropped_cols (a int, b_drop int, c int);
 	ALTER TABLE dropped_cols REPLICA IDENTITY FULL;
-	CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
-	ALTER TABLE generated_cols REPLICA IDENTITY FULL;
-	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols, generated_cols;
+	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols;
 	-- some initial data
 	INSERT INTO dropped_cols VALUES (1, 1, 1);
-	INSERT INTO generated_cols (a, c) VALUES (1, 1);
 ));
 
 $node_subscriber->safe_psql(
 	'postgres', qq(
 	 CREATE TABLE dropped_cols (a int, b_drop int, c int);
-	 CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
 ));
 
 $publisher_connstr = $node_publisher->connstr . ' dbname=postgres';
@@ -421,7 +417,6 @@ $node_subscriber->safe_psql(
 $node_publisher->safe_psql(
 	'postgres', qq(
 		UPDATE dropped_cols SET a = 100;
-		UPDATE generated_cols SET a = 100;
 ));
 $node_publisher->wait_for_catchup('sub_dropped_cols');
 
@@ -430,11 +425,6 @@ is( $node_subscriber->safe_psql(
 	qq(1),
 	'replication with RI FULL and dropped columns');
 
-is( $node_subscriber->safe_psql(
-		'postgres', "SELECT count(*) FROM generated_cols WHERE a = 100"),
-	qq(1),
-	'replication with RI FULL and generated columns');
-
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');
 
-- 
2.34.1

From f9b85b8f376456205d6e410ee08d3e083bff4fed Mon Sep 17 00:00:00 2001
From: Shlok Kyal <shlok.kyal.oss@gmail.com>
Date: Tue, 5 Nov 2024 11:11:25 +0530
Subject: [PATCH v7] Disallow UPDATE/DELETE on table with generated column as
 REPLICA IDENTITY

UPDATE/DELETE on table having unpublished generated column as REPLICA
IDENTITY is allowed. UPDATE/DELETE on such tables should not be allowed
---
 src/backend/commands/publicationcmds.c    | 76 +++++++++++++++++++++++
 src/backend/executor/execReplication.c    | 12 ++++
 src/backend/utils/cache/relcache.c        | 19 ++++++
 src/include/catalog/pg_publication.h      |  7 +++
 src/include/commands/publicationcmds.h    |  2 +
 src/test/regress/expected/publication.out | 25 ++++++++
 src/test/regress/sql/publication.sql      | 26 ++++++++
 src/test/subscription/t/100_bugs.pl       | 16 +----
 8 files changed, 170 insertions(+), 13 deletions(-)

diff --git a/src/backend/commands/publicationcmds.c b/src/backend/commands/publicationcmds.c
index 0129db18c6..84abcef208 100644
--- a/src/backend/commands/publicationcmds.c
+++ b/src/backend/commands/publicationcmds.c
@@ -444,6 +444,82 @@ pub_collist_contains_invalid_column(Oid pubid, Relation relation, List *ancestor
 	return result;
 }
 
+/*
+ * Check if REPLICA IDENTITY consists of any unpublished generated column.
+ *
+ * Returns true if any replica identity column is an unpublished generated
+ * column.
+ */
+bool
+replident_has_unpublished_gen_col(Oid pubid, Relation relation, List *ancestors,
+								  bool pubviaroot)
+{
+	Oid			relid = RelationGetRelid(relation);
+	Oid			publish_as_relid;
+	bool		result = false;
+	bool		found;
+	Publication *pub;
+
+	/*
+	 * For a partition, if pubviaroot is true, find the topmost ancestor that
+	 * is published via this publication as we need to use its column list for
+	 * the changes.
+	 *
+	 * Note that even though the column list used is for an ancestor, the
+	 * REPLICA IDENTITY used will be for the actual child table.
+	 */
+	if (pubviaroot && relation->rd_rel->relispartition)
+	{
+		publish_as_relid = GetTopMostAncestorInPublication(pubid, ancestors, NULL);
+
+		if (!OidIsValid(publish_as_relid))
+			publish_as_relid = relid;
+	}
+	else
+		publish_as_relid = relid;
+
+	pub = GetPublication(pubid);
+	found = check_and_fetch_column_list(pub, publish_as_relid, NULL, NULL);
+
+	if (!found)
+	{
+		int			x;
+		Bitmapset  *idattrs = NULL;
+		TupleDesc	desc = RelationGetDescr(relation);
+
+		/*
+		 * REPLICA IDENTITY can be FULL only if there is no column list for
+		 * publication. If REPLICA IDENTITY is set as FULL and relation has a
+		 * generated column we should error out.
+		 */
+		if (relation->rd_rel->relreplident == REPLICA_IDENTITY_FULL &&
+			relation->rd_att->constr &&
+			relation->rd_att->constr->has_generated_stored)
+			return true;
+
+		idattrs = RelationGetIndexAttrBitmap(relation,
+											 INDEX_ATTR_BITMAP_IDENTITY_KEY);
+
+		for (int i = 0; i < desc->natts; i++)
+		{
+			Form_pg_attribute att = TupleDescAttr(desc, i);
+
+			/* check if generated column is part of REPLICA IDENTITY */
+			if (!att->attisdropped && att->attgenerated &&
+				bms_is_member(att->attnum - FirstLowInvalidHeapAttributeNumber,
+							  idattrs))
+			{
+				result = true;
+				break;
+			}
+		}
+
+		bms_free(idattrs);
+	}
+
+	return result;
+}
+
 /* check_functions_in_node callback */
 static bool
 contain_mutable_or_user_functions_checker(Oid func_id, void *context)
diff --git a/src/backend/executor/execReplication.c b/src/backend/executor/execReplication.c
index 54025c9f15..57599df153 100644
--- a/src/backend/executor/execReplication.c
+++ b/src/backend/executor/execReplication.c
@@ -809,6 +809,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot update table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_UPDATE && !pubdesc.replident_valid_for_update)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot update table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("Replica identity consists of an unpublished generated column.")));
 	else if (cmd == CMD_DELETE && !pubdesc.rf_valid_for_delete)
 		ereport(ERROR,
 				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
@@ -821,6 +827,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot delete from table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_DELETE && !pubdesc.replident_valid_for_delete)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot delete from table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("Replica identity consists of an unpublished generated column.")));
 
 	/* If relation has replica identity we are always good. */
 	if (OidIsValid(RelationGetReplicaIndex(rel)))
diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
index 342467fd18..ef3203a682 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -5736,6 +5736,8 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		pubdesc->rf_valid_for_delete = true;
 		pubdesc->cols_valid_for_update = true;
 		pubdesc->cols_valid_for_delete = true;
+		pubdesc->replident_valid_for_update = true;
+		pubdesc->replident_valid_for_delete = true;
 		return;
 	}
 
@@ -5750,6 +5752,8 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 	pubdesc->rf_valid_for_delete = true;
 	pubdesc->cols_valid_for_update = true;
 	pubdesc->cols_valid_for_delete = true;
+	pubdesc->replident_valid_for_update = true;
+	pubdesc->replident_valid_for_delete = true;
 
 	/* Fetch the publication membership info. */
 	puboids = GetRelationPublications(relid);
@@ -5827,6 +5831,21 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 				pubdesc->cols_valid_for_delete = false;
 		}
 
+		/*
+		 * Check if all columns which are part of the REPLICA IDENTITY is
+		 * published.
+		 */
+		if (!pubform->pubgencols &&
+			(pubform->pubupdate || pubform->pubdelete) &&
+			replident_has_unpublished_gen_col(pubid, relation, ancestors,
+											  pubform->pubviaroot))
+		{
+			if (pubform->pubupdate)
+				pubdesc->replident_valid_for_update = false;
+			if (pubform->pubdelete)
+				pubdesc->replident_valid_for_delete = false;
+		}
+
 		ReleaseSysCache(tup);
 
 		/*
diff --git a/src/include/catalog/pg_publication.h b/src/include/catalog/pg_publication.h
index 9a83a72d6b..486f609a9a 100644
--- a/src/include/catalog/pg_publication.h
+++ b/src/include/catalog/pg_publication.h
@@ -98,6 +98,13 @@ typedef struct PublicationDesc
 	 */
 	bool		cols_valid_for_update;
 	bool		cols_valid_for_delete;
+
+	/*
+	 * true if all generated columns which are part of replica identity are
+	 * published or the publication actions do not include UPDATE or DELETE.
+	 */
+	bool		replident_valid_for_update;
+	bool		replident_valid_for_delete;
 } PublicationDesc;
 
 typedef struct Publication
diff --git a/src/include/commands/publicationcmds.h b/src/include/commands/publicationcmds.h
index 5487c571f6..b18e576b77 100644
--- a/src/include/commands/publicationcmds.h
+++ b/src/include/commands/publicationcmds.h
@@ -35,5 +35,7 @@ extern bool pub_rf_contains_invalid_column(Oid pubid, Relation relation,
 										   List *ancestors, bool pubviaroot);
 extern bool pub_collist_contains_invalid_column(Oid pubid, Relation relation,
 												List *ancestors, bool pubviaroot);
+extern bool replident_has_unpublished_gen_col(Oid pubid, Relation relation,
+											  List *ancestors, bool pubviaroot);
 
 #endif							/* PUBLICATIONCMDS_H */
diff --git a/src/test/regress/expected/publication.out b/src/test/regress/expected/publication.out
index 5de2d64d01..caf6e31554 100644
--- a/src/test/regress/expected/publication.out
+++ b/src/test/regress/expected/publication.out
@@ -672,6 +672,31 @@ DROP TABLE rf_tbl_abcd_pk;
 DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  Replica identity consists of an unpublished generated column.
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  Replica identity consists of an unpublished generated column.
+DROP PUBLICATION pub_gencol;
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+DROP INDEX testpub_gencol_idx;
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/regress/sql/publication.sql b/src/test/regress/sql/publication.sql
index 48e68bcca2..3b1c661440 100644
--- a/src/test/regress/sql/publication.sql
+++ b/src/test/regress/sql/publication.sql
@@ -396,6 +396,32 @@ DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
 
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+
+DROP INDEX testpub_gencol_idx;
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
+
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/subscription/t/100_bugs.pl b/src/test/subscription/t/100_bugs.pl
index cb36ca7b16..794b928f50 100644
--- a/src/test/subscription/t/100_bugs.pl
+++ b/src/test/subscription/t/100_bugs.pl
@@ -377,8 +377,8 @@ $node_publisher->safe_psql('postgres', "DROP PUBLICATION tap_pub_sch");
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');
 
-# The bug was that when the REPLICA IDENTITY FULL is used with dropped or
-# generated columns, we fail to apply updates and deletes
+# The bug was that when the REPLICA IDENTITY FULL is used with dropped
+# we fail to apply updates and deletes
 $node_publisher->rotate_logfile();
 $node_publisher->start();
 
@@ -389,18 +389,14 @@ $node_publisher->safe_psql(
 	'postgres', qq(
 	CREATE TABLE dropped_cols (a int, b_drop int, c int);
 	ALTER TABLE dropped_cols REPLICA IDENTITY FULL;
-	CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
-	ALTER TABLE generated_cols REPLICA IDENTITY FULL;
-	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols, generated_cols;
+	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols;
 	-- some initial data
 	INSERT INTO dropped_cols VALUES (1, 1, 1);
-	INSERT INTO generated_cols (a, c) VALUES (1, 1);
 ));
 
 $node_subscriber->safe_psql(
 	'postgres', qq(
 	 CREATE TABLE dropped_cols (a int, b_drop int, c int);
-	 CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
 ));
 
 $publisher_connstr = $node_publisher->connstr . ' dbname=postgres';
@@ -421,7 +417,6 @@ $node_subscriber->safe_psql(
 $node_publisher->safe_psql(
 	'postgres', qq(
 		UPDATE dropped_cols SET a = 100;
-		UPDATE generated_cols SET a = 100;
 ));
 $node_publisher->wait_for_catchup('sub_dropped_cols');
 
@@ -430,11 +425,6 @@ is( $node_subscriber->safe_psql(
 	qq(1),
 	'replication with RI FULL and dropped columns');
 
-is( $node_subscriber->safe_psql(
-		'postgres', "SELECT count(*) FROM generated_cols WHERE a = 100"),
-	qq(1),
-	'replication with RI FULL and generated columns');
-
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');
 
-- 
2.34.1

From f995e00c4527bb88e609f90939366d39cd6e9025 Mon Sep 17 00:00:00 2001
From: Shlok Kyal <shlok.kyal.oss@gmail.com>
Date: Tue, 5 Nov 2024 11:11:25 +0530
Subject: [PATCH v8] Disallow UPDATE/DELETE on table with generated column as
 REPLICA IDENTITY

UPDATE/DELETE on table having unpublished generated column as REPLICA
IDENTITY is allowed. UPDATE/DELETE on such tables should not be allowed
---
 src/backend/commands/publicationcmds.c    | 75 +++++++++++++++++++++++
 src/backend/executor/execReplication.c    | 12 ++++
 src/backend/utils/cache/relcache.c        | 30 +++++++++
 src/include/catalog/pg_publication.h      |  7 +++
 src/include/commands/publicationcmds.h    |  2 +
 src/test/regress/expected/publication.out | 25 ++++++++
 src/test/regress/sql/publication.sql      | 26 ++++++++
 src/test/subscription/t/100_bugs.pl       | 16 +----
 8 files changed, 180 insertions(+), 13 deletions(-)

diff --git a/src/backend/commands/publicationcmds.c b/src/backend/commands/publicationcmds.c
index 0129db18c6..8e6a61c997 100644
--- a/src/backend/commands/publicationcmds.c
+++ b/src/backend/commands/publicationcmds.c
@@ -444,6 +444,81 @@ pub_collist_contains_invalid_column(Oid pubid, Relation relation, List *ancestor
 	return result;
 }
 
+/*
+ * Check if REPLICA IDENTITY consists of any unpublished generated column.
+ *
+ * Returns true if any replica identity column is an unpublished generated
+ * column.
+ */
+bool
+replident_has_unpublished_gen_col(Oid pubid, Relation relation, List *ancestors,
+								  bool pubviaroot)
+{
+	Oid			relid = RelationGetRelid(relation);
+	Oid			publish_as_relid;
+	bool		result = false;
+	bool		found;
+	Publication *pub;
+
+	/*
+	 * For a partition, if pubviaroot is true, find the topmost ancestor that
+	 * is published via this publication as we need to use its column list for
+	 * the changes.
+	 *
+	 * Note that even though the column list used is for an ancestor, the
+	 * REPLICA IDENTITY used will be for the actual child table.
+	 */
+	if (pubviaroot && relation->rd_rel->relispartition)
+	{
+		publish_as_relid = GetTopMostAncestorInPublication(pubid, ancestors, NULL);
+
+		if (!OidIsValid(publish_as_relid))
+			publish_as_relid = relid;
+	}
+	else
+		publish_as_relid = relid;
+
+	pub = GetPublication(pubid);
+	found = check_and_fetch_column_list(pub, publish_as_relid, NULL, NULL);
+
+	if (!found)
+	{
+		Bitmapset  *idattrs = NULL;
+		TupleDesc	desc = RelationGetDescr(relation);
+
+		/*
+		 * REPLICA IDENTITY can be FULL only if there is no column list for
+		 * publication. If REPLICA IDENTITY is set as FULL and relation has a
+		 * generated column we should error out.
+		 */
+		if (relation->rd_rel->relreplident == REPLICA_IDENTITY_FULL &&
+			relation->rd_att->constr &&
+			relation->rd_att->constr->has_generated_stored)
+			return true;
+
+		idattrs = RelationGetIndexAttrBitmap(relation,
+											 INDEX_ATTR_BITMAP_IDENTITY_KEY);
+
+		for (int i = 0; i < desc->natts; i++)
+		{
+			Form_pg_attribute att = TupleDescAttr(desc, i);
+
+			/* check if generated column is part of REPLICA IDENTITY */
+			if (!att->attisdropped && att->attgenerated &&
+				bms_is_member(att->attnum - FirstLowInvalidHeapAttributeNumber,
+							  idattrs))
+			{
+				result = true;
+				break;
+			}
+		}
+
+		bms_free(idattrs);
+	}
+
+	return result;
+}
+
 /* check_functions_in_node callback */
 static bool
 contain_mutable_or_user_functions_checker(Oid func_id, void *context)
diff --git a/src/backend/executor/execReplication.c b/src/backend/executor/execReplication.c
index 54025c9f15..57599df153 100644
--- a/src/backend/executor/execReplication.c
+++ b/src/backend/executor/execReplication.c
@@ -809,6 +809,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot update table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_UPDATE && !pubdesc.replident_valid_for_update)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot update table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("Replica identity consists of an unpublished generated column.")));
 	else if (cmd == CMD_DELETE && !pubdesc.rf_valid_for_delete)
 		ereport(ERROR,
 				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
@@ -821,6 +827,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot delete from table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_DELETE && !pubdesc.replident_valid_for_delete)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot delete from table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("Replica identity consists of an unpublished generated column.")));
 
 	/* If relation has replica identity we are always good. */
 	if (OidIsValid(RelationGetReplicaIndex(rel)))
diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
index 342467fd18..1e4cf99e85 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -5736,6 +5736,8 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		pubdesc->rf_valid_for_delete = true;
 		pubdesc->cols_valid_for_update = true;
 		pubdesc->cols_valid_for_delete = true;
+		pubdesc->replident_valid_for_update = true;
+		pubdesc->replident_valid_for_delete = true;
 		return;
 	}
 
@@ -5750,6 +5752,8 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 	pubdesc->rf_valid_for_delete = true;
 	pubdesc->cols_valid_for_update = true;
 	pubdesc->cols_valid_for_delete = true;
+	pubdesc->replident_valid_for_update = true;
+	pubdesc->replident_valid_for_delete = true;
 
 	/* Fetch the publication membership info. */
 	puboids = GetRelationPublications(relid);
@@ -5827,6 +5831,21 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 				pubdesc->cols_valid_for_delete = false;
 		}
 
+		/*
+		 * Check if all columns which are part of the REPLICA IDENTITY is
+		 * published.
+		 */
+		if (!pubform->pubgencols &&
+			(pubform->pubupdate || pubform->pubdelete) &&
+			replident_has_unpublished_gen_col(pubid, relation, ancestors,
+											  pubform->pubviaroot))
+		{
+			if (pubform->pubupdate)
+				pubdesc->replident_valid_for_update = false;
+			if (pubform->pubdelete)
+				pubdesc->replident_valid_for_delete = false;
+		}
+
 		ReleaseSysCache(tup);
 
 		/*
@@ -5848,6 +5867,17 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 			pubdesc->pubactions.pubdelete && pubdesc->pubactions.pubtruncate &&
 			!pubdesc->cols_valid_for_update && !pubdesc->cols_valid_for_delete)
 			break;
+
+		/*
+		 * If we know everything is replicated and replica identity has an
+		 * unpublished generated column, there is no point to check for other
+		 * publications.
+		 */
+		if (pubdesc->pubactions.pubinsert && pubdesc->pubactions.pubupdate &&
+			pubdesc->pubactions.pubdelete && pubdesc->pubactions.pubtruncate &&
+			!pubdesc->replident_valid_for_update &&
+			!pubdesc->replident_valid_for_delete)
+			break;
 	}
 
 	if (relation->rd_pubdesc)
diff --git a/src/include/catalog/pg_publication.h b/src/include/catalog/pg_publication.h
index 9a83a72d6b..486f609a9a 100644
--- a/src/include/catalog/pg_publication.h
+++ b/src/include/catalog/pg_publication.h
@@ -98,6 +98,13 @@ typedef struct PublicationDesc
 	 */
 	bool		cols_valid_for_update;
 	bool		cols_valid_for_delete;
+
+	/*
+	 * true if all generated columns which are part of replica identity are
+	 * published or the publication actions do not include UPDATE or DELETE.
+	 */
+	bool		replident_valid_for_update;
+	bool		replident_valid_for_delete;
 } PublicationDesc;
 
 typedef struct Publication
diff --git a/src/include/commands/publicationcmds.h b/src/include/commands/publicationcmds.h
index 5487c571f6..b18e576b77 100644
--- a/src/include/commands/publicationcmds.h
+++ b/src/include/commands/publicationcmds.h
@@ -35,5 +35,7 @@ extern bool pub_rf_contains_invalid_column(Oid pubid, Relation relation,
 										   List *ancestors, bool pubviaroot);
 extern bool pub_collist_contains_invalid_column(Oid pubid, Relation relation,
 												List *ancestors, bool pubviaroot);
+extern bool replident_has_unpublished_gen_col(Oid pubid, Relation relation,
+											  List *ancestors, bool pubviaroot);
 
 #endif							/* PUBLICATIONCMDS_H */
diff --git a/src/test/regress/expected/publication.out b/src/test/regress/expected/publication.out
index 5de2d64d01..caf6e31554 100644
--- a/src/test/regress/expected/publication.out
+++ b/src/test/regress/expected/publication.out
@@ -672,6 +672,31 @@ DROP TABLE rf_tbl_abcd_pk;
 DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  Replica identity consists of an unpublished generated column.
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  Replica identity consists of an unpublished generated column.
+DROP PUBLICATION pub_gencol;
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+DROP INDEX testpub_gencol_idx;
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/regress/sql/publication.sql b/src/test/regress/sql/publication.sql
index 48e68bcca2..3b1c661440 100644
--- a/src/test/regress/sql/publication.sql
+++ b/src/test/regress/sql/publication.sql
@@ -396,6 +396,32 @@ DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
 
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+
+DROP INDEX testpub_gencol_idx;
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
+
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/subscription/t/100_bugs.pl b/src/test/subscription/t/100_bugs.pl
index cb36ca7b16..794b928f50 100644
--- a/src/test/subscription/t/100_bugs.pl
+++ b/src/test/subscription/t/100_bugs.pl
@@ -377,8 +377,8 @@ $node_publisher->safe_psql('postgres', "DROP PUBLICATION tap_pub_sch");
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');
 
-# The bug was that when the REPLICA IDENTITY FULL is used with dropped or
-# generated columns, we fail to apply updates and deletes
+# The bug was that when the REPLICA IDENTITY FULL is used with dropped
+# we fail to apply updates and deletes
 $node_publisher->rotate_logfile();
 $node_publisher->start();
 
@@ -389,18 +389,14 @@ $node_publisher->safe_psql(
 	'postgres', qq(
 	CREATE TABLE dropped_cols (a int, b_drop int, c int);
 	ALTER TABLE dropped_cols REPLICA IDENTITY FULL;
-	CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
-	ALTER TABLE generated_cols REPLICA IDENTITY FULL;
-	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols, generated_cols;
+	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols;
 	-- some initial data
 	INSERT INTO dropped_cols VALUES (1, 1, 1);
-	INSERT INTO generated_cols (a, c) VALUES (1, 1);
 ));
 
 $node_subscriber->safe_psql(
 	'postgres', qq(
 	 CREATE TABLE dropped_cols (a int, b_drop int, c int);
-	 CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
 ));
 
 $publisher_connstr = $node_publisher->connstr . ' dbname=postgres';
@@ -421,7 +417,6 @@ $node_subscriber->safe_psql(
 $node_publisher->safe_psql(
 	'postgres', qq(
 		UPDATE dropped_cols SET a = 100;
-		UPDATE generated_cols SET a = 100;
 ));
 $node_publisher->wait_for_catchup('sub_dropped_cols');
 
@@ -430,11 +425,6 @@ is( $node_subscriber->safe_psql(
 	qq(1),
 	'replication with RI FULL and dropped columns');
 
-is( $node_subscriber->safe_psql(
-		'postgres', "SELECT count(*) FROM generated_cols WHERE a = 100"),
-	qq(1),
-	'replication with RI FULL and generated columns');
-
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');
 
-- 
2.34.1

From 553dc53ec14b6ed8b898992f6b2085cbe12e5408 Mon Sep 17 00:00:00 2001
From: Hou Zhijie <houzj.fnst@cn.fujitsu.com>
Date: Tue, 19 Nov 2024 11:13:13 +0800
Subject: [PATCH v2] improve logic

---
 src/backend/commands/publicationcmds.c | 29 ++++++++++++++------------
 src/backend/utils/cache/relcache.c     |  2 +-
 2 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/src/backend/commands/publicationcmds.c b/src/backend/commands/publicationcmds.c
index 8e6a61c997..053877c524 100644
--- a/src/backend/commands/publicationcmds.c
+++ b/src/backend/commands/publicationcmds.c
@@ -455,11 +455,16 @@ replident_has_unpublished_gen_col(Oid pubid, Relation relation, List *ancestors,
 								  bool pubviaroot)
 {
 	Oid			relid = RelationGetRelid(relation);
-	Oid			publish_as_relid;
+	Oid			publish_as_relid = RelationGetRelid(relation);
 	bool		result = false;
 	bool		found;
 	Publication *pub;
 
+	/* Return if the table does not contain any generated columns */
+	if (!relation->rd_att->constr ||
+		!relation->rd_att->constr->has_generated_stored)
+		return false;
+
 	/*
 	 * For a partition, if pubviaroot is true, find the topmost ancestor that
 	 * is published via this publication as we need to use its column list for
@@ -475,38 +480,36 @@ replident_has_unpublished_gen_col(Oid pubid, Relation relation, List *ancestors,
 		if (!OidIsValid(publish_as_relid))
 			publish_as_relid = relid;
 	}
-	else
-		publish_as_relid = relid;
 
 	pub = GetPublication(pubid);
 	found = check_and_fetch_column_list(pub, publish_as_relid, NULL, NULL);
 
 	if (!found)
 	{
-		Bitmapset  *idattrs = NULL;
 		TupleDesc	desc = RelationGetDescr(relation);
+		Bitmapset  *idattrs;
+		int			x;
 
 		/*
 		 * REPLICA IDENTITY can be FULL only if there is no column list for
 		 * publication. If REPLICA IDENTITY is set as FULL and relation has a
 		 * generated column we should error out.
 		 */
-		if (relation->rd_rel->relreplident == REPLICA_IDENTITY_FULL &&
-			relation->rd_att->constr &&
-			relation->rd_att->constr->has_generated_stored)
+		if (relation->rd_rel->relreplident == REPLICA_IDENTITY_FULL)
 			return true;
 
+		/* Remember columns that are part of the REPLICA IDENTITY */
 		idattrs = RelationGetIndexAttrBitmap(relation,
 											 INDEX_ATTR_BITMAP_IDENTITY_KEY);
 
-		for (int i = 0; i < desc->natts; i++)
+		x = -1;
+		while ((x = bms_next_member(idattrs, x)) >= 0)
 		{
-			Form_pg_attribute att = TupleDescAttr(desc, i);
+			AttrNumber	attnum = (x + FirstLowInvalidHeapAttributeNumber);
+			Form_pg_attribute att = TupleDescAttr(desc, attnum - 1);
 
-			/* check if generated column is part of REPLICA IDENTITY */
-			if (!att->attisdropped && att->attgenerated &&
-				bms_is_member(att->attnum - FirstLowInvalidHeapAttributeNumber,
-							  idattrs))
+			/* Check if generated column is part of REPLICA IDENTITY */
+			if (!att->attisdropped && att->attgenerated)
 			{
 				result = true;
 				break;
diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
index 1e4cf99e85..be8f8eea8f 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -5832,7 +5832,7 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		}
 
 		/*
-		 * Check if all columns which are part of the REPLICA IDENTITY is
+		 * Check if all generated columns included in the REPLICA IDENTITY are
 		 * published.
 		 */
 		if (!pubform->pubgencols &&
-- 
2.30.0.windows.2

From fc2a6043cef8a48bf3145e5e1bca90ac33e7efa9 Mon Sep 17 00:00:00 2001
From: Shlok Kyal <shlok.kyal.oss@gmail.com>
Date: Tue, 5 Nov 2024 11:11:25 +0530
Subject: [PATCH v9] Disallow UPDATE/DELETE on table with generated column as
 REPLICA IDENTITY

UPDATE/DELETE on table having unpublished generated column as REPLICA
IDENTITY is allowed. UPDATE/DELETE on such tables should not be allowed
---
 src/backend/commands/publicationcmds.c    | 78 +++++++++++++++++++++++
 src/backend/executor/execReplication.c    | 12 ++++
 src/backend/utils/cache/relcache.c        | 30 +++++++++
 src/include/catalog/pg_publication.h      |  7 ++
 src/include/commands/publicationcmds.h    |  2 +
 src/test/regress/expected/publication.out | 25 ++++++++
 src/test/regress/sql/publication.sql      | 26 ++++++++
 src/test/subscription/t/100_bugs.pl       |  4 +-
 8 files changed, 182 insertions(+), 2 deletions(-)

diff --git a/src/backend/commands/publicationcmds.c b/src/backend/commands/publicationcmds.c
index 0129db18c6..053877c524 100644
--- a/src/backend/commands/publicationcmds.c
+++ b/src/backend/commands/publicationcmds.c
@@ -444,6 +444,84 @@ pub_collist_contains_invalid_column(Oid pubid, Relation relation, List *ancestor
 	return result;
 }
 
+/*
+ * Check if REPLICA IDENTITY consists of any unpublished generated column.
+ *
+ * Returns true if any replica identity column is an unpublished generated
+ * column.
+ */
+bool
+replident_has_unpublished_gen_col(Oid pubid, Relation relation, List *ancestors,
+								  bool pubviaroot)
+{
+	Oid			relid = RelationGetRelid(relation);
+	Oid			publish_as_relid = RelationGetRelid(relation);
+	bool		result = false;
+	bool		found;
+	Publication *pub;
+
+	/* Return if the table does not contain any generated columns */
+	if (!relation->rd_att->constr ||
+		!relation->rd_att->constr->has_generated_stored)
+		return false;
+
+	/*
+	 * For a partition, if pubviaroot is true, find the topmost ancestor that
+	 * is published via this publication as we need to use its column list for
+	 * the changes.
+	 *
+	 * Note that even though the column list used is for an ancestor, the
+	 * REPLICA IDENTITY used will be for the actual child table.
+	 */
+	if (pubviaroot && relation->rd_rel->relispartition)
+	{
+		publish_as_relid = GetTopMostAncestorInPublication(pubid, ancestors, NULL);
+
+		if (!OidIsValid(publish_as_relid))
+			publish_as_relid = relid;
+	}
+
+	pub = GetPublication(pubid);
+	found = check_and_fetch_column_list(pub, publish_as_relid, NULL, NULL);
+
+	if (!found)
+	{
+		TupleDesc	desc = RelationGetDescr(relation);
+		Bitmapset  *idattrs;
+		int			x;
+
+		/*
+		 * REPLICA IDENTITY can be FULL only if there is no column list for
+		 * publication. If REPLICA IDENTITY is set as FULL and relation has a
+		 * generated column we should error out.
+		 */
+		if (relation->rd_rel->relreplident == REPLICA_IDENTITY_FULL)
+			return true;
+
+		/* Remember columns that are part of the REPLICA IDENTITY */
+		idattrs = RelationGetIndexAttrBitmap(relation,
+											 INDEX_ATTR_BITMAP_IDENTITY_KEY);
+
+		x = -1;
+		while ((x = bms_next_member(idattrs, x)) >= 0)
+		{
+			AttrNumber	attnum = (x + FirstLowInvalidHeapAttributeNumber);
+			Form_pg_attribute att = TupleDescAttr(desc, attnum - 1);
+
+			/* Check if generated column is part of REPLICA IDENTITY */
+			if (!att->attisdropped && att->attgenerated)
+			{
+				result = true;
+				break;
+			}
+		}
+
+		bms_free(idattrs);
+	}
+
+	return result;
+}
+
 /* check_functions_in_node callback */
 static bool
 contain_mutable_or_user_functions_checker(Oid func_id, void *context)
diff --git a/src/backend/executor/execReplication.c b/src/backend/executor/execReplication.c
index 54025c9f15..57599df153 100644
--- a/src/backend/executor/execReplication.c
+++ b/src/backend/executor/execReplication.c
@@ -809,6 +809,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot update table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_UPDATE && !pubdesc.replident_valid_for_update)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot update table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("Replica identity consists of an unpublished generated column.")));
 	else if (cmd == CMD_DELETE && !pubdesc.rf_valid_for_delete)
 		ereport(ERROR,
 				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
@@ -821,6 +827,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot delete from table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_DELETE && !pubdesc.replident_valid_for_delete)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot delete from table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("Replica identity consists of an unpublished generated column.")));
 
 	/* If relation has replica identity we are always good. */
 	if (OidIsValid(RelationGetReplicaIndex(rel)))
diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
index 342467fd18..be8f8eea8f 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -5736,6 +5736,8 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		pubdesc->rf_valid_for_delete = true;
 		pubdesc->cols_valid_for_update = true;
 		pubdesc->cols_valid_for_delete = true;
+		pubdesc->replident_valid_for_update = true;
+		pubdesc->replident_valid_for_delete = true;
 		return;
 	}
 
@@ -5750,6 +5752,8 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 	pubdesc->rf_valid_for_delete = true;
 	pubdesc->cols_valid_for_update = true;
 	pubdesc->cols_valid_for_delete = true;
+	pubdesc->replident_valid_for_update = true;
+	pubdesc->replident_valid_for_delete = true;
 
 	/* Fetch the publication membership info. */
 	puboids = GetRelationPublications(relid);
@@ -5827,6 +5831,21 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 				pubdesc->cols_valid_for_delete = false;
 		}
 
+		/*
+		 * Check if all generated columns included in the REPLICA IDENTITY are
+		 * published.
+		 */
+		if (!pubform->pubgencols &&
+			(pubform->pubupdate || pubform->pubdelete) &&
+			replident_has_unpublished_gen_col(pubid, relation, ancestors,
+											  pubform->pubviaroot))
+		{
+			if (pubform->pubupdate)
+				pubdesc->replident_valid_for_update = false;
+			if (pubform->pubdelete)
+				pubdesc->replident_valid_for_delete = false;
+		}
+
 		ReleaseSysCache(tup);
 
 		/*
@@ -5848,6 +5867,17 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 			pubdesc->pubactions.pubdelete && pubdesc->pubactions.pubtruncate &&
 			!pubdesc->cols_valid_for_update && !pubdesc->cols_valid_for_delete)
 			break;
+
+		/*
+		 * If we know everything is replicated and replica identity has an
+		 * unpublished generated column, there is no point to check for other
+		 * publications.
+		 */
+		if (pubdesc->pubactions.pubinsert && pubdesc->pubactions.pubupdate &&
+			pubdesc->pubactions.pubdelete && pubdesc->pubactions.pubtruncate &&
+			!pubdesc->replident_valid_for_update &&
+			!pubdesc->replident_valid_for_delete)
+			break;
 	}
 
 	if (relation->rd_pubdesc)
diff --git a/src/include/catalog/pg_publication.h b/src/include/catalog/pg_publication.h
index 9a83a72d6b..486f609a9a 100644
--- a/src/include/catalog/pg_publication.h
+++ b/src/include/catalog/pg_publication.h
@@ -98,6 +98,13 @@ typedef struct PublicationDesc
 	 */
 	bool		cols_valid_for_update;
 	bool		cols_valid_for_delete;
+
+	/*
+	 * true if all generated columns which are part of replica identity are
+	 * published or the publication actions do not include UPDATE or DELETE.
+	 */
+	bool		replident_valid_for_update;
+	bool		replident_valid_for_delete;
 } PublicationDesc;
 
 typedef struct Publication
diff --git a/src/include/commands/publicationcmds.h b/src/include/commands/publicationcmds.h
index 5487c571f6..b18e576b77 100644
--- a/src/include/commands/publicationcmds.h
+++ b/src/include/commands/publicationcmds.h
@@ -35,5 +35,7 @@ extern bool pub_rf_contains_invalid_column(Oid pubid, Relation relation,
 										   List *ancestors, bool pubviaroot);
 extern bool pub_collist_contains_invalid_column(Oid pubid, Relation relation,
 												List *ancestors, bool pubviaroot);
+extern bool replident_has_unpublished_gen_col(Oid pubid, Relation relation,
+											  List *ancestors, bool pubviaroot);
 
 #endif							/* PUBLICATIONCMDS_H */
diff --git a/src/test/regress/expected/publication.out b/src/test/regress/expected/publication.out
index 5de2d64d01..caf6e31554 100644
--- a/src/test/regress/expected/publication.out
+++ b/src/test/regress/expected/publication.out
@@ -672,6 +672,31 @@ DROP TABLE rf_tbl_abcd_pk;
 DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  Replica identity consists of an unpublished generated column.
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  Replica identity consists of an unpublished generated column.
+DROP PUBLICATION pub_gencol;
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+DROP INDEX testpub_gencol_idx;
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/regress/sql/publication.sql b/src/test/regress/sql/publication.sql
index 48e68bcca2..3b1c661440 100644
--- a/src/test/regress/sql/publication.sql
+++ b/src/test/regress/sql/publication.sql
@@ -396,6 +396,32 @@ DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
 
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+
+DROP INDEX testpub_gencol_idx;
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
+
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/subscription/t/100_bugs.pl b/src/test/subscription/t/100_bugs.pl
index cb36ca7b16..64b902db73 100644
--- a/src/test/subscription/t/100_bugs.pl
+++ b/src/test/subscription/t/100_bugs.pl
@@ -391,7 +391,7 @@ $node_publisher->safe_psql(
 	ALTER TABLE dropped_cols REPLICA IDENTITY FULL;
 	CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
 	ALTER TABLE generated_cols REPLICA IDENTITY FULL;
-	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols, generated_cols;
+	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols, generated_cols with (publish_generated_columns = true);
 	-- some initial data
 	INSERT INTO dropped_cols VALUES (1, 1, 1);
 	INSERT INTO generated_cols (a, c) VALUES (1, 1);
@@ -400,7 +400,7 @@ $node_publisher->safe_psql(
 $node_subscriber->safe_psql(
 	'postgres', qq(
 	 CREATE TABLE dropped_cols (a int, b_drop int, c int);
-	 CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
+	 CREATE TABLE generated_cols (a int, b_gen int, c int);
 ));
 
 $publisher_connstr = $node_publisher->connstr . ' dbname=postgres';
-- 
2.34.1

From 0d50d01253feeade771f4bc591ac6e61162d08a4 Mon Sep 17 00:00:00 2001
From: Shlok Kyal <shlok.kyal.oss@gmail.com>
Date: Tue, 5 Nov 2024 11:11:25 +0530
Subject: [PATCH v10] Disallow UPDATE/DELETE on table with generated column as
 REPLICA IDENTITY

UPDATE or DELETE operations on tables with unpublished generated columns
set as REPLICA IDENTITY are not permitted. This patch ensures that if an
UPDATE or DELETE command is executed on such tables, an error will be
thrown.
With this patch, the behavior has changed for the test added in commit
adedf54. Additionally, there is no other way to trigger the bug that was
fixed by commit adedf54, so the test has been removed.
---
 src/backend/commands/publicationcmds.c    | 78 +++++++++++++++++++++++
 src/backend/executor/execReplication.c    | 12 ++++
 src/backend/utils/cache/relcache.c        | 30 +++++++++
 src/include/catalog/pg_publication.h      |  7 ++
 src/include/commands/publicationcmds.h    |  2 +
 src/test/regress/expected/publication.out | 25 ++++++++
 src/test/regress/sql/publication.sql      | 26 ++++++++
 src/test/subscription/t/100_bugs.pl       | 16 +----
 8 files changed, 183 insertions(+), 13 deletions(-)

diff --git a/src/backend/commands/publicationcmds.c b/src/backend/commands/publicationcmds.c
index 0129db18c6..053877c524 100644
--- a/src/backend/commands/publicationcmds.c
+++ b/src/backend/commands/publicationcmds.c
@@ -444,6 +444,84 @@ pub_collist_contains_invalid_column(Oid pubid, Relation relation, List *ancestor
 	return result;
 }
 
+/*
+ * Check if REPLICA IDENTITY consists of any unpublished generated column.
+ *
+ * Returns true if any replica identity column is an unpublished generated
+ * column.
+ */
+bool
+replident_has_unpublished_gen_col(Oid pubid, Relation relation, List *ancestors,
+								  bool pubviaroot)
+{
+	Oid			relid = RelationGetRelid(relation);
+	Oid			publish_as_relid = RelationGetRelid(relation);
+	bool		result = false;
+	bool		found;
+	Publication *pub;
+
+	/* Return if the table does not contain any generated columns */
+	if (!relation->rd_att->constr ||
+		!relation->rd_att->constr->has_generated_stored)
+		return false;
+
+	/*
+	 * For a partition, if pubviaroot is true, find the topmost ancestor that
+	 * is published via this publication as we need to use its column list for
+	 * the changes.
+	 *
+	 * Note that even though the column list used is for an ancestor, the
+	 * REPLICA IDENTITY used will be for the actual child table.
+	 */
+	if (pubviaroot && relation->rd_rel->relispartition)
+	{
+		publish_as_relid = GetTopMostAncestorInPublication(pubid, ancestors, NULL);
+
+		if (!OidIsValid(publish_as_relid))
+			publish_as_relid = relid;
+	}
+
+	pub = GetPublication(pubid);
+	found = check_and_fetch_column_list(pub, publish_as_relid, NULL, NULL);
+
+	if (!found)
+	{
+		TupleDesc	desc = RelationGetDescr(relation);
+		Bitmapset  *idattrs;
+		int			x;
+
+		/*
+		 * REPLICA IDENTITY can be FULL only if there is no column list for
+		 * publication. If REPLICA IDENTITY is set as FULL and relation has a
+		 * generated column we should error out.
+		 */
+		if (relation->rd_rel->relreplident == REPLICA_IDENTITY_FULL)
+			return true;
+
+		/* Remember columns that are part of the REPLICA IDENTITY */
+		idattrs = RelationGetIndexAttrBitmap(relation,
+											 INDEX_ATTR_BITMAP_IDENTITY_KEY);
+
+		x = -1;
+		while ((x = bms_next_member(idattrs, x)) >= 0)
+		{
+			AttrNumber	attnum = (x + FirstLowInvalidHeapAttributeNumber);
+			Form_pg_attribute att = TupleDescAttr(desc, attnum - 1);
+
+			/* Check if generated column is part of REPLICA IDENTITY */
+			if (!att->attisdropped && att->attgenerated)
+			{
+				result = true;
+				break;
+			}
+		}
+
+		bms_free(idattrs);
+	}
+
+	return result;
+}
+
 /* check_functions_in_node callback */
 static bool
 contain_mutable_or_user_functions_checker(Oid func_id, void *context)
diff --git a/src/backend/executor/execReplication.c b/src/backend/executor/execReplication.c
index 54025c9f15..57599df153 100644
--- a/src/backend/executor/execReplication.c
+++ b/src/backend/executor/execReplication.c
@@ -809,6 +809,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot update table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_UPDATE && !pubdesc.replident_valid_for_update)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot update table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("Replica identity consists of an unpublished generated column.")));
 	else if (cmd == CMD_DELETE && !pubdesc.rf_valid_for_delete)
 		ereport(ERROR,
 				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
@@ -821,6 +827,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot delete from table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_DELETE && !pubdesc.replident_valid_for_delete)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot delete from table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("Replica identity consists of an unpublished generated column.")));
 
 	/* If relation has replica identity we are always good. */
 	if (OidIsValid(RelationGetReplicaIndex(rel)))
diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
index 342467fd18..be8f8eea8f 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -5736,6 +5736,8 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		pubdesc->rf_valid_for_delete = true;
 		pubdesc->cols_valid_for_update = true;
 		pubdesc->cols_valid_for_delete = true;
+		pubdesc->replident_valid_for_update = true;
+		pubdesc->replident_valid_for_delete = true;
 		return;
 	}
 
@@ -5750,6 +5752,8 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 	pubdesc->rf_valid_for_delete = true;
 	pubdesc->cols_valid_for_update = true;
 	pubdesc->cols_valid_for_delete = true;
+	pubdesc->replident_valid_for_update = true;
+	pubdesc->replident_valid_for_delete = true;
 
 	/* Fetch the publication membership info. */
 	puboids = GetRelationPublications(relid);
@@ -5827,6 +5831,21 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 				pubdesc->cols_valid_for_delete = false;
 		}
 
+		/*
+		 * Check if all generated columns included in the REPLICA IDENTITY are
+		 * published.
+		 */
+		if (!pubform->pubgencols &&
+			(pubform->pubupdate || pubform->pubdelete) &&
+			replident_has_unpublished_gen_col(pubid, relation, ancestors,
+											  pubform->pubviaroot))
+		{
+			if (pubform->pubupdate)
+				pubdesc->replident_valid_for_update = false;
+			if (pubform->pubdelete)
+				pubdesc->replident_valid_for_delete = false;
+		}
+
 		ReleaseSysCache(tup);
 
 		/*
@@ -5848,6 +5867,17 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 			pubdesc->pubactions.pubdelete && pubdesc->pubactions.pubtruncate &&
 			!pubdesc->cols_valid_for_update && !pubdesc->cols_valid_for_delete)
 			break;
+
+		/*
+		 * If we know everything is replicated and replica identity has an
+		 * unpublished generated column, there is no point to check for other
+		 * publications.
+		 */
+		if (pubdesc->pubactions.pubinsert && pubdesc->pubactions.pubupdate &&
+			pubdesc->pubactions.pubdelete && pubdesc->pubactions.pubtruncate &&
+			!pubdesc->replident_valid_for_update &&
+			!pubdesc->replident_valid_for_delete)
+			break;
 	}
 
 	if (relation->rd_pubdesc)
diff --git a/src/include/catalog/pg_publication.h b/src/include/catalog/pg_publication.h
index 9a83a72d6b..486f609a9a 100644
--- a/src/include/catalog/pg_publication.h
+++ b/src/include/catalog/pg_publication.h
@@ -98,6 +98,13 @@ typedef struct PublicationDesc
 	 */
 	bool		cols_valid_for_update;
 	bool		cols_valid_for_delete;
+
+	/*
+	 * true if all generated columns which are part of replica identity are
+	 * published or the publication actions do not include UPDATE or DELETE.
+	 */
+	bool		replident_valid_for_update;
+	bool		replident_valid_for_delete;
 } PublicationDesc;
 
 typedef struct Publication
diff --git a/src/include/commands/publicationcmds.h b/src/include/commands/publicationcmds.h
index 5487c571f6..b18e576b77 100644
--- a/src/include/commands/publicationcmds.h
+++ b/src/include/commands/publicationcmds.h
@@ -35,5 +35,7 @@ extern bool pub_rf_contains_invalid_column(Oid pubid, Relation relation,
 										   List *ancestors, bool pubviaroot);
 extern bool pub_collist_contains_invalid_column(Oid pubid, Relation relation,
 												List *ancestors, bool pubviaroot);
+extern bool replident_has_unpublished_gen_col(Oid pubid, Relation relation,
+											  List *ancestors, bool pubviaroot);
 
 #endif							/* PUBLICATIONCMDS_H */
diff --git a/src/test/regress/expected/publication.out b/src/test/regress/expected/publication.out
index 5de2d64d01..caf6e31554 100644
--- a/src/test/regress/expected/publication.out
+++ b/src/test/regress/expected/publication.out
@@ -672,6 +672,31 @@ DROP TABLE rf_tbl_abcd_pk;
 DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  Replica identity consists of an unpublished generated column.
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  Replica identity consists of an unpublished generated column.
+DROP PUBLICATION pub_gencol;
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+DROP INDEX testpub_gencol_idx;
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/regress/sql/publication.sql b/src/test/regress/sql/publication.sql
index 48e68bcca2..3b1c661440 100644
--- a/src/test/regress/sql/publication.sql
+++ b/src/test/regress/sql/publication.sql
@@ -396,6 +396,32 @@ DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
 
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+
+DROP INDEX testpub_gencol_idx;
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
+
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/subscription/t/100_bugs.pl b/src/test/subscription/t/100_bugs.pl
index cb36ca7b16..794b928f50 100644
--- a/src/test/subscription/t/100_bugs.pl
+++ b/src/test/subscription/t/100_bugs.pl
@@ -377,8 +377,8 @@ $node_publisher->safe_psql('postgres', "DROP PUBLICATION tap_pub_sch");
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');
 
-# The bug was that when the REPLICA IDENTITY FULL is used with dropped or
-# generated columns, we fail to apply updates and deletes
+# The bug was that when the REPLICA IDENTITY FULL is used with dropped
+# we fail to apply updates and deletes
 $node_publisher->rotate_logfile();
 $node_publisher->start();
 
@@ -389,18 +389,14 @@ $node_publisher->safe_psql(
 	'postgres', qq(
 	CREATE TABLE dropped_cols (a int, b_drop int, c int);
 	ALTER TABLE dropped_cols REPLICA IDENTITY FULL;
-	CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
-	ALTER TABLE generated_cols REPLICA IDENTITY FULL;
-	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols, generated_cols;
+	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols;
 	-- some initial data
 	INSERT INTO dropped_cols VALUES (1, 1, 1);
-	INSERT INTO generated_cols (a, c) VALUES (1, 1);
 ));
 
 $node_subscriber->safe_psql(
 	'postgres', qq(
 	 CREATE TABLE dropped_cols (a int, b_drop int, c int);
-	 CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
 ));
 
 $publisher_connstr = $node_publisher->connstr . ' dbname=postgres';
@@ -421,7 +417,6 @@ $node_subscriber->safe_psql(
 $node_publisher->safe_psql(
 	'postgres', qq(
 		UPDATE dropped_cols SET a = 100;
-		UPDATE generated_cols SET a = 100;
 ));
 $node_publisher->wait_for_catchup('sub_dropped_cols');
 
@@ -430,11 +425,6 @@ is( $node_subscriber->safe_psql(
 	qq(1),
 	'replication with RI FULL and dropped columns');
 
-is( $node_subscriber->safe_psql(
-		'postgres', "SELECT count(*) FROM generated_cols WHERE a = 100"),
-	qq(1),
-	'replication with RI FULL and generated columns');
-
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');
 
-- 
2.34.1

From 8640a4605f32f8072d093902da0af23dfa6d119b Mon Sep 17 00:00:00 2001
From: Hou Zhijie <houzj.fnst@cn.fujitsu.com>
Date: Wed, 20 Nov 2024 16:34:04 +0800
Subject: [PATCH v10] combine functions

---
 src/backend/commands/publicationcmds.c | 219 ++++++++++---------------
 src/backend/utils/cache/relcache.c     |  38 ++---
 src/include/commands/publicationcmds.h |   7 +-
 3 files changed, 105 insertions(+), 159 deletions(-)

diff --git a/src/backend/commands/publicationcmds.c b/src/backend/commands/publicationcmds.c
index 053877c524..0d5daf7626 100644
--- a/src/backend/commands/publicationcmds.c
+++ b/src/backend/commands/publicationcmds.c
@@ -336,21 +336,36 @@ pub_rf_contains_invalid_column(Oid pubid, Relation relation, List *ancestors,
 }
 
 /*
- * Check if all columns referenced in the REPLICA IDENTITY are covered by
- * the column list.
+ * Check for invalid columns in the publication table definition.
  *
- * Returns true if any replica identity column is not covered by column list.
+ * This function evaluates two conditions:
+ *
+ * 1. Ensures that all columns referenced in the REPLICA IDENTITY are covered
+ *    by the column list. If any column is missing, *invalid_column_list is set
+ *    to true.
+ *
+ * 2. Ensures that the REPLICA IDENTITY does not contain unpublished generated
+ *    columns. If an unpublished generated column is found,
+ *    *unpublished_gen_col is set to true.
+ *
+ * Returns true if any of the above conditions are not met.
  */
 bool
-pub_collist_contains_invalid_column(Oid pubid, Relation relation, List *ancestors,
-									bool pubviaroot)
+pub_contains_invalid_column(Oid pubid, Relation relation, List *ancestors,
+							bool pubviaroot, bool pubgencols,
+							bool *invalid_column_list,
+							bool *unpublished_gen_col)
 {
-	HeapTuple	tuple;
 	Oid			relid = RelationGetRelid(relation);
 	Oid			publish_as_relid = RelationGetRelid(relation);
-	bool		result = false;
-	Datum		datum;
-	bool		isnull;
+	Bitmapset  *idattrs;
+	Bitmapset  *columns = NULL;
+	TupleDesc	desc = RelationGetDescr(relation);
+	Publication *pub;
+	int			x;
+
+	*invalid_column_list = false;
+	*unpublished_gen_col = false;
 
 	/*
 	 * For a partition, if pubviaroot is true, find the topmost ancestor that
@@ -368,158 +383,90 @@ pub_collist_contains_invalid_column(Oid pubid, Relation relation, List *ancestor
 			publish_as_relid = relid;
 	}
 
-	tuple = SearchSysCache2(PUBLICATIONRELMAP,
-							ObjectIdGetDatum(publish_as_relid),
-							ObjectIdGetDatum(pubid));
-
-	if (!HeapTupleIsValid(tuple))
-		return false;
-
-	datum = SysCacheGetAttr(PUBLICATIONRELMAP, tuple,
-							Anum_pg_publication_rel_prattrs,
-							&isnull);
+	/* Fetch the column list */
+	pub = GetPublication(pubid);
+	check_and_fetch_column_list(pub, publish_as_relid, NULL, &columns);
 
-	if (!isnull)
+	if (relation->rd_rel->relreplident == REPLICA_IDENTITY_FULL)
 	{
-		int			x;
-		Bitmapset  *idattrs;
-		Bitmapset  *columns = NULL;
-
 		/* With REPLICA IDENTITY FULL, no column list is allowed. */
-		if (relation->rd_rel->relreplident == REPLICA_IDENTITY_FULL)
-			result = true;
-
-		/* Transform the column list datum to a bitmapset. */
-		columns = pub_collist_to_bitmapset(NULL, datum, NULL);
-
-		/* Remember columns that are part of the REPLICA IDENTITY */
-		idattrs = RelationGetIndexAttrBitmap(relation,
-											 INDEX_ATTR_BITMAP_IDENTITY_KEY);
+		*invalid_column_list = (columns != NULL);
 
 		/*
-		 * Attnums in the bitmap returned by RelationGetIndexAttrBitmap are
-		 * offset (to handle system columns the usual way), while column list
-		 * does not use offset, so we can't do bms_is_subset(). Instead, we
-		 * have to loop over the idattrs and check all of them are in the
-		 * list.
+		 * REPLICA IDENTITY can be FULL only if there is no column list for
+		 * publication. When REPLICA IDENTITY is FULL and the relation
+		 * includes a generated column, but the publish_generated_columns
+		 * option is set to false, this scenario is invalid.
 		 */
-		x = -1;
-		while ((x = bms_next_member(idattrs, x)) >= 0)
-		{
-			AttrNumber	attnum = (x + FirstLowInvalidHeapAttributeNumber);
-
-			/*
-			 * If pubviaroot is true, we are validating the column list of the
-			 * parent table, but the bitmap contains the replica identity
-			 * information of the child table. The parent/child attnums may
-			 * not match, so translate them to the parent - get the attname
-			 * from the child, and look it up in the parent.
-			 */
-			if (pubviaroot)
-			{
-				/* attribute name in the child table */
-				char	   *colname = get_attname(relid, attnum, false);
-
-				/*
-				 * Determine the attnum for the attribute name in parent (we
-				 * are using the column list defined on the parent).
-				 */
-				attnum = get_attnum(publish_as_relid, colname);
-			}
-
-			/* replica identity column, not covered by the column list */
-			if (!bms_is_member(attnum, columns))
-			{
-				result = true;
-				break;
-			}
-		}
+		if (!pubgencols && relation->rd_att->constr &&
+			relation->rd_att->constr->has_generated_stored)
+			*unpublished_gen_col = true;
 
-		bms_free(idattrs);
-		bms_free(columns);
+		if (*unpublished_gen_col && *invalid_column_list)
+			return true;
 	}
 
-	ReleaseSysCache(tuple);
-
-	return result;
-}
-
-/*
- * Check if REPLICA IDENTITY consists of any unpublished generated column.
- *
- * Returns true if any replica identity column is an unpublished generated
- * column.
- */
-bool
-replident_has_unpublished_gen_col(Oid pubid, Relation relation, List *ancestors,
-								  bool pubviaroot)
-{
-	Oid			relid = RelationGetRelid(relation);
-	Oid			publish_as_relid = RelationGetRelid(relation);
-	bool		result = false;
-	bool		found;
-	Publication *pub;
-
-	/* Return if the table does not contain any generated columns */
-	if (!relation->rd_att->constr ||
-		!relation->rd_att->constr->has_generated_stored)
-		return false;
+	/* Remember columns that are part of the REPLICA IDENTITY */
+	idattrs = RelationGetIndexAttrBitmap(relation,
+										 INDEX_ATTR_BITMAP_IDENTITY_KEY);
 
 	/*
-	 * For a partition, if pubviaroot is true, find the topmost ancestor that
-	 * is published via this publication as we need to use its column list for
-	 * the changes.
-	 *
-	 * Note that even though the column list used is for an ancestor, the
-	 * REPLICA IDENTITY used will be for the actual child table.
+	 * Attnums in the bitmap returned by RelationGetIndexAttrBitmap are offset
+	 * (to handle system columns the usual way), while column list does not
+	 * use offset, so we can't do bms_is_subset(). Instead, we have to loop
+	 * over the idattrs and check all of them are in the list.
 	 */
-	if (pubviaroot && relation->rd_rel->relispartition)
+	x = -1;
+	while ((x = bms_next_member(idattrs, x)) >= 0)
 	{
-		publish_as_relid = GetTopMostAncestorInPublication(pubid, ancestors, NULL);
+		AttrNumber	attnum = (x + FirstLowInvalidHeapAttributeNumber);
+		Form_pg_attribute att = TupleDescAttr(desc, attnum - 1);
 
-		if (!OidIsValid(publish_as_relid))
-			publish_as_relid = relid;
-	}
+		Assert(!att->attisdropped);
 
-	pub = GetPublication(pubid);
-	found = check_and_fetch_column_list(pub, publish_as_relid, NULL, NULL);
+		/* Check if generated column is part of REPLICA IDENTITY */
+		*unpublished_gen_col |= att->attgenerated;
 
-	if (!found)
-	{
-		TupleDesc	desc = RelationGetDescr(relation);
-		Bitmapset  *idattrs;
-		int			x;
+		if (columns == NULL)
+		{
+			/* Break the loop if unpublished generated columns exist. */
+			if (*unpublished_gen_col)
+				break;
+
+			/* Skip validating the column list since it is not defined */
+			continue;
+		}
 
 		/*
-		 * REPLICA IDENTITY can be FULL only if there is no column list for
-		 * publication. If REPLICA IDENTITY is set as FULL and relation has a
-		 * generated column we should error out.
+		 * If pubviaroot is true, we are validating the column list of the
+		 * parent table, but the bitmap contains the replica identity
+		 * information of the child table. The parent/child attnums may not
+		 * match, so translate them to the parent - get the attname from the
+		 * child, and look it up in the parent.
 		 */
-		if (relation->rd_rel->relreplident == REPLICA_IDENTITY_FULL)
-			return true;
-
-		/* Remember columns that are part of the REPLICA IDENTITY */
-		idattrs = RelationGetIndexAttrBitmap(relation,
-											 INDEX_ATTR_BITMAP_IDENTITY_KEY);
-
-		x = -1;
-		while ((x = bms_next_member(idattrs, x)) >= 0)
+		if (pubviaroot)
 		{
-			AttrNumber	attnum = (x + FirstLowInvalidHeapAttributeNumber);
-			Form_pg_attribute att = TupleDescAttr(desc, attnum - 1);
+			/* attribute name in the child table */
+			char	   *colname = get_attname(relid, attnum, false);
 
-			/* Check if generated column is part of REPLICA IDENTITY */
-			if (!att->attisdropped && att->attgenerated)
-			{
-				result = true;
-				break;
-			}
+			/*
+			 * Determine the attnum for the attribute name in parent (we are
+			 * using the column list defined on the parent).
+			 */
+			attnum = get_attnum(publish_as_relid, colname);
 		}
 
-		bms_free(idattrs);
+		/* replica identity column, not covered by the column list */
+		*invalid_column_list |= !bms_is_member(attnum, columns);
+
+		if (*invalid_column_list && *unpublished_gen_col)
+			break;
 	}
 
-	return result;
+	bms_free(columns);
+	bms_free(idattrs);
+
+	return *invalid_column_list || *unpublished_gen_col;
 }
 
 /* check_functions_in_node callback */
diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
index be8f8eea8f..c1199596bf 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -5783,6 +5783,8 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		Oid			pubid = lfirst_oid(lc);
 		HeapTuple	tup;
 		Form_pg_publication pubform;
+		bool		invalid_column_list;
+		bool		unpublished_gen_col;
 
 		tup = SearchSysCache1(PUBLICATIONOID, ObjectIdGetDatum(pubid));
 
@@ -5817,33 +5819,27 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		/*
 		 * Check if all columns are part of the REPLICA IDENTITY index or not.
 		 *
-		 * If the publication is FOR ALL TABLES then it means the table has no
-		 * column list and we can skip the validation.
-		 */
-		if (!pubform->puballtables &&
-			(pubform->pubupdate || pubform->pubdelete) &&
-			pub_collist_contains_invalid_column(pubid, relation, ancestors,
-												pubform->pubviaroot))
-		{
-			if (pubform->pubupdate)
-				pubdesc->cols_valid_for_update = false;
-			if (pubform->pubdelete)
-				pubdesc->cols_valid_for_delete = false;
-		}
-
-		/*
 		 * Check if all generated columns included in the REPLICA IDENTITY are
 		 * published.
 		 */
-		if (!pubform->pubgencols &&
-			(pubform->pubupdate || pubform->pubdelete) &&
-			replident_has_unpublished_gen_col(pubid, relation, ancestors,
-											  pubform->pubviaroot))
+		if ((pubform->pubupdate || pubform->pubdelete) &&
+			pub_contains_invalid_column(pubid, relation, ancestors,
+										pubform->pubviaroot,
+										pubform->pubgencols,
+										&invalid_column_list,
+										&unpublished_gen_col))
 		{
 			if (pubform->pubupdate)
-				pubdesc->replident_valid_for_update = false;
+			{
+				pubdesc->cols_valid_for_update = !invalid_column_list;
+				pubdesc->replident_valid_for_update = !unpublished_gen_col;
+			}
+
 			if (pubform->pubdelete)
-				pubdesc->replident_valid_for_delete = false;
+			{
+				pubdesc->cols_valid_for_delete = !invalid_column_list;
+				pubdesc->replident_valid_for_delete = !unpublished_gen_col;
+			}
 		}
 
 		ReleaseSysCache(tup);
diff --git a/src/include/commands/publicationcmds.h b/src/include/commands/publicationcmds.h
index b18e576b77..fbf78ea08e 100644
--- a/src/include/commands/publicationcmds.h
+++ b/src/include/commands/publicationcmds.h
@@ -33,8 +33,11 @@ extern void AlterPublicationOwner_oid(Oid subid, Oid newOwnerId);
 extern void InvalidatePublicationRels(List *relids);
 extern bool pub_rf_contains_invalid_column(Oid pubid, Relation relation,
 										   List *ancestors, bool pubviaroot);
-extern bool pub_collist_contains_invalid_column(Oid pubid, Relation relation,
-												List *ancestors, bool pubviaroot);
+extern bool pub_contains_invalid_column(Oid pubid, Relation relation,
+										List *ancestors, bool pubviaroot,
+										bool pubgencols,
+										bool *invalid_column_list,
+										bool *unpublished_gen_col);
 extern bool replident_has_unpublished_gen_col(Oid pubid, Relation relation,
 											  List *ancestors, bool pubviaroot);
 
-- 
2.30.0.windows.2

From 0e9b329ddb84ebe3f622b3cfdca7800516e68601 Mon Sep 17 00:00:00 2001
From: Shlok Kyal <shlok.kyal.oss@gmail.com>
Date: Tue, 5 Nov 2024 11:11:25 +0530
Subject: [PATCH v11] Disallow UPDATE/DELETE on table with generated column as
 REPLICA IDENTITY

UPDATE or DELETE operations on tables with unpublished generated columns
set as REPLICA IDENTITY are not permitted. This patch ensures that if an
UPDATE or DELETE command is executed on such tables, an error will be
thrown.
With this patch, the behavior has changed for the test added in commit
adedf54. Additionally, there is no other way to trigger the bug that was
fixed by commit adedf54, so the test has been removed.
---
 src/backend/commands/publicationcmds.c    | 78 +++++++++++++++++++++++
 src/backend/executor/execReplication.c    | 12 ++++
 src/backend/utils/cache/relcache.c        | 30 +++++++++
 src/include/catalog/pg_publication.h      |  7 ++
 src/include/commands/publicationcmds.h    |  2 +
 src/test/regress/expected/publication.out | 24 +++++++
 src/test/regress/sql/publication.sql      | 25 ++++++++
 src/test/subscription/t/100_bugs.pl       | 16 +----
 8 files changed, 181 insertions(+), 13 deletions(-)

diff --git a/src/backend/commands/publicationcmds.c b/src/backend/commands/publicationcmds.c
index 0129db18c6..bad5fbba5c 100644
--- a/src/backend/commands/publicationcmds.c
+++ b/src/backend/commands/publicationcmds.c
@@ -444,6 +444,84 @@ pub_collist_contains_invalid_column(Oid pubid, Relation relation, List *ancestor
 	return result;
 }
 
+/*
+ * Check if REPLICA IDENTITY consists of any unpublished generated column.
+ *
+ * Returns true if any replica identity column is an unpublished generated
+ * column.
+ */
+bool
+replident_has_unpublished_gen_col(Oid pubid, Relation relation, List *ancestors,
+								  bool pubviaroot)
+{
+	Oid			relid = RelationGetRelid(relation);
+	Oid			publish_as_relid = RelationGetRelid(relation);
+	bool		result = false;
+	bool		found;
+	Publication *pub;
+
+	/* Return if the table does not contain any generated columns */
+	if (!relation->rd_att->constr ||
+		!relation->rd_att->constr->has_generated_stored)
+		return false;
+
+	/*
+	 * For a partition, if pubviaroot is true, find the topmost ancestor that
+	 * is published via this publication as we need to check if its column list
+	 * is specified.
+	 *
+	 * Note that even though the column list of an ancestor is checked, the
+	 * REPLICA IDENTITY used will be for the actual child table.
+	 */
+	if (pubviaroot && relation->rd_rel->relispartition)
+	{
+		publish_as_relid = GetTopMostAncestorInPublication(pubid, ancestors, NULL);
+
+		if (!OidIsValid(publish_as_relid))
+			publish_as_relid = relid;
+	}
+
+	pub = GetPublication(pubid);
+	found = check_and_fetch_column_list(pub, publish_as_relid, NULL, NULL);
+
+	if (!found)
+	{
+		TupleDesc	desc = RelationGetDescr(relation);
+		Bitmapset  *idattrs;
+		int			x;
+
+		/*
+		 * REPLICA IDENTITY can be FULL only if there is no column list for
+		 * publication. If REPLICA IDENTITY is set as FULL and relation has a
+		 * generated column we should error out.
+		 */
+		if (relation->rd_rel->relreplident == REPLICA_IDENTITY_FULL)
+			return true;
+
+		/* Remember columns that are part of the REPLICA IDENTITY */
+		idattrs = RelationGetIndexAttrBitmap(relation,
+											 INDEX_ATTR_BITMAP_IDENTITY_KEY);
+
+		x = -1;
+		while ((x = bms_next_member(idattrs, x)) >= 0)
+		{
+			AttrNumber	attnum = (x + FirstLowInvalidHeapAttributeNumber);
+			Form_pg_attribute att = TupleDescAttr(desc, attnum - 1);
+
+			/* Check if generated column is part of REPLICA IDENTITY */
+			if (!att->attisdropped && att->attgenerated)
+			{
+				result = true;
+				break;
+			}
+		}
+
+		bms_free(idattrs);
+	}
+
+	return result;
+}
+
 /* check_functions_in_node callback */
 static bool
 contain_mutable_or_user_functions_checker(Oid func_id, void *context)
diff --git a/src/backend/executor/execReplication.c b/src/backend/executor/execReplication.c
index 54025c9f15..57599df153 100644
--- a/src/backend/executor/execReplication.c
+++ b/src/backend/executor/execReplication.c
@@ -809,6 +809,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot update table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_UPDATE && !pubdesc.replident_valid_for_update)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot update table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("Replica identity consists of an unpublished generated column.")));
 	else if (cmd == CMD_DELETE && !pubdesc.rf_valid_for_delete)
 		ereport(ERROR,
 				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
@@ -821,6 +827,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot delete from table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_DELETE && !pubdesc.replident_valid_for_delete)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot delete from table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("Replica identity consists of an unpublished generated column.")));
 
 	/* If relation has replica identity we are always good. */
 	if (OidIsValid(RelationGetReplicaIndex(rel)))
diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
index 342467fd18..be8f8eea8f 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -5736,6 +5736,8 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		pubdesc->rf_valid_for_delete = true;
 		pubdesc->cols_valid_for_update = true;
 		pubdesc->cols_valid_for_delete = true;
+		pubdesc->replident_valid_for_update = true;
+		pubdesc->replident_valid_for_delete = true;
 		return;
 	}
 
@@ -5750,6 +5752,8 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 	pubdesc->rf_valid_for_delete = true;
 	pubdesc->cols_valid_for_update = true;
 	pubdesc->cols_valid_for_delete = true;
+	pubdesc->replident_valid_for_update = true;
+	pubdesc->replident_valid_for_delete = true;
 
 	/* Fetch the publication membership info. */
 	puboids = GetRelationPublications(relid);
@@ -5827,6 +5831,21 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 				pubdesc->cols_valid_for_delete = false;
 		}
 
+		/*
+		 * Check if all generated columns included in the REPLICA IDENTITY are
+		 * published.
+		 */
+		if (!pubform->pubgencols &&
+			(pubform->pubupdate || pubform->pubdelete) &&
+			replident_has_unpublished_gen_col(pubid, relation, ancestors,
+											  pubform->pubviaroot))
+		{
+			if (pubform->pubupdate)
+				pubdesc->replident_valid_for_update = false;
+			if (pubform->pubdelete)
+				pubdesc->replident_valid_for_delete = false;
+		}
+
 		ReleaseSysCache(tup);
 
 		/*
@@ -5848,6 +5867,17 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 			pubdesc->pubactions.pubdelete && pubdesc->pubactions.pubtruncate &&
 			!pubdesc->cols_valid_for_update && !pubdesc->cols_valid_for_delete)
 			break;
+
+		/*
+		 * If we know everything is replicated and replica identity has an
+		 * unpublished generated column, there is no point to check for other
+		 * publications.
+		 */
+		if (pubdesc->pubactions.pubinsert && pubdesc->pubactions.pubupdate &&
+			pubdesc->pubactions.pubdelete && pubdesc->pubactions.pubtruncate &&
+			!pubdesc->replident_valid_for_update &&
+			!pubdesc->replident_valid_for_delete)
+			break;
 	}
 
 	if (relation->rd_pubdesc)
diff --git a/src/include/catalog/pg_publication.h b/src/include/catalog/pg_publication.h
index 9a83a72d6b..486f609a9a 100644
--- a/src/include/catalog/pg_publication.h
+++ b/src/include/catalog/pg_publication.h
@@ -98,6 +98,13 @@ typedef struct PublicationDesc
 	 */
 	bool		cols_valid_for_update;
 	bool		cols_valid_for_delete;
+
+	/*
+	 * true if all generated columns which are part of replica identity are
+	 * published or the publication actions do not include UPDATE or DELETE.
+	 */
+	bool		replident_valid_for_update;
+	bool		replident_valid_for_delete;
 } PublicationDesc;
 
 typedef struct Publication
diff --git a/src/include/commands/publicationcmds.h b/src/include/commands/publicationcmds.h
index 5487c571f6..b18e576b77 100644
--- a/src/include/commands/publicationcmds.h
+++ b/src/include/commands/publicationcmds.h
@@ -35,5 +35,7 @@ extern bool pub_rf_contains_invalid_column(Oid pubid, Relation relation,
 										   List *ancestors, bool pubviaroot);
 extern bool pub_collist_contains_invalid_column(Oid pubid, Relation relation,
 												List *ancestors, bool pubviaroot);
+extern bool replident_has_unpublished_gen_col(Oid pubid, Relation relation,
+											  List *ancestors, bool pubviaroot);
 
 #endif							/* PUBLICATIONCMDS_H */
diff --git a/src/test/regress/expected/publication.out b/src/test/regress/expected/publication.out
index 5de2d64d01..12d0611a7c 100644
--- a/src/test/regress/expected/publication.out
+++ b/src/test/regress/expected/publication.out
@@ -672,6 +672,30 @@ DROP TABLE rf_tbl_abcd_pk;
 DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  Replica identity consists of an unpublished generated column.
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  Replica identity consists of an unpublished generated column.
+DROP PUBLICATION pub_gencol;
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/regress/sql/publication.sql b/src/test/regress/sql/publication.sql
index 48e68bcca2..a29587b45d 100644
--- a/src/test/regress/sql/publication.sql
+++ b/src/test/regress/sql/publication.sql
@@ -396,6 +396,31 @@ DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
 
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
+
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/subscription/t/100_bugs.pl b/src/test/subscription/t/100_bugs.pl
index cb36ca7b16..794b928f50 100644
--- a/src/test/subscription/t/100_bugs.pl
+++ b/src/test/subscription/t/100_bugs.pl
@@ -377,8 +377,8 @@ $node_publisher->safe_psql('postgres', "DROP PUBLICATION tap_pub_sch");
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');
 
-# The bug was that when the REPLICA IDENTITY FULL is used with dropped or
-# generated columns, we fail to apply updates and deletes
+# The bug was that when the REPLICA IDENTITY FULL is used with dropped
+# we fail to apply updates and deletes
 $node_publisher->rotate_logfile();
 $node_publisher->start();
 
@@ -389,18 +389,14 @@ $node_publisher->safe_psql(
 	'postgres', qq(
 	CREATE TABLE dropped_cols (a int, b_drop int, c int);
 	ALTER TABLE dropped_cols REPLICA IDENTITY FULL;
-	CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
-	ALTER TABLE generated_cols REPLICA IDENTITY FULL;
-	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols, generated_cols;
+	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols;
 	-- some initial data
 	INSERT INTO dropped_cols VALUES (1, 1, 1);
-	INSERT INTO generated_cols (a, c) VALUES (1, 1);
 ));
 
 $node_subscriber->safe_psql(
 	'postgres', qq(
 	 CREATE TABLE dropped_cols (a int, b_drop int, c int);
-	 CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
 ));
 
 $publisher_connstr = $node_publisher->connstr . ' dbname=postgres';
@@ -421,7 +417,6 @@ $node_subscriber->safe_psql(
 $node_publisher->safe_psql(
 	'postgres', qq(
 		UPDATE dropped_cols SET a = 100;
-		UPDATE generated_cols SET a = 100;
 ));
 $node_publisher->wait_for_catchup('sub_dropped_cols');
 
@@ -430,11 +425,6 @@ is( $node_subscriber->safe_psql(
 	qq(1),
 	'replication with RI FULL and dropped columns');
 
-is( $node_subscriber->safe_psql(
-		'postgres', "SELECT count(*) FROM generated_cols WHERE a = 100"),
-	qq(1),
-	'replication with RI FULL and generated columns');
-
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');
 
-- 
2.34.1

From 4dd53e651323c0574531a9b4d79ed31f6af1b2d0 Mon Sep 17 00:00:00 2001
From: Shlok Kyal <shlok.kyal.oss@gmail.com>
Date: Tue, 5 Nov 2024 11:11:25 +0530
Subject: [PATCH v12] Disallow UPDATE/DELETE on table with generated column as
 REPLICA IDENTITY

UPDATE or DELETE operations on tables with unpublished generated columns
set as REPLICA IDENTITY are not permitted. This patch ensures that if an
UPDATE or DELETE command is executed on such tables, an error will be
thrown.
With this patch, the behavior has changed for the test added in commit
adedf54. Additionally, there is no other way to trigger the bug that was
fixed by commit adedf54, so the test has been removed.
---
 src/backend/commands/publicationcmds.c    | 153 +++++++++++++---------
 src/backend/executor/execReplication.c    |  32 ++++-
 src/backend/utils/cache/relcache.c        |  57 ++++++--
 src/include/catalog/pg_publication.h      |   7 +
 src/include/commands/publicationcmds.h    |   7 +-
 src/test/regress/expected/publication.out |  24 ++++
 src/test/regress/sql/publication.sql      |  25 ++++
 src/test/subscription/t/100_bugs.pl       |  16 +--
 8 files changed, 226 insertions(+), 95 deletions(-)

diff --git a/src/backend/commands/publicationcmds.c b/src/backend/commands/publicationcmds.c
index 031c84ec29..d153ee329a 100644
--- a/src/backend/commands/publicationcmds.c
+++ b/src/backend/commands/publicationcmds.c
@@ -336,21 +336,36 @@ pub_rf_contains_invalid_column(Oid pubid, Relation relation, List *ancestors,
 }
 
 /*
- * Check if all columns referenced in the REPLICA IDENTITY are covered by
- * the column list.
+ * Check for invalid columns in the publication table definition.
  *
- * Returns true if any replica identity column is not covered by column list.
+ * This function evaluates two conditions:
+ *
+ * 1. Ensures that all columns referenced in the REPLICA IDENTITY are covered
+ *    by the column list. If any column is missing, *invalid_column_list is set
+ *    to true.
+ *
+ * 2. Ensures that the REPLICA IDENTITY does not contain unpublished generated
+ *    columns. If an unpublished generated column is found,
+ *    *unpublished_gen_col is set to true.
+ *
+ * Returns true if any of the above conditions are not met.
  */
 bool
-pub_collist_contains_invalid_column(Oid pubid, Relation relation, List *ancestors,
-									bool pubviaroot)
+pub_contains_invalid_column(Oid pubid, Relation relation, List *ancestors,
+							bool pubviaroot, bool pubgencols,
+							bool *invalid_column_list,
+							bool *unpublished_gen_col)
 {
-	HeapTuple	tuple;
 	Oid			relid = RelationGetRelid(relation);
 	Oid			publish_as_relid = RelationGetRelid(relation);
-	bool		result = false;
-	Datum		datum;
-	bool		isnull;
+	Bitmapset  *idattrs;
+	Bitmapset  *columns = NULL;
+	TupleDesc	desc = RelationGetDescr(relation);
+	Publication *pub;
+	int			x;
+
+	*invalid_column_list = false;
+	*unpublished_gen_col = false;
 
 	/*
 	 * For a partition, if pubviaroot is true, find the topmost ancestor that
@@ -368,80 +383,90 @@ pub_collist_contains_invalid_column(Oid pubid, Relation relation, List *ancestor
 			publish_as_relid = relid;
 	}
 
-	tuple = SearchSysCache2(PUBLICATIONRELMAP,
-							ObjectIdGetDatum(publish_as_relid),
-							ObjectIdGetDatum(pubid));
+	/* Fetch the column list */
+	pub = GetPublication(pubid);
+	check_and_fetch_column_list(pub, publish_as_relid, NULL, &columns);
 
-	if (!HeapTupleIsValid(tuple))
-		return false;
+	if (relation->rd_rel->relreplident == REPLICA_IDENTITY_FULL)
+	{
+		/* With REPLICA IDENTITY FULL, no column list is allowed. */
+		*invalid_column_list = (columns != NULL);
 
-	datum = SysCacheGetAttr(PUBLICATIONRELMAP, tuple,
-							Anum_pg_publication_rel_prattrs,
-							&isnull);
+		/*
+		 * REPLICA IDENTITY can be FULL only if there is no column list for
+		 * publication. When REPLICA IDENTITY is FULL and the relation
+		 * includes a generated column, but the publish_generated_columns
+		 * option is set to false, this scenario is invalid.
+		 */
+		if (!pubgencols && relation->rd_att->constr &&
+			relation->rd_att->constr->has_generated_stored)
+			*unpublished_gen_col = true;
 
-	if (!isnull)
+		if (*unpublished_gen_col && *invalid_column_list)
+			return true;
+	}
+
+	/* Remember columns that are part of the REPLICA IDENTITY */
+	idattrs = RelationGetIndexAttrBitmap(relation,
+										 INDEX_ATTR_BITMAP_IDENTITY_KEY);
+
+	/*
+	 * Attnums in the bitmap returned by RelationGetIndexAttrBitmap are offset
+	 * (to handle system columns the usual way), while column list does not
+	 * use offset, so we can't do bms_is_subset(). Instead, we have to loop
+	 * over the idattrs and check all of them are in the list.
+	 */
+	x = -1;
+	while ((x = bms_next_member(idattrs, x)) >= 0)
 	{
-		int			x;
-		Bitmapset  *idattrs;
-		Bitmapset  *columns = NULL;
+		AttrNumber	attnum = (x + FirstLowInvalidHeapAttributeNumber);
+		Form_pg_attribute att = TupleDescAttr(desc, attnum - 1);
 
-		/* With REPLICA IDENTITY FULL, no column list is allowed. */
-		if (relation->rd_rel->relreplident == REPLICA_IDENTITY_FULL)
-			result = true;
+		Assert(!att->attisdropped);
 
-		/* Transform the column list datum to a bitmapset. */
-		columns = pub_collist_to_bitmapset(NULL, datum, NULL);
+		/* Check if generated column is part of REPLICA IDENTITY */
+		*unpublished_gen_col |= att->attgenerated;
 
-		/* Remember columns that are part of the REPLICA IDENTITY */
-		idattrs = RelationGetIndexAttrBitmap(relation,
-											 INDEX_ATTR_BITMAP_IDENTITY_KEY);
+		if (columns == NULL)
+		{
+			/* Break the loop if unpublished generated columns exist. */
+			if (*unpublished_gen_col)
+				break;
+
+			/* Skip validating the column list since it is not defined */
+			continue;
+		}
 
 		/*
-		 * Attnums in the bitmap returned by RelationGetIndexAttrBitmap are
-		 * offset (to handle system columns the usual way), while column list
-		 * does not use offset, so we can't do bms_is_subset(). Instead, we
-		 * have to loop over the idattrs and check all of them are in the
-		 * list.
+		 * If pubviaroot is true, we are validating the column list of the
+		 * parent table, but the bitmap contains the replica identity
+		 * information of the child table. The parent/child attnums may not
+		 * match, so translate them to the parent - get the attname from the
+		 * child, and look it up in the parent.
 		 */
-		x = -1;
-		while ((x = bms_next_member(idattrs, x)) >= 0)
+		if (pubviaroot)
 		{
-			AttrNumber	attnum = (x + FirstLowInvalidHeapAttributeNumber);
+			/* attribute name in the child table */
+			char	   *colname = get_attname(relid, attnum, false);
 
 			/*
-			 * If pubviaroot is true, we are validating the column list of the
-			 * parent table, but the bitmap contains the replica identity
-			 * information of the child table. The parent/child attnums may
-			 * not match, so translate them to the parent - get the attname
-			 * from the child, and look it up in the parent.
+			 * Determine the attnum for the attribute name in parent (we are
+			 * using the column list defined on the parent).
 			 */
-			if (pubviaroot)
-			{
-				/* attribute name in the child table */
-				char	   *colname = get_attname(relid, attnum, false);
-
-				/*
-				 * Determine the attnum for the attribute name in parent (we
-				 * are using the column list defined on the parent).
-				 */
-				attnum = get_attnum(publish_as_relid, colname);
-			}
-
-			/* replica identity column, not covered by the column list */
-			if (!bms_is_member(attnum, columns))
-			{
-				result = true;
-				break;
-			}
+			attnum = get_attnum(publish_as_relid, colname);
 		}
 
-		bms_free(idattrs);
-		bms_free(columns);
+		/* replica identity column, not covered by the column list */
+		*invalid_column_list |= !bms_is_member(attnum, columns);
+
+		if (*invalid_column_list && *unpublished_gen_col)
+			break;
 	}
 
-	ReleaseSysCache(tuple);
+	bms_free(columns);
+	bms_free(idattrs);
 
-	return result;
+	return *invalid_column_list || *unpublished_gen_col;
 }
 
 /* check_functions_in_node callback */
diff --git a/src/backend/executor/execReplication.c b/src/backend/executor/execReplication.c
index 54025c9f15..227a8aeea0 100644
--- a/src/backend/executor/execReplication.c
+++ b/src/backend/executor/execReplication.c
@@ -785,10 +785,22 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 		return;
 
 	/*
-	 * 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 - i.e. when all referenced columns are part of REPLICA IDENTITY
-	 * or the table does not publish UPDATEs or DELETEs.
+	 * It is only safe to execute UPDATE/DELETE when:
+	 *
+	 * 1. All columns, referenced in the row filters from publications which
+	 * the relation is in, are valid - i.e. when all referenced columns are
+	 * part of REPLICA IDENTITY or the table does not publish UPDATEs or
+	 * DELETEs.
+	 *
+	 * 2. All columns, referenced in the column lists from publications which
+	 * the relation is in, are valid - i.e. when all referenced columns are
+	 * part of REPLICA IDENTITY or the table does not publish UPDATEs or
+	 * DELETEs.
+	 *
+	 * 3. All generated columns in REPLICA IDENTITY of the relation, for all
+	 * the publications which the relation is in, are valid - i.e. when
+	 * unpublished generated columns are not part of REPLICA IDENTITY or the
+	 * table does not publish UPDATEs or DELETEs.
 	 *
 	 * XXX We could optimize it by first checking whether any of the
 	 * publications have a row filter for this relation. If not and relation
@@ -809,6 +821,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot update table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_UPDATE && !pubdesc.gencols_valid_for_update)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot update table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("Replica identity consists of an unpublished generated column.")));
 	else if (cmd == CMD_DELETE && !pubdesc.rf_valid_for_delete)
 		ereport(ERROR,
 				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
@@ -821,6 +839,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot delete from table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_DELETE && !pubdesc.gencols_valid_for_delete)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot delete from table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("Replica identity consists of an unpublished generated column.")));
 
 	/* If relation has replica identity we are always good. */
 	if (OidIsValid(RelationGetReplicaIndex(rel)))
diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
index d0892cee24..9205ea86c4 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -5706,12 +5706,19 @@ RelationGetExclusionInfo(Relation indexRelation,
  * Get the publication information for the given relation.
  *
  * Traverse all the publications which the relation is in to get the
- * publication actions and validate the row filter expressions for such
- * publications if any. We consider the row filter expression as invalid if it
- * references any column which is not part of REPLICA IDENTITY.
+ * publication actions and validate:
+ * 1. The row filter expressions for such publications if any. We consider the
+ *    row filter expression as invalid if it references any column which is not
+ *    part of REPLICA IDENTITY.
+ * 2. The column list for such publication if any. We consider column list as
+ * 	  invalid if it references any column which is not part of REPLICA IDENTITY.
+ * 3. The generated columns of the relation for such publications. We consider
+ *    any reference of an unpublished generated column in REPLICA IDENTITY as
+ *    invalid.
  *
  * To avoid fetching the publication information repeatedly, we cache the
- * publication actions and row filter validation information.
+ * publication actions, row filter validation information, column list
+ * validation information and generated column validation information.
  */
 void
 RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
@@ -5734,6 +5741,8 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		pubdesc->rf_valid_for_delete = true;
 		pubdesc->cols_valid_for_update = true;
 		pubdesc->cols_valid_for_delete = true;
+		pubdesc->gencols_valid_for_update = true;
+		pubdesc->gencols_valid_for_delete = true;
 		return;
 	}
 
@@ -5748,6 +5757,8 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 	pubdesc->rf_valid_for_delete = true;
 	pubdesc->cols_valid_for_update = true;
 	pubdesc->cols_valid_for_delete = true;
+	pubdesc->gencols_valid_for_update = true;
+	pubdesc->gencols_valid_for_delete = true;
 
 	/* Fetch the publication membership info. */
 	puboids = GetRelationPublications(relid);
@@ -5777,6 +5788,8 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		Oid			pubid = lfirst_oid(lc);
 		HeapTuple	tup;
 		Form_pg_publication pubform;
+		bool		invalid_column_list;
+		bool		unpublished_gen_col;
 
 		tup = SearchSysCache1(PUBLICATIONOID, ObjectIdGetDatum(pubid));
 
@@ -5811,18 +5824,27 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		/*
 		 * Check if all columns are part of the REPLICA IDENTITY index or not.
 		 *
-		 * If the publication is FOR ALL TABLES then it means the table has no
-		 * column list and we can skip the validation.
+		 * Check if all generated columns included in the REPLICA IDENTITY are
+		 * published.
 		 */
-		if (!pubform->puballtables &&
-			(pubform->pubupdate || pubform->pubdelete) &&
-			pub_collist_contains_invalid_column(pubid, relation, ancestors,
-												pubform->pubviaroot))
+		if ((pubform->pubupdate || pubform->pubdelete) &&
+			pub_contains_invalid_column(pubid, relation, ancestors,
+										pubform->pubviaroot,
+										pubform->pubgencols,
+										&invalid_column_list,
+										&unpublished_gen_col))
 		{
 			if (pubform->pubupdate)
-				pubdesc->cols_valid_for_update = false;
+			{
+				pubdesc->cols_valid_for_update = !invalid_column_list;
+				pubdesc->gencols_valid_for_update = !unpublished_gen_col;
+			}
+
 			if (pubform->pubdelete)
-				pubdesc->cols_valid_for_delete = false;
+			{
+				pubdesc->cols_valid_for_delete = !invalid_column_list;
+				pubdesc->gencols_valid_for_delete = !unpublished_gen_col;
+			}
 		}
 
 		ReleaseSysCache(tup);
@@ -5846,6 +5868,17 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 			pubdesc->pubactions.pubdelete && pubdesc->pubactions.pubtruncate &&
 			!pubdesc->cols_valid_for_update && !pubdesc->cols_valid_for_delete)
 			break;
+
+		/*
+		 * If we know everything is replicated and replica identity has an
+		 * unpublished generated column, there is no point to check for other
+		 * publications.
+		 */
+		if (pubdesc->pubactions.pubinsert && pubdesc->pubactions.pubupdate &&
+			pubdesc->pubactions.pubdelete && pubdesc->pubactions.pubtruncate &&
+			!pubdesc->gencols_valid_for_update &&
+			!pubdesc->gencols_valid_for_delete)
+			break;
 	}
 
 	if (relation->rd_pubdesc)
diff --git a/src/include/catalog/pg_publication.h b/src/include/catalog/pg_publication.h
index 9a83a72d6b..e167b34461 100644
--- a/src/include/catalog/pg_publication.h
+++ b/src/include/catalog/pg_publication.h
@@ -98,6 +98,13 @@ typedef struct PublicationDesc
 	 */
 	bool		cols_valid_for_update;
 	bool		cols_valid_for_delete;
+
+	/*
+	 * true if all generated columns which are part of replica identity are
+	 * published or the publication actions do not include UPDATE or DELETE.
+	 */
+	bool		gencols_valid_for_update;
+	bool		gencols_valid_for_delete;
 } PublicationDesc;
 
 typedef struct Publication
diff --git a/src/include/commands/publicationcmds.h b/src/include/commands/publicationcmds.h
index 5487c571f6..c200c19c6d 100644
--- a/src/include/commands/publicationcmds.h
+++ b/src/include/commands/publicationcmds.h
@@ -33,7 +33,10 @@ extern void AlterPublicationOwner_oid(Oid subid, Oid newOwnerId);
 extern void InvalidatePublicationRels(List *relids);
 extern bool pub_rf_contains_invalid_column(Oid pubid, Relation relation,
 										   List *ancestors, bool pubviaroot);
-extern bool pub_collist_contains_invalid_column(Oid pubid, Relation relation,
-												List *ancestors, bool pubviaroot);
+extern bool pub_contains_invalid_column(Oid pubid, Relation relation,
+										List *ancestors, bool pubviaroot,
+										bool pubgencols,
+										bool *invalid_column_list,
+										bool *unpublished_gen_col);
 
 #endif							/* PUBLICATIONCMDS_H */
diff --git a/src/test/regress/expected/publication.out b/src/test/regress/expected/publication.out
index 5de2d64d01..12d0611a7c 100644
--- a/src/test/regress/expected/publication.out
+++ b/src/test/regress/expected/publication.out
@@ -672,6 +672,30 @@ DROP TABLE rf_tbl_abcd_pk;
 DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  Replica identity consists of an unpublished generated column.
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  Replica identity consists of an unpublished generated column.
+DROP PUBLICATION pub_gencol;
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/regress/sql/publication.sql b/src/test/regress/sql/publication.sql
index 48e68bcca2..a29587b45d 100644
--- a/src/test/regress/sql/publication.sql
+++ b/src/test/regress/sql/publication.sql
@@ -396,6 +396,31 @@ DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
 
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
+
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/subscription/t/100_bugs.pl b/src/test/subscription/t/100_bugs.pl
index cb36ca7b16..794b928f50 100644
--- a/src/test/subscription/t/100_bugs.pl
+++ b/src/test/subscription/t/100_bugs.pl
@@ -377,8 +377,8 @@ $node_publisher->safe_psql('postgres', "DROP PUBLICATION tap_pub_sch");
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');
 
-# The bug was that when the REPLICA IDENTITY FULL is used with dropped or
-# generated columns, we fail to apply updates and deletes
+# The bug was that when the REPLICA IDENTITY FULL is used with dropped
+# we fail to apply updates and deletes
 $node_publisher->rotate_logfile();
 $node_publisher->start();
 
@@ -389,18 +389,14 @@ $node_publisher->safe_psql(
 	'postgres', qq(
 	CREATE TABLE dropped_cols (a int, b_drop int, c int);
 	ALTER TABLE dropped_cols REPLICA IDENTITY FULL;
-	CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
-	ALTER TABLE generated_cols REPLICA IDENTITY FULL;
-	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols, generated_cols;
+	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols;
 	-- some initial data
 	INSERT INTO dropped_cols VALUES (1, 1, 1);
-	INSERT INTO generated_cols (a, c) VALUES (1, 1);
 ));
 
 $node_subscriber->safe_psql(
 	'postgres', qq(
 	 CREATE TABLE dropped_cols (a int, b_drop int, c int);
-	 CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
 ));
 
 $publisher_connstr = $node_publisher->connstr . ' dbname=postgres';
@@ -421,7 +417,6 @@ $node_subscriber->safe_psql(
 $node_publisher->safe_psql(
 	'postgres', qq(
 		UPDATE dropped_cols SET a = 100;
-		UPDATE generated_cols SET a = 100;
 ));
 $node_publisher->wait_for_catchup('sub_dropped_cols');
 
@@ -430,11 +425,6 @@ is( $node_subscriber->safe_psql(
 	qq(1),
 	'replication with RI FULL and dropped columns');
 
-is( $node_subscriber->safe_psql(
-		'postgres', "SELECT count(*) FROM generated_cols WHERE a = 100"),
-	qq(1),
-	'replication with RI FULL and generated columns');
-
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');
 
-- 
2.34.1

From 362d0a20ccdace50055e37f996f5e168f6f86873 Mon Sep 17 00:00:00 2001
From: Shlok Kyal <shlok.kyal.oss@gmail.com>
Date: Tue, 5 Nov 2024 11:11:25 +0530
Subject: [PATCH v13] Disallow UPDATE/DELETE on table with generated column as
 REPLICA IDENTITY

UPDATE or DELETE operations on tables with unpublished generated columns
set as REPLICA IDENTITY are not permitted. This patch ensures that if an
UPDATE or DELETE command is executed on such tables, an error will be
thrown.
With this patch, the behavior has changed for the test added in commit
adedf54. Additionally, there is no other way to trigger the bug that was
fixed by commit adedf54, so the test has been removed.
---
 src/backend/commands/publicationcmds.c    | 159 +++++++++++++---------
 src/backend/executor/execReplication.c    |  32 ++++-
 src/backend/utils/cache/relcache.c        |  57 ++++++--
 src/include/catalog/pg_publication.h      |   7 +
 src/include/commands/publicationcmds.h    |   7 +-
 src/test/regress/expected/publication.out |  24 ++++
 src/test/regress/sql/publication.sql      |  25 ++++
 src/test/subscription/t/100_bugs.pl       |  16 +--
 8 files changed, 232 insertions(+), 95 deletions(-)

diff --git a/src/backend/commands/publicationcmds.c b/src/backend/commands/publicationcmds.c
index 031c84ec29..d2b4c8e9a6 100644
--- a/src/backend/commands/publicationcmds.c
+++ b/src/backend/commands/publicationcmds.c
@@ -336,21 +336,36 @@ pub_rf_contains_invalid_column(Oid pubid, Relation relation, List *ancestors,
 }
 
 /*
- * Check if all columns referenced in the REPLICA IDENTITY are covered by
- * the column list.
+ * Check for invalid columns in the publication table definition.
  *
- * Returns true if any replica identity column is not covered by column list.
+ * This function evaluates two conditions:
+ *
+ * 1. Ensures that all columns referenced in the REPLICA IDENTITY are covered
+ *    by the column list. If any column is missing, *invalid_column_list is set
+ *    to true.
+ *
+ * 2. Ensures that the REPLICA IDENTITY does not contain unpublished generated
+ *    columns. If an unpublished generated column is found,
+ *    *unpublished_gen_col is set to true.
+ *
+ * Returns true if any of the above conditions are not met.
  */
 bool
-pub_collist_contains_invalid_column(Oid pubid, Relation relation, List *ancestors,
-									bool pubviaroot)
+pub_contains_invalid_column(Oid pubid, Relation relation, List *ancestors,
+							bool pubviaroot, bool pubgencols,
+							bool *invalid_column_list,
+							bool *unpublished_gen_col)
 {
-	HeapTuple	tuple;
 	Oid			relid = RelationGetRelid(relation);
 	Oid			publish_as_relid = RelationGetRelid(relation);
-	bool		result = false;
-	Datum		datum;
-	bool		isnull;
+	Bitmapset  *idattrs;
+	Bitmapset  *columns = NULL;
+	TupleDesc	desc = RelationGetDescr(relation);
+	Publication *pub;
+	int			x;
+
+	*invalid_column_list = false;
+	*unpublished_gen_col = false;
 
 	/*
 	 * For a partition, if pubviaroot is true, find the topmost ancestor that
@@ -368,80 +383,96 @@ pub_collist_contains_invalid_column(Oid pubid, Relation relation, List *ancestor
 			publish_as_relid = relid;
 	}
 
-	tuple = SearchSysCache2(PUBLICATIONRELMAP,
-							ObjectIdGetDatum(publish_as_relid),
-							ObjectIdGetDatum(pubid));
+	/* Fetch the column list */
+	pub = GetPublication(pubid);
+	check_and_fetch_column_list(pub, publish_as_relid, NULL, &columns);
 
-	if (!HeapTupleIsValid(tuple))
-		return false;
+	if (relation->rd_rel->relreplident == REPLICA_IDENTITY_FULL)
+	{
+		/* With REPLICA IDENTITY FULL, no column list is allowed. */
+		*invalid_column_list = (columns != NULL);
 
-	datum = SysCacheGetAttr(PUBLICATIONRELMAP, tuple,
-							Anum_pg_publication_rel_prattrs,
-							&isnull);
+		/*
+		 * REPLICA IDENTITY can be FULL only if there is no column list for
+		 * publication. When REPLICA IDENTITY is FULL and the relation
+		 * includes a generated column, but the publish_generated_columns
+		 * option is set to false, this scenario is invalid.
+		 */
+		if (!pubgencols && relation->rd_att->constr &&
+			relation->rd_att->constr->has_generated_stored)
+			*unpublished_gen_col = true;
 
-	if (!isnull)
-	{
-		int			x;
-		Bitmapset  *idattrs;
-		Bitmapset  *columns = NULL;
+		if (*unpublished_gen_col && *invalid_column_list)
+			return true;
+	}
 
-		/* With REPLICA IDENTITY FULL, no column list is allowed. */
-		if (relation->rd_rel->relreplident == REPLICA_IDENTITY_FULL)
-			result = true;
+	/* Remember columns that are part of the REPLICA IDENTITY */
+	idattrs = RelationGetIndexAttrBitmap(relation,
+										 INDEX_ATTR_BITMAP_IDENTITY_KEY);
 
-		/* Transform the column list datum to a bitmapset. */
-		columns = pub_collist_to_bitmapset(NULL, datum, NULL);
+	/*
+	 * Attnums in the bitmap returned by RelationGetIndexAttrBitmap are offset
+	 * (to handle system columns the usual way), while column list does not
+	 * use offset, so we can't do bms_is_subset(). Instead, we have to loop
+	 * over the idattrs and check all of them are in the list.
+	 */
+	x = -1;
+	while ((x = bms_next_member(idattrs, x)) >= 0)
+	{
+		AttrNumber	attnum = (x + FirstLowInvalidHeapAttributeNumber);
+		Form_pg_attribute att = TupleDescAttr(desc, attnum - 1);
 
-		/* Remember columns that are part of the REPLICA IDENTITY */
-		idattrs = RelationGetIndexAttrBitmap(relation,
-											 INDEX_ATTR_BITMAP_IDENTITY_KEY);
+		Assert(!att->attisdropped);
 
-		/*
-		 * Attnums in the bitmap returned by RelationGetIndexAttrBitmap are
-		 * offset (to handle system columns the usual way), while column list
-		 * does not use offset, so we can't do bms_is_subset(). Instead, we
-		 * have to loop over the idattrs and check all of them are in the
-		 * list.
-		 */
-		x = -1;
-		while ((x = bms_next_member(idattrs, x)) >= 0)
+		if (columns == NULL)
 		{
-			AttrNumber	attnum = (x + FirstLowInvalidHeapAttributeNumber);
-
 			/*
-			 * If pubviaroot is true, we are validating the column list of the
-			 * parent table, but the bitmap contains the replica identity
-			 * information of the child table. The parent/child attnums may
-			 * not match, so translate them to the parent - get the attname
-			 * from the child, and look it up in the parent.
+			 * Check if pubgencols is false and generated column is part of
+			 * REPLICA IDENTITY
 			 */
-			if (pubviaroot)
+			if (!pubgencols)
 			{
-				/* attribute name in the child table */
-				char	   *colname = get_attname(relid, attnum, false);
+				*unpublished_gen_col |= att->attgenerated;
 
-				/*
-				 * Determine the attnum for the attribute name in parent (we
-				 * are using the column list defined on the parent).
-				 */
-				attnum = get_attnum(publish_as_relid, colname);
+				/* Break the loop if unpublished generated columns exist. */
+				if (*unpublished_gen_col)
+					break;
 			}
 
-			/* replica identity column, not covered by the column list */
-			if (!bms_is_member(attnum, columns))
-			{
-				result = true;
-				break;
-			}
+			/* Skip validating the column list since it is not defined */
+			continue;
+		}
+
+		/*
+		 * If pubviaroot is true, we are validating the column list of the
+		 * parent table, but the bitmap contains the replica identity
+		 * information of the child table. The parent/child attnums may not
+		 * match, so translate them to the parent - get the attname from the
+		 * child, and look it up in the parent.
+		 */
+		if (pubviaroot)
+		{
+			/* attribute name in the child table */
+			char	   *colname = get_attname(relid, attnum, false);
+
+			/*
+			 * Determine the attnum for the attribute name in parent (we are
+			 * using the column list defined on the parent).
+			 */
+			attnum = get_attnum(publish_as_relid, colname);
 		}
 
-		bms_free(idattrs);
-		bms_free(columns);
+		/* replica identity column, not covered by the column list */
+		*invalid_column_list |= !bms_is_member(attnum, columns);
+
+		if (*invalid_column_list && *unpublished_gen_col)
+			break;
 	}
 
-	ReleaseSysCache(tuple);
+	bms_free(columns);
+	bms_free(idattrs);
 
-	return result;
+	return *invalid_column_list || *unpublished_gen_col;
 }
 
 /* check_functions_in_node callback */
diff --git a/src/backend/executor/execReplication.c b/src/backend/executor/execReplication.c
index 54025c9f15..227a8aeea0 100644
--- a/src/backend/executor/execReplication.c
+++ b/src/backend/executor/execReplication.c
@@ -785,10 +785,22 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 		return;
 
 	/*
-	 * 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 - i.e. when all referenced columns are part of REPLICA IDENTITY
-	 * or the table does not publish UPDATEs or DELETEs.
+	 * It is only safe to execute UPDATE/DELETE when:
+	 *
+	 * 1. All columns, referenced in the row filters from publications which
+	 * the relation is in, are valid - i.e. when all referenced columns are
+	 * part of REPLICA IDENTITY or the table does not publish UPDATEs or
+	 * DELETEs.
+	 *
+	 * 2. All columns, referenced in the column lists from publications which
+	 * the relation is in, are valid - i.e. when all referenced columns are
+	 * part of REPLICA IDENTITY or the table does not publish UPDATEs or
+	 * DELETEs.
+	 *
+	 * 3. All generated columns in REPLICA IDENTITY of the relation, for all
+	 * the publications which the relation is in, are valid - i.e. when
+	 * unpublished generated columns are not part of REPLICA IDENTITY or the
+	 * table does not publish UPDATEs or DELETEs.
 	 *
 	 * XXX We could optimize it by first checking whether any of the
 	 * publications have a row filter for this relation. If not and relation
@@ -809,6 +821,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot update table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_UPDATE && !pubdesc.gencols_valid_for_update)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot update table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("Replica identity consists of an unpublished generated column.")));
 	else if (cmd == CMD_DELETE && !pubdesc.rf_valid_for_delete)
 		ereport(ERROR,
 				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
@@ -821,6 +839,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot delete from table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_DELETE && !pubdesc.gencols_valid_for_delete)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot delete from table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("Replica identity consists of an unpublished generated column.")));
 
 	/* If relation has replica identity we are always good. */
 	if (OidIsValid(RelationGetReplicaIndex(rel)))
diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
index d0892cee24..9205ea86c4 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -5706,12 +5706,19 @@ RelationGetExclusionInfo(Relation indexRelation,
  * Get the publication information for the given relation.
  *
  * Traverse all the publications which the relation is in to get the
- * publication actions and validate the row filter expressions for such
- * publications if any. We consider the row filter expression as invalid if it
- * references any column which is not part of REPLICA IDENTITY.
+ * publication actions and validate:
+ * 1. The row filter expressions for such publications if any. We consider the
+ *    row filter expression as invalid if it references any column which is not
+ *    part of REPLICA IDENTITY.
+ * 2. The column list for such publication if any. We consider column list as
+ * 	  invalid if it references any column which is not part of REPLICA IDENTITY.
+ * 3. The generated columns of the relation for such publications. We consider
+ *    any reference of an unpublished generated column in REPLICA IDENTITY as
+ *    invalid.
  *
  * To avoid fetching the publication information repeatedly, we cache the
- * publication actions and row filter validation information.
+ * publication actions, row filter validation information, column list
+ * validation information and generated column validation information.
  */
 void
 RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
@@ -5734,6 +5741,8 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		pubdesc->rf_valid_for_delete = true;
 		pubdesc->cols_valid_for_update = true;
 		pubdesc->cols_valid_for_delete = true;
+		pubdesc->gencols_valid_for_update = true;
+		pubdesc->gencols_valid_for_delete = true;
 		return;
 	}
 
@@ -5748,6 +5757,8 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 	pubdesc->rf_valid_for_delete = true;
 	pubdesc->cols_valid_for_update = true;
 	pubdesc->cols_valid_for_delete = true;
+	pubdesc->gencols_valid_for_update = true;
+	pubdesc->gencols_valid_for_delete = true;
 
 	/* Fetch the publication membership info. */
 	puboids = GetRelationPublications(relid);
@@ -5777,6 +5788,8 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		Oid			pubid = lfirst_oid(lc);
 		HeapTuple	tup;
 		Form_pg_publication pubform;
+		bool		invalid_column_list;
+		bool		unpublished_gen_col;
 
 		tup = SearchSysCache1(PUBLICATIONOID, ObjectIdGetDatum(pubid));
 
@@ -5811,18 +5824,27 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		/*
 		 * Check if all columns are part of the REPLICA IDENTITY index or not.
 		 *
-		 * If the publication is FOR ALL TABLES then it means the table has no
-		 * column list and we can skip the validation.
+		 * Check if all generated columns included in the REPLICA IDENTITY are
+		 * published.
 		 */
-		if (!pubform->puballtables &&
-			(pubform->pubupdate || pubform->pubdelete) &&
-			pub_collist_contains_invalid_column(pubid, relation, ancestors,
-												pubform->pubviaroot))
+		if ((pubform->pubupdate || pubform->pubdelete) &&
+			pub_contains_invalid_column(pubid, relation, ancestors,
+										pubform->pubviaroot,
+										pubform->pubgencols,
+										&invalid_column_list,
+										&unpublished_gen_col))
 		{
 			if (pubform->pubupdate)
-				pubdesc->cols_valid_for_update = false;
+			{
+				pubdesc->cols_valid_for_update = !invalid_column_list;
+				pubdesc->gencols_valid_for_update = !unpublished_gen_col;
+			}
+
 			if (pubform->pubdelete)
-				pubdesc->cols_valid_for_delete = false;
+			{
+				pubdesc->cols_valid_for_delete = !invalid_column_list;
+				pubdesc->gencols_valid_for_delete = !unpublished_gen_col;
+			}
 		}
 
 		ReleaseSysCache(tup);
@@ -5846,6 +5868,17 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 			pubdesc->pubactions.pubdelete && pubdesc->pubactions.pubtruncate &&
 			!pubdesc->cols_valid_for_update && !pubdesc->cols_valid_for_delete)
 			break;
+
+		/*
+		 * If we know everything is replicated and replica identity has an
+		 * unpublished generated column, there is no point to check for other
+		 * publications.
+		 */
+		if (pubdesc->pubactions.pubinsert && pubdesc->pubactions.pubupdate &&
+			pubdesc->pubactions.pubdelete && pubdesc->pubactions.pubtruncate &&
+			!pubdesc->gencols_valid_for_update &&
+			!pubdesc->gencols_valid_for_delete)
+			break;
 	}
 
 	if (relation->rd_pubdesc)
diff --git a/src/include/catalog/pg_publication.h b/src/include/catalog/pg_publication.h
index 9a83a72d6b..e167b34461 100644
--- a/src/include/catalog/pg_publication.h
+++ b/src/include/catalog/pg_publication.h
@@ -98,6 +98,13 @@ typedef struct PublicationDesc
 	 */
 	bool		cols_valid_for_update;
 	bool		cols_valid_for_delete;
+
+	/*
+	 * true if all generated columns which are part of replica identity are
+	 * published or the publication actions do not include UPDATE or DELETE.
+	 */
+	bool		gencols_valid_for_update;
+	bool		gencols_valid_for_delete;
 } PublicationDesc;
 
 typedef struct Publication
diff --git a/src/include/commands/publicationcmds.h b/src/include/commands/publicationcmds.h
index 5487c571f6..c200c19c6d 100644
--- a/src/include/commands/publicationcmds.h
+++ b/src/include/commands/publicationcmds.h
@@ -33,7 +33,10 @@ extern void AlterPublicationOwner_oid(Oid subid, Oid newOwnerId);
 extern void InvalidatePublicationRels(List *relids);
 extern bool pub_rf_contains_invalid_column(Oid pubid, Relation relation,
 										   List *ancestors, bool pubviaroot);
-extern bool pub_collist_contains_invalid_column(Oid pubid, Relation relation,
-												List *ancestors, bool pubviaroot);
+extern bool pub_contains_invalid_column(Oid pubid, Relation relation,
+										List *ancestors, bool pubviaroot,
+										bool pubgencols,
+										bool *invalid_column_list,
+										bool *unpublished_gen_col);
 
 #endif							/* PUBLICATIONCMDS_H */
diff --git a/src/test/regress/expected/publication.out b/src/test/regress/expected/publication.out
index 5de2d64d01..12d0611a7c 100644
--- a/src/test/regress/expected/publication.out
+++ b/src/test/regress/expected/publication.out
@@ -672,6 +672,30 @@ DROP TABLE rf_tbl_abcd_pk;
 DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  Replica identity consists of an unpublished generated column.
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  Replica identity consists of an unpublished generated column.
+DROP PUBLICATION pub_gencol;
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/regress/sql/publication.sql b/src/test/regress/sql/publication.sql
index 48e68bcca2..a29587b45d 100644
--- a/src/test/regress/sql/publication.sql
+++ b/src/test/regress/sql/publication.sql
@@ -396,6 +396,31 @@ DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
 
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
+
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/subscription/t/100_bugs.pl b/src/test/subscription/t/100_bugs.pl
index cb36ca7b16..794b928f50 100644
--- a/src/test/subscription/t/100_bugs.pl
+++ b/src/test/subscription/t/100_bugs.pl
@@ -377,8 +377,8 @@ $node_publisher->safe_psql('postgres', "DROP PUBLICATION tap_pub_sch");
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');
 
-# The bug was that when the REPLICA IDENTITY FULL is used with dropped or
-# generated columns, we fail to apply updates and deletes
+# The bug was that when the REPLICA IDENTITY FULL is used with dropped
+# we fail to apply updates and deletes
 $node_publisher->rotate_logfile();
 $node_publisher->start();
 
@@ -389,18 +389,14 @@ $node_publisher->safe_psql(
 	'postgres', qq(
 	CREATE TABLE dropped_cols (a int, b_drop int, c int);
 	ALTER TABLE dropped_cols REPLICA IDENTITY FULL;
-	CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
-	ALTER TABLE generated_cols REPLICA IDENTITY FULL;
-	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols, generated_cols;
+	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols;
 	-- some initial data
 	INSERT INTO dropped_cols VALUES (1, 1, 1);
-	INSERT INTO generated_cols (a, c) VALUES (1, 1);
 ));
 
 $node_subscriber->safe_psql(
 	'postgres', qq(
 	 CREATE TABLE dropped_cols (a int, b_drop int, c int);
-	 CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
 ));
 
 $publisher_connstr = $node_publisher->connstr . ' dbname=postgres';
@@ -421,7 +417,6 @@ $node_subscriber->safe_psql(
 $node_publisher->safe_psql(
 	'postgres', qq(
 		UPDATE dropped_cols SET a = 100;
-		UPDATE generated_cols SET a = 100;
 ));
 $node_publisher->wait_for_catchup('sub_dropped_cols');
 
@@ -430,11 +425,6 @@ is( $node_subscriber->safe_psql(
 	qq(1),
 	'replication with RI FULL and dropped columns');
 
-is( $node_subscriber->safe_psql(
-		'postgres', "SELECT count(*) FROM generated_cols WHERE a = 100"),
-	qq(1),
-	'replication with RI FULL and generated columns');
-
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');
 
-- 
2.34.1

From abdc729ab5e880543e4702d65e8ea66533325d71 Mon Sep 17 00:00:00 2001
From: Hou Zhijie <houzj.fnst@cn.fujitsu.com>
Date: Tue, 3 Dec 2024 12:19:50 +0800
Subject: [PATCH] improvements

---
 doc/src/sgml/ref/create_publication.sgml |  9 ++++++++
 src/backend/commands/publicationcmds.c   | 13 +++++-------
 src/backend/executor/execReplication.c   | 27 ++++++++++++------------
 3 files changed, 27 insertions(+), 22 deletions(-)

diff --git a/doc/src/sgml/ref/create_publication.sgml b/doc/src/sgml/ref/create_publication.sgml
index f8e217d661..ae21018697 100644
--- a/doc/src/sgml/ref/create_publication.sgml
+++ b/doc/src/sgml/ref/create_publication.sgml
@@ -311,6 +311,15 @@ CREATE PUBLICATION <replaceable class="parameter">name</replaceable>
    system columns.
   </para>
 
+  <para>
+   If tables added to a publication include generated columns that are part of
+   the <literal>REPLICA IDENTITY</literal>, it is essential to publish these
+   columns by explicitly listing them in the column list or by enabling the
+   <literal>publish_generated_columns</literal> option. Otherwise,
+   <command>UPDATE</command> or <command>DELETE</command> operations will be
+   disallowed on those tables.
+  </para>
+
   <para>
    The row filter on a table becomes redundant if
    <literal>FOR TABLES IN SCHEMA</literal> is specified and the table
diff --git a/src/backend/commands/publicationcmds.c b/src/backend/commands/publicationcmds.c
index d2b4c8e9a6..323f59fae1 100644
--- a/src/backend/commands/publicationcmds.c
+++ b/src/backend/commands/publicationcmds.c
@@ -427,16 +427,13 @@ pub_contains_invalid_column(Oid pubid, Relation relation, List *ancestors,
 		if (columns == NULL)
 		{
 			/*
-			 * Check if pubgencols is false and generated column is part of
-			 * REPLICA IDENTITY
+			 * Break the loop if an unpublished generated column is part of the
+			 * REPLICA IDENTITY.
 			 */
-			if (!pubgencols)
+			if (!pubgencols && att->attgenerated)
 			{
-				*unpublished_gen_col |= att->attgenerated;
-
-				/* Break the loop if unpublished generated columns exist. */
-				if (*unpublished_gen_col)
-					break;
+				*unpublished_gen_col = true;
+				break;
 			}
 
 			/* Skip validating the column list since it is not defined */
diff --git a/src/backend/executor/execReplication.c b/src/backend/executor/execReplication.c
index 227a8aeea0..f66ff21159 100644
--- a/src/backend/executor/execReplication.c
+++ b/src/backend/executor/execReplication.c
@@ -785,27 +785,26 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 		return;
 
 	/*
-	 * It is only safe to execute UPDATE/DELETE when:
+	 * It is only safe to execute UPDATE/DELETE if the relation does not
+	 * publish UPDATEs or DELETEs, or all the following conditions are
+	 * satisfied:
 	 *
 	 * 1. All columns, referenced in the row filters from publications which
-	 * the relation is in, are valid - i.e. when all referenced columns are
-	 * part of REPLICA IDENTITY or the table does not publish UPDATEs or
-	 * DELETEs.
+	 *    the relation is in, are valid - i.e. when all referenced columns are
+	 *    part of REPLICA IDENTITY.
 	 *
 	 * 2. All columns, referenced in the column lists from publications which
-	 * the relation is in, are valid - i.e. when all referenced columns are
-	 * part of REPLICA IDENTITY or the table does not publish UPDATEs or
-	 * DELETEs.
+	 *    the relation is in, are valid - i.e. when all columns referenced in
+	 *    the REPLICA IDENTITY are covered by the column list.
 	 *
-	 * 3. All generated columns in REPLICA IDENTITY of the relation, for all
-	 * the publications which the relation is in, are valid - i.e. when
-	 * unpublished generated columns are not part of REPLICA IDENTITY or the
-	 * table does not publish UPDATEs or DELETEs.
+	 * 3. All generated columns in REPLICA IDENTITY of the relation, are valid
+	 *    - i.e. when all these generated columns are published.
 	 *
 	 * XXX We could optimize it by first checking whether any of the
-	 * publications have a row filter for this relation. If not and relation
-	 * has replica identity then we can avoid building the descriptor but as
-	 * this happens only one time it doesn't seem worth the additional
+	 * publications have a row filter or column list for this relation, or if
+	 * the relation contains a generated column. If none of these exist and the
+	 * relation has replica identity then we can avoid building the descriptor
+	 * but as this happens only one time it doesn't seem worth the additional
 	 * complexity.
 	 */
 	RelationBuildPublicationDesc(rel, &pubdesc);
-- 
2.30.0.windows.2

From 0421de2cd268104b9bb87f12452fa20d5acdaeeb Mon Sep 17 00:00:00 2001
From: Shlok Kyal <shlok.kyal.oss@gmail.com>
Date: Tue, 5 Nov 2024 11:11:25 +0530
Subject: [PATCH v14] Disallow UPDATE/DELETE on table with generated column as
 REPLICA IDENTITY

UPDATE or DELETE operations on tables with unpublished generated columns
set as REPLICA IDENTITY are not permitted. This patch ensures that if an
UPDATE or DELETE command is executed on such tables, an error will be
thrown.
With this patch, the behavior has changed for the test added in commit
adedf54. Additionally, there is no other way to trigger the bug that was
fixed by commit adedf54, so the test has been removed.
---
 doc/src/sgml/ref/create_publication.sgml  |   9 ++
 src/backend/commands/publicationcmds.c    | 156 +++++++++++++---------
 src/backend/executor/execReplication.c    |  37 ++++-
 src/backend/utils/cache/relcache.c        |  57 ++++++--
 src/include/catalog/pg_publication.h      |   7 +
 src/include/commands/publicationcmds.h    |   7 +-
 src/test/regress/expected/publication.out |  24 ++++
 src/test/regress/sql/publication.sql      |  25 ++++
 src/test/subscription/t/100_bugs.pl       |  16 +--
 9 files changed, 240 insertions(+), 98 deletions(-)

diff --git a/doc/src/sgml/ref/create_publication.sgml b/doc/src/sgml/ref/create_publication.sgml
index f8e217d661..ae21018697 100644
--- a/doc/src/sgml/ref/create_publication.sgml
+++ b/doc/src/sgml/ref/create_publication.sgml
@@ -311,6 +311,15 @@ CREATE PUBLICATION <replaceable class="parameter">name</replaceable>
    system columns.
   </para>
 
+  <para>
+   If tables added to a publication include generated columns that are part of
+   the <literal>REPLICA IDENTITY</literal>, it is essential to publish these
+   columns by explicitly listing them in the column list or by enabling the
+   <literal>publish_generated_columns</literal> option. Otherwise,
+   <command>UPDATE</command> or <command>DELETE</command> operations will be
+   disallowed on those tables.
+  </para>
+
   <para>
    The row filter on a table becomes redundant if
    <literal>FOR TABLES IN SCHEMA</literal> is specified and the table
diff --git a/src/backend/commands/publicationcmds.c b/src/backend/commands/publicationcmds.c
index 031c84ec29..323f59fae1 100644
--- a/src/backend/commands/publicationcmds.c
+++ b/src/backend/commands/publicationcmds.c
@@ -336,21 +336,36 @@ pub_rf_contains_invalid_column(Oid pubid, Relation relation, List *ancestors,
 }
 
 /*
- * Check if all columns referenced in the REPLICA IDENTITY are covered by
- * the column list.
+ * Check for invalid columns in the publication table definition.
  *
- * Returns true if any replica identity column is not covered by column list.
+ * This function evaluates two conditions:
+ *
+ * 1. Ensures that all columns referenced in the REPLICA IDENTITY are covered
+ *    by the column list. If any column is missing, *invalid_column_list is set
+ *    to true.
+ *
+ * 2. Ensures that the REPLICA IDENTITY does not contain unpublished generated
+ *    columns. If an unpublished generated column is found,
+ *    *unpublished_gen_col is set to true.
+ *
+ * Returns true if any of the above conditions are not met.
  */
 bool
-pub_collist_contains_invalid_column(Oid pubid, Relation relation, List *ancestors,
-									bool pubviaroot)
+pub_contains_invalid_column(Oid pubid, Relation relation, List *ancestors,
+							bool pubviaroot, bool pubgencols,
+							bool *invalid_column_list,
+							bool *unpublished_gen_col)
 {
-	HeapTuple	tuple;
 	Oid			relid = RelationGetRelid(relation);
 	Oid			publish_as_relid = RelationGetRelid(relation);
-	bool		result = false;
-	Datum		datum;
-	bool		isnull;
+	Bitmapset  *idattrs;
+	Bitmapset  *columns = NULL;
+	TupleDesc	desc = RelationGetDescr(relation);
+	Publication *pub;
+	int			x;
+
+	*invalid_column_list = false;
+	*unpublished_gen_col = false;
 
 	/*
 	 * For a partition, if pubviaroot is true, find the topmost ancestor that
@@ -368,80 +383,93 @@ pub_collist_contains_invalid_column(Oid pubid, Relation relation, List *ancestor
 			publish_as_relid = relid;
 	}
 
-	tuple = SearchSysCache2(PUBLICATIONRELMAP,
-							ObjectIdGetDatum(publish_as_relid),
-							ObjectIdGetDatum(pubid));
+	/* Fetch the column list */
+	pub = GetPublication(pubid);
+	check_and_fetch_column_list(pub, publish_as_relid, NULL, &columns);
 
-	if (!HeapTupleIsValid(tuple))
-		return false;
+	if (relation->rd_rel->relreplident == REPLICA_IDENTITY_FULL)
+	{
+		/* With REPLICA IDENTITY FULL, no column list is allowed. */
+		*invalid_column_list = (columns != NULL);
 
-	datum = SysCacheGetAttr(PUBLICATIONRELMAP, tuple,
-							Anum_pg_publication_rel_prattrs,
-							&isnull);
+		/*
+		 * REPLICA IDENTITY can be FULL only if there is no column list for
+		 * publication. When REPLICA IDENTITY is FULL and the relation
+		 * includes a generated column, but the publish_generated_columns
+		 * option is set to false, this scenario is invalid.
+		 */
+		if (!pubgencols && relation->rd_att->constr &&
+			relation->rd_att->constr->has_generated_stored)
+			*unpublished_gen_col = true;
 
-	if (!isnull)
+		if (*unpublished_gen_col && *invalid_column_list)
+			return true;
+	}
+
+	/* Remember columns that are part of the REPLICA IDENTITY */
+	idattrs = RelationGetIndexAttrBitmap(relation,
+										 INDEX_ATTR_BITMAP_IDENTITY_KEY);
+
+	/*
+	 * Attnums in the bitmap returned by RelationGetIndexAttrBitmap are offset
+	 * (to handle system columns the usual way), while column list does not
+	 * use offset, so we can't do bms_is_subset(). Instead, we have to loop
+	 * over the idattrs and check all of them are in the list.
+	 */
+	x = -1;
+	while ((x = bms_next_member(idattrs, x)) >= 0)
 	{
-		int			x;
-		Bitmapset  *idattrs;
-		Bitmapset  *columns = NULL;
+		AttrNumber	attnum = (x + FirstLowInvalidHeapAttributeNumber);
+		Form_pg_attribute att = TupleDescAttr(desc, attnum - 1);
 
-		/* With REPLICA IDENTITY FULL, no column list is allowed. */
-		if (relation->rd_rel->relreplident == REPLICA_IDENTITY_FULL)
-			result = true;
+		Assert(!att->attisdropped);
 
-		/* Transform the column list datum to a bitmapset. */
-		columns = pub_collist_to_bitmapset(NULL, datum, NULL);
+		if (columns == NULL)
+		{
+			/*
+			 * Break the loop if an unpublished generated column is part of the
+			 * REPLICA IDENTITY.
+			 */
+			if (!pubgencols && att->attgenerated)
+			{
+				*unpublished_gen_col = true;
+				break;
+			}
 
-		/* Remember columns that are part of the REPLICA IDENTITY */
-		idattrs = RelationGetIndexAttrBitmap(relation,
-											 INDEX_ATTR_BITMAP_IDENTITY_KEY);
+			/* Skip validating the column list since it is not defined */
+			continue;
+		}
 
 		/*
-		 * Attnums in the bitmap returned by RelationGetIndexAttrBitmap are
-		 * offset (to handle system columns the usual way), while column list
-		 * does not use offset, so we can't do bms_is_subset(). Instead, we
-		 * have to loop over the idattrs and check all of them are in the
-		 * list.
+		 * If pubviaroot is true, we are validating the column list of the
+		 * parent table, but the bitmap contains the replica identity
+		 * information of the child table. The parent/child attnums may not
+		 * match, so translate them to the parent - get the attname from the
+		 * child, and look it up in the parent.
 		 */
-		x = -1;
-		while ((x = bms_next_member(idattrs, x)) >= 0)
+		if (pubviaroot)
 		{
-			AttrNumber	attnum = (x + FirstLowInvalidHeapAttributeNumber);
+			/* attribute name in the child table */
+			char	   *colname = get_attname(relid, attnum, false);
 
 			/*
-			 * If pubviaroot is true, we are validating the column list of the
-			 * parent table, but the bitmap contains the replica identity
-			 * information of the child table. The parent/child attnums may
-			 * not match, so translate them to the parent - get the attname
-			 * from the child, and look it up in the parent.
+			 * Determine the attnum for the attribute name in parent (we are
+			 * using the column list defined on the parent).
 			 */
-			if (pubviaroot)
-			{
-				/* attribute name in the child table */
-				char	   *colname = get_attname(relid, attnum, false);
-
-				/*
-				 * Determine the attnum for the attribute name in parent (we
-				 * are using the column list defined on the parent).
-				 */
-				attnum = get_attnum(publish_as_relid, colname);
-			}
-
-			/* replica identity column, not covered by the column list */
-			if (!bms_is_member(attnum, columns))
-			{
-				result = true;
-				break;
-			}
+			attnum = get_attnum(publish_as_relid, colname);
 		}
 
-		bms_free(idattrs);
-		bms_free(columns);
+		/* replica identity column, not covered by the column list */
+		*invalid_column_list |= !bms_is_member(attnum, columns);
+
+		if (*invalid_column_list && *unpublished_gen_col)
+			break;
 	}
 
-	ReleaseSysCache(tuple);
+	bms_free(columns);
+	bms_free(idattrs);
 
-	return result;
+	return *invalid_column_list || *unpublished_gen_col;
 }
 
 /* check_functions_in_node callback */
diff --git a/src/backend/executor/execReplication.c b/src/backend/executor/execReplication.c
index 54025c9f15..f66ff21159 100644
--- a/src/backend/executor/execReplication.c
+++ b/src/backend/executor/execReplication.c
@@ -785,15 +785,26 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 		return;
 
 	/*
-	 * 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 - i.e. when all referenced columns are part of REPLICA IDENTITY
-	 * or the table does not publish UPDATEs or DELETEs.
+	 * It is only safe to execute UPDATE/DELETE if the relation does not
+	 * publish UPDATEs or DELETEs, or all the following conditions are
+	 * satisfied:
+	 *
+	 * 1. All columns, referenced in the row filters from publications which
+	 *    the relation is in, are valid - i.e. when all referenced columns are
+	 *    part of REPLICA IDENTITY.
+	 *
+	 * 2. All columns, referenced in the column lists from publications which
+	 *    the relation is in, are valid - i.e. when all columns referenced in
+	 *    the REPLICA IDENTITY are covered by the column list.
+	 *
+	 * 3. All generated columns in REPLICA IDENTITY of the relation, are valid
+	 *    - i.e. when all these generated columns are published.
 	 *
 	 * XXX We could optimize it by first checking whether any of the
-	 * publications have a row filter for this relation. If not and relation
-	 * has replica identity then we can avoid building the descriptor but as
-	 * this happens only one time it doesn't seem worth the additional
+	 * publications have a row filter or column list for this relation, or if
+	 * the relation contains a generated column. If none of these exist and the
+	 * relation has replica identity then we can avoid building the descriptor
+	 * but as this happens only one time it doesn't seem worth the additional
 	 * complexity.
 	 */
 	RelationBuildPublicationDesc(rel, &pubdesc);
@@ -809,6 +820,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot update table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_UPDATE && !pubdesc.gencols_valid_for_update)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot update table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("Replica identity consists of an unpublished generated column.")));
 	else if (cmd == CMD_DELETE && !pubdesc.rf_valid_for_delete)
 		ereport(ERROR,
 				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
@@ -821,6 +838,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot delete from table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_DELETE && !pubdesc.gencols_valid_for_delete)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot delete from table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("Replica identity consists of an unpublished generated column.")));
 
 	/* If relation has replica identity we are always good. */
 	if (OidIsValid(RelationGetReplicaIndex(rel)))
diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
index d0892cee24..9205ea86c4 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -5706,12 +5706,19 @@ RelationGetExclusionInfo(Relation indexRelation,
  * Get the publication information for the given relation.
  *
  * Traverse all the publications which the relation is in to get the
- * publication actions and validate the row filter expressions for such
- * publications if any. We consider the row filter expression as invalid if it
- * references any column which is not part of REPLICA IDENTITY.
+ * publication actions and validate:
+ * 1. The row filter expressions for such publications if any. We consider the
+ *    row filter expression as invalid if it references any column which is not
+ *    part of REPLICA IDENTITY.
+ * 2. The column list for such publication if any. We consider column list as
+ * 	  invalid if it references any column which is not part of REPLICA IDENTITY.
+ * 3. The generated columns of the relation for such publications. We consider
+ *    any reference of an unpublished generated column in REPLICA IDENTITY as
+ *    invalid.
  *
  * To avoid fetching the publication information repeatedly, we cache the
- * publication actions and row filter validation information.
+ * publication actions, row filter validation information, column list
+ * validation information and generated column validation information.
  */
 void
 RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
@@ -5734,6 +5741,8 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		pubdesc->rf_valid_for_delete = true;
 		pubdesc->cols_valid_for_update = true;
 		pubdesc->cols_valid_for_delete = true;
+		pubdesc->gencols_valid_for_update = true;
+		pubdesc->gencols_valid_for_delete = true;
 		return;
 	}
 
@@ -5748,6 +5757,8 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 	pubdesc->rf_valid_for_delete = true;
 	pubdesc->cols_valid_for_update = true;
 	pubdesc->cols_valid_for_delete = true;
+	pubdesc->gencols_valid_for_update = true;
+	pubdesc->gencols_valid_for_delete = true;
 
 	/* Fetch the publication membership info. */
 	puboids = GetRelationPublications(relid);
@@ -5777,6 +5788,8 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		Oid			pubid = lfirst_oid(lc);
 		HeapTuple	tup;
 		Form_pg_publication pubform;
+		bool		invalid_column_list;
+		bool		unpublished_gen_col;
 
 		tup = SearchSysCache1(PUBLICATIONOID, ObjectIdGetDatum(pubid));
 
@@ -5811,18 +5824,27 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		/*
 		 * Check if all columns are part of the REPLICA IDENTITY index or not.
 		 *
-		 * If the publication is FOR ALL TABLES then it means the table has no
-		 * column list and we can skip the validation.
+		 * Check if all generated columns included in the REPLICA IDENTITY are
+		 * published.
 		 */
-		if (!pubform->puballtables &&
-			(pubform->pubupdate || pubform->pubdelete) &&
-			pub_collist_contains_invalid_column(pubid, relation, ancestors,
-												pubform->pubviaroot))
+		if ((pubform->pubupdate || pubform->pubdelete) &&
+			pub_contains_invalid_column(pubid, relation, ancestors,
+										pubform->pubviaroot,
+										pubform->pubgencols,
+										&invalid_column_list,
+										&unpublished_gen_col))
 		{
 			if (pubform->pubupdate)
-				pubdesc->cols_valid_for_update = false;
+			{
+				pubdesc->cols_valid_for_update = !invalid_column_list;
+				pubdesc->gencols_valid_for_update = !unpublished_gen_col;
+			}
+
 			if (pubform->pubdelete)
-				pubdesc->cols_valid_for_delete = false;
+			{
+				pubdesc->cols_valid_for_delete = !invalid_column_list;
+				pubdesc->gencols_valid_for_delete = !unpublished_gen_col;
+			}
 		}
 
 		ReleaseSysCache(tup);
@@ -5846,6 +5868,17 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 			pubdesc->pubactions.pubdelete && pubdesc->pubactions.pubtruncate &&
 			!pubdesc->cols_valid_for_update && !pubdesc->cols_valid_for_delete)
 			break;
+
+		/*
+		 * If we know everything is replicated and replica identity has an
+		 * unpublished generated column, there is no point to check for other
+		 * publications.
+		 */
+		if (pubdesc->pubactions.pubinsert && pubdesc->pubactions.pubupdate &&
+			pubdesc->pubactions.pubdelete && pubdesc->pubactions.pubtruncate &&
+			!pubdesc->gencols_valid_for_update &&
+			!pubdesc->gencols_valid_for_delete)
+			break;
 	}
 
 	if (relation->rd_pubdesc)
diff --git a/src/include/catalog/pg_publication.h b/src/include/catalog/pg_publication.h
index 9a83a72d6b..e167b34461 100644
--- a/src/include/catalog/pg_publication.h
+++ b/src/include/catalog/pg_publication.h
@@ -98,6 +98,13 @@ typedef struct PublicationDesc
 	 */
 	bool		cols_valid_for_update;
 	bool		cols_valid_for_delete;
+
+	/*
+	 * true if all generated columns which are part of replica identity are
+	 * published or the publication actions do not include UPDATE or DELETE.
+	 */
+	bool		gencols_valid_for_update;
+	bool		gencols_valid_for_delete;
 } PublicationDesc;
 
 typedef struct Publication
diff --git a/src/include/commands/publicationcmds.h b/src/include/commands/publicationcmds.h
index 5487c571f6..c200c19c6d 100644
--- a/src/include/commands/publicationcmds.h
+++ b/src/include/commands/publicationcmds.h
@@ -33,7 +33,10 @@ extern void AlterPublicationOwner_oid(Oid subid, Oid newOwnerId);
 extern void InvalidatePublicationRels(List *relids);
 extern bool pub_rf_contains_invalid_column(Oid pubid, Relation relation,
 										   List *ancestors, bool pubviaroot);
-extern bool pub_collist_contains_invalid_column(Oid pubid, Relation relation,
-												List *ancestors, bool pubviaroot);
+extern bool pub_contains_invalid_column(Oid pubid, Relation relation,
+										List *ancestors, bool pubviaroot,
+										bool pubgencols,
+										bool *invalid_column_list,
+										bool *unpublished_gen_col);
 
 #endif							/* PUBLICATIONCMDS_H */
diff --git a/src/test/regress/expected/publication.out b/src/test/regress/expected/publication.out
index 5de2d64d01..12d0611a7c 100644
--- a/src/test/regress/expected/publication.out
+++ b/src/test/regress/expected/publication.out
@@ -672,6 +672,30 @@ DROP TABLE rf_tbl_abcd_pk;
 DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  Replica identity consists of an unpublished generated column.
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  Replica identity consists of an unpublished generated column.
+DROP PUBLICATION pub_gencol;
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/regress/sql/publication.sql b/src/test/regress/sql/publication.sql
index 48e68bcca2..a29587b45d 100644
--- a/src/test/regress/sql/publication.sql
+++ b/src/test/regress/sql/publication.sql
@@ -396,6 +396,31 @@ DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
 
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
+
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/subscription/t/100_bugs.pl b/src/test/subscription/t/100_bugs.pl
index cb36ca7b16..794b928f50 100644
--- a/src/test/subscription/t/100_bugs.pl
+++ b/src/test/subscription/t/100_bugs.pl
@@ -377,8 +377,8 @@ $node_publisher->safe_psql('postgres', "DROP PUBLICATION tap_pub_sch");
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');
 
-# The bug was that when the REPLICA IDENTITY FULL is used with dropped or
-# generated columns, we fail to apply updates and deletes
+# The bug was that when the REPLICA IDENTITY FULL is used with dropped
+# we fail to apply updates and deletes
 $node_publisher->rotate_logfile();
 $node_publisher->start();
 
@@ -389,18 +389,14 @@ $node_publisher->safe_psql(
 	'postgres', qq(
 	CREATE TABLE dropped_cols (a int, b_drop int, c int);
 	ALTER TABLE dropped_cols REPLICA IDENTITY FULL;
-	CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
-	ALTER TABLE generated_cols REPLICA IDENTITY FULL;
-	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols, generated_cols;
+	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols;
 	-- some initial data
 	INSERT INTO dropped_cols VALUES (1, 1, 1);
-	INSERT INTO generated_cols (a, c) VALUES (1, 1);
 ));
 
 $node_subscriber->safe_psql(
 	'postgres', qq(
 	 CREATE TABLE dropped_cols (a int, b_drop int, c int);
-	 CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
 ));
 
 $publisher_connstr = $node_publisher->connstr . ' dbname=postgres';
@@ -421,7 +417,6 @@ $node_subscriber->safe_psql(
 $node_publisher->safe_psql(
 	'postgres', qq(
 		UPDATE dropped_cols SET a = 100;
-		UPDATE generated_cols SET a = 100;
 ));
 $node_publisher->wait_for_catchup('sub_dropped_cols');
 
@@ -430,11 +425,6 @@ is( $node_subscriber->safe_psql(
 	qq(1),
 	'replication with RI FULL and dropped columns');
 
-is( $node_subscriber->safe_psql(
-		'postgres', "SELECT count(*) FROM generated_cols WHERE a = 100"),
-	qq(1),
-	'replication with RI FULL and generated columns');
-
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');
 
-- 
2.34.1

From 11bfcaa237a9c46d5c6b798d06c188d4e16025da Mon Sep 17 00:00:00 2001
From: Shlok Kyal <shlok.kyal.oss@gmail.com>
Date: Tue, 5 Nov 2024 11:11:25 +0530
Subject: [PATCH v15] Disallow UPDATE/DELETE on table with generated column as
 REPLICA IDENTITY

UPDATE or DELETE operations on tables with unpublished generated columns
set as REPLICA IDENTITY are not permitted. This patch ensures that if an
UPDATE or DELETE command is executed on such tables, an error will be
thrown.
With this patch, the behavior has changed for the test added in commit
adedf54. Additionally, there is no other way to trigger the bug that was
fixed by commit adedf54, so the test has been removed.
---
 doc/src/sgml/ref/create_publication.sgml  |   8 ++
 src/backend/commands/publicationcmds.c    | 156 +++++++++++++---------
 src/backend/executor/execReplication.c    |  37 ++++-
 src/backend/utils/cache/relcache.c        |  57 ++++++--
 src/include/catalog/pg_publication.h      |   7 +
 src/include/commands/publicationcmds.h    |   7 +-
 src/test/regress/expected/publication.out |  24 ++++
 src/test/regress/sql/publication.sql      |  25 ++++
 src/test/subscription/t/100_bugs.pl       |  16 +--
 9 files changed, 238 insertions(+), 99 deletions(-)

diff --git a/doc/src/sgml/ref/create_publication.sgml b/doc/src/sgml/ref/create_publication.sgml
index f8e217d661..5e25536554 100644
--- a/doc/src/sgml/ref/create_publication.sgml
+++ b/doc/src/sgml/ref/create_publication.sgml
@@ -311,6 +311,14 @@ CREATE PUBLICATION <replaceable class="parameter">name</replaceable>
    system columns.
   </para>
 
+  <para>
+   The generated columns that are part of <literal>REPLICA IDENTITY</literal>
+   must be published explicitly either by listing them in the column list or
+   by enabling the <literal>publish_generated_columns</literal> option, in
+   order for <command>UPDATE</command> and <command>DELETE</command> operations
+   to be published.
+  </para>
+
   <para>
    The row filter on a table becomes redundant if
    <literal>FOR TABLES IN SCHEMA</literal> is specified and the table
diff --git a/src/backend/commands/publicationcmds.c b/src/backend/commands/publicationcmds.c
index 031c84ec29..5050057a7e 100644
--- a/src/backend/commands/publicationcmds.c
+++ b/src/backend/commands/publicationcmds.c
@@ -336,21 +336,36 @@ pub_rf_contains_invalid_column(Oid pubid, Relation relation, List *ancestors,
 }
 
 /*
- * Check if all columns referenced in the REPLICA IDENTITY are covered by
- * the column list.
+ * Check for invalid columns in the publication table definition.
  *
- * Returns true if any replica identity column is not covered by column list.
+ * This function evaluates two conditions:
+ *
+ * 1. Ensures that all columns referenced in the REPLICA IDENTITY are covered
+ *    by the column list. If any column is missing, *invalid_column_list is set
+ *    to true.
+ * 2. Ensures that all the generated columns referenced in the REPLICA IDENTITY
+ *    are published either by listing them in the column list or by enabling
+ *    publish_generated_columns option. If any unpublished generated column is
+ *    found, *invalid_gen_col is set to true.
+ *
+ * Returns true if any of the above conditions are not met.
  */
 bool
-pub_collist_contains_invalid_column(Oid pubid, Relation relation, List *ancestors,
-									bool pubviaroot)
+pub_contains_invalid_column(Oid pubid, Relation relation, List *ancestors,
+							bool pubviaroot, bool pubgencols,
+							bool *invalid_column_list,
+							bool *invalid_gen_col)
 {
-	HeapTuple	tuple;
 	Oid			relid = RelationGetRelid(relation);
 	Oid			publish_as_relid = RelationGetRelid(relation);
-	bool		result = false;
-	Datum		datum;
-	bool		isnull;
+	Bitmapset  *idattrs;
+	Bitmapset  *columns = NULL;
+	TupleDesc	desc = RelationGetDescr(relation);
+	Publication *pub;
+	int			x;
+
+	*invalid_column_list = false;
+	*invalid_gen_col = false;
 
 	/*
 	 * For a partition, if pubviaroot is true, find the topmost ancestor that
@@ -368,80 +383,91 @@ pub_collist_contains_invalid_column(Oid pubid, Relation relation, List *ancestor
 			publish_as_relid = relid;
 	}
 
-	tuple = SearchSysCache2(PUBLICATIONRELMAP,
-							ObjectIdGetDatum(publish_as_relid),
-							ObjectIdGetDatum(pubid));
+	/* Fetch the column list */
+	pub = GetPublication(pubid);
+	check_and_fetch_column_list(pub, publish_as_relid, NULL, &columns);
 
-	if (!HeapTupleIsValid(tuple))
-		return false;
+	if (relation->rd_rel->relreplident == REPLICA_IDENTITY_FULL)
+	{
+		/* With REPLICA IDENTITY FULL, no column list is allowed. */
+		*invalid_column_list = (columns != NULL);
 
-	datum = SysCacheGetAttr(PUBLICATIONRELMAP, tuple,
-							Anum_pg_publication_rel_prattrs,
-							&isnull);
+		/*
+		 * As we don't allow a column list with REPLICA IDENTITY FULL, the
+		 * publish_generated_columns option must be set to true if the table
+		 * has any stored generated columns.
+		 */
+		if (!pubgencols &&
+			relation->rd_att->constr &&
+			relation->rd_att->constr->has_generated_stored)
+			*invalid_gen_col = true;
 
-	if (!isnull)
-	{
-		int			x;
-		Bitmapset  *idattrs;
-		Bitmapset  *columns = NULL;
+		if (*invalid_gen_col && *invalid_column_list)
+			return true;
+	}
 
-		/* With REPLICA IDENTITY FULL, no column list is allowed. */
-		if (relation->rd_rel->relreplident == REPLICA_IDENTITY_FULL)
-			result = true;
+	/* Remember columns that are part of the REPLICA IDENTITY */
+	idattrs = RelationGetIndexAttrBitmap(relation,
+										 INDEX_ATTR_BITMAP_IDENTITY_KEY);
 
-		/* Transform the column list datum to a bitmapset. */
-		columns = pub_collist_to_bitmapset(NULL, datum, NULL);
+	/*
+	 * Attnums in the bitmap returned by RelationGetIndexAttrBitmap are offset
+	 * (to handle system columns the usual way), while column list does not
+	 * use offset, so we can't do bms_is_subset(). Instead, we have to loop
+	 * over the idattrs and check all of them are in the list.
+	 */
+	x = -1;
+	while ((x = bms_next_member(idattrs, x)) >= 0)
+	{
+		AttrNumber	attnum = (x + FirstLowInvalidHeapAttributeNumber);
+		Form_pg_attribute att = TupleDescAttr(desc, attnum - 1);
 
-		/* Remember columns that are part of the REPLICA IDENTITY */
-		idattrs = RelationGetIndexAttrBitmap(relation,
-											 INDEX_ATTR_BITMAP_IDENTITY_KEY);
+		if (columns == NULL)
+		{
+			/*
+			 * The publish_generated_columns option must be set to true if the
+			 * REPLICA IDENTITY contains any stored generated column.
+			 */
+			if (!pubgencols && att->attgenerated)
+			{
+				*invalid_gen_col = true;
+				break;
+			}
+
+			/* Skip validating the column list since it is not defined */
+			continue;
+		}
 
 		/*
-		 * Attnums in the bitmap returned by RelationGetIndexAttrBitmap are
-		 * offset (to handle system columns the usual way), while column list
-		 * does not use offset, so we can't do bms_is_subset(). Instead, we
-		 * have to loop over the idattrs and check all of them are in the
-		 * list.
+		 * If pubviaroot is true, we are validating the column list of the
+		 * parent table, but the bitmap contains the replica identity
+		 * information of the child table. The parent/child attnums may not
+		 * match, so translate them to the parent - get the attname from the
+		 * child, and look it up in the parent.
 		 */
-		x = -1;
-		while ((x = bms_next_member(idattrs, x)) >= 0)
+		if (pubviaroot)
 		{
-			AttrNumber	attnum = (x + FirstLowInvalidHeapAttributeNumber);
+			/* attribute name in the child table */
+			char	   *colname = get_attname(relid, attnum, false);
 
 			/*
-			 * If pubviaroot is true, we are validating the column list of the
-			 * parent table, but the bitmap contains the replica identity
-			 * information of the child table. The parent/child attnums may
-			 * not match, so translate them to the parent - get the attname
-			 * from the child, and look it up in the parent.
+			 * Determine the attnum for the attribute name in parent (we are
+			 * using the column list defined on the parent).
 			 */
-			if (pubviaroot)
-			{
-				/* attribute name in the child table */
-				char	   *colname = get_attname(relid, attnum, false);
-
-				/*
-				 * Determine the attnum for the attribute name in parent (we
-				 * are using the column list defined on the parent).
-				 */
-				attnum = get_attnum(publish_as_relid, colname);
-			}
-
-			/* replica identity column, not covered by the column list */
-			if (!bms_is_member(attnum, columns))
-			{
-				result = true;
-				break;
-			}
+			attnum = get_attnum(publish_as_relid, colname);
 		}
 
-		bms_free(idattrs);
-		bms_free(columns);
+		/* replica identity column, not covered by the column list */
+		*invalid_column_list |= !bms_is_member(attnum, columns);
+
+		if (*invalid_column_list && *invalid_gen_col)
+			break;
 	}
 
-	ReleaseSysCache(tuple);
+	bms_free(columns);
+	bms_free(idattrs);
 
-	return result;
+	return *invalid_column_list || *invalid_gen_col;
 }
 
 /* check_functions_in_node callback */
diff --git a/src/backend/executor/execReplication.c b/src/backend/executor/execReplication.c
index 54025c9f15..7f90cd022b 100644
--- a/src/backend/executor/execReplication.c
+++ b/src/backend/executor/execReplication.c
@@ -785,15 +785,26 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 		return;
 
 	/*
-	 * 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 - i.e. when all referenced columns are part of REPLICA IDENTITY
-	 * or the table does not publish UPDATEs or DELETEs.
+	 * It is only safe to execute UPDATE/DELETE if the relation does not
+	 * publish UPDATEs or DELETEs, or all the following conditions are
+	 * satisfied:
+	 *
+	 * 1. All columns, referenced in the row filters from publications which
+	 *    the relation is in, are valid - i.e. when all referenced columns are
+	 *    part of REPLICA IDENTITY.
+	 *
+	 * 2. All columns, referenced in the column lists are valid - i.e. when all
+	 *    columns referenced in the REPLICA IDENTITY are covered by the column
+	 *    list.
+	 *
+	 * 3. All generated columns in REPLICA IDENTITY of the relation, are valid
+	 *    - i.e. when all these generated columns are published.
 	 *
 	 * XXX We could optimize it by first checking whether any of the
-	 * publications have a row filter for this relation. If not and relation
-	 * has replica identity then we can avoid building the descriptor but as
-	 * this happens only one time it doesn't seem worth the additional
+	 * publications have a row filter or column list for this relation, or if
+	 * the relation contains a generated column. If none of these exist and the
+	 * relation has replica identity then we can avoid building the descriptor
+	 * but as this happens only one time it doesn't seem worth the additional
 	 * complexity.
 	 */
 	RelationBuildPublicationDesc(rel, &pubdesc);
@@ -809,6 +820,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot update table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_UPDATE && !pubdesc.gencols_valid_for_update)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot update table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("Replica identity consists of an unpublished generated column.")));
 	else if (cmd == CMD_DELETE && !pubdesc.rf_valid_for_delete)
 		ereport(ERROR,
 				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
@@ -821,6 +838,12 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				 errmsg("cannot delete from table \"%s\"",
 						RelationGetRelationName(rel)),
 				 errdetail("Column list used by the publication does not cover the replica identity.")));
+	else if (cmd == CMD_DELETE && !pubdesc.gencols_valid_for_delete)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+				 errmsg("cannot delete from table \"%s\"",
+						RelationGetRelationName(rel)),
+				 errdetail("Replica identity consists of an unpublished generated column.")));
 
 	/* If relation has replica identity we are always good. */
 	if (OidIsValid(RelationGetReplicaIndex(rel)))
diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
index d0892cee24..422509f18d 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -5706,12 +5706,19 @@ RelationGetExclusionInfo(Relation indexRelation,
  * Get the publication information for the given relation.
  *
  * Traverse all the publications which the relation is in to get the
- * publication actions and validate the row filter expressions for such
- * publications if any. We consider the row filter expression as invalid if it
- * references any column which is not part of REPLICA IDENTITY.
+ * publication actions and validate:
+ * 1. The row filter expressions for such publications if any. We consider the
+ *    row filter expression as invalid if it references any column which is not
+ *    part of REPLICA IDENTITY.
+ * 2. The column list for such publication if any. We consider the column list
+ * 	  invalid if REPLICA IDENTITY contains any column that is not part of it.
+ * 3. The generated columns of the relation for such publications. We consider
+ *    any reference of an unpublished generated column in REPLICA IDENTITY as
+ *    invalid.
  *
  * To avoid fetching the publication information repeatedly, we cache the
- * publication actions and row filter validation information.
+ * publication actions, row filter validation information, column list
+ * validation information, and generated column validation information.
  */
 void
 RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
@@ -5734,6 +5741,8 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		pubdesc->rf_valid_for_delete = true;
 		pubdesc->cols_valid_for_update = true;
 		pubdesc->cols_valid_for_delete = true;
+		pubdesc->gencols_valid_for_update = true;
+		pubdesc->gencols_valid_for_delete = true;
 		return;
 	}
 
@@ -5748,6 +5757,8 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 	pubdesc->rf_valid_for_delete = true;
 	pubdesc->cols_valid_for_update = true;
 	pubdesc->cols_valid_for_delete = true;
+	pubdesc->gencols_valid_for_update = true;
+	pubdesc->gencols_valid_for_delete = true;
 
 	/* Fetch the publication membership info. */
 	puboids = GetRelationPublications(relid);
@@ -5777,6 +5788,8 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		Oid			pubid = lfirst_oid(lc);
 		HeapTuple	tup;
 		Form_pg_publication pubform;
+		bool		invalid_column_list;
+		bool		invalid_gen_col;
 
 		tup = SearchSysCache1(PUBLICATIONOID, ObjectIdGetDatum(pubid));
 
@@ -5811,18 +5824,27 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 		/*
 		 * Check if all columns are part of the REPLICA IDENTITY index or not.
 		 *
-		 * If the publication is FOR ALL TABLES then it means the table has no
-		 * column list and we can skip the validation.
+		 * Check if all generated columns included in the REPLICA IDENTITY are
+		 * published.
 		 */
-		if (!pubform->puballtables &&
-			(pubform->pubupdate || pubform->pubdelete) &&
-			pub_collist_contains_invalid_column(pubid, relation, ancestors,
-												pubform->pubviaroot))
+		if ((pubform->pubupdate || pubform->pubdelete) &&
+			pub_contains_invalid_column(pubid, relation, ancestors,
+										pubform->pubviaroot,
+										pubform->pubgencols,
+										&invalid_column_list,
+										&invalid_gen_col))
 		{
 			if (pubform->pubupdate)
-				pubdesc->cols_valid_for_update = false;
+			{
+				pubdesc->cols_valid_for_update = !invalid_column_list;
+				pubdesc->gencols_valid_for_update = !invalid_gen_col;
+			}
+
 			if (pubform->pubdelete)
-				pubdesc->cols_valid_for_delete = false;
+			{
+				pubdesc->cols_valid_for_delete = !invalid_column_list;
+				pubdesc->gencols_valid_for_delete = !invalid_gen_col;
+			}
 		}
 
 		ReleaseSysCache(tup);
@@ -5846,6 +5868,17 @@ RelationBuildPublicationDesc(Relation relation, PublicationDesc *pubdesc)
 			pubdesc->pubactions.pubdelete && pubdesc->pubactions.pubtruncate &&
 			!pubdesc->cols_valid_for_update && !pubdesc->cols_valid_for_delete)
 			break;
+
+		/*
+		 * If we know everything is replicated and replica identity has an
+		 * unpublished generated column, there is no point to check for other
+		 * publications.
+		 */
+		if (pubdesc->pubactions.pubinsert && pubdesc->pubactions.pubupdate &&
+			pubdesc->pubactions.pubdelete && pubdesc->pubactions.pubtruncate &&
+			!pubdesc->gencols_valid_for_update &&
+			!pubdesc->gencols_valid_for_delete)
+			break;
 	}
 
 	if (relation->rd_pubdesc)
diff --git a/src/include/catalog/pg_publication.h b/src/include/catalog/pg_publication.h
index 9a83a72d6b..e2d894a2ff 100644
--- a/src/include/catalog/pg_publication.h
+++ b/src/include/catalog/pg_publication.h
@@ -98,6 +98,13 @@ typedef struct PublicationDesc
 	 */
 	bool		cols_valid_for_update;
 	bool		cols_valid_for_delete;
+
+	/*
+	 * true if all generated columns that are part of replica identity are
+	 * published or the publication actions do not include UPDATE or DELETE.
+	 */
+	bool		gencols_valid_for_update;
+	bool		gencols_valid_for_delete;
 } PublicationDesc;
 
 typedef struct Publication
diff --git a/src/include/commands/publicationcmds.h b/src/include/commands/publicationcmds.h
index 5487c571f6..19037518e8 100644
--- a/src/include/commands/publicationcmds.h
+++ b/src/include/commands/publicationcmds.h
@@ -33,7 +33,10 @@ extern void AlterPublicationOwner_oid(Oid subid, Oid newOwnerId);
 extern void InvalidatePublicationRels(List *relids);
 extern bool pub_rf_contains_invalid_column(Oid pubid, Relation relation,
 										   List *ancestors, bool pubviaroot);
-extern bool pub_collist_contains_invalid_column(Oid pubid, Relation relation,
-												List *ancestors, bool pubviaroot);
+extern bool pub_contains_invalid_column(Oid pubid, Relation relation,
+										List *ancestors, bool pubviaroot,
+										bool pubgencols,
+										bool *invalid_column_list,
+										bool *invalid_gen_col);
 
 #endif							/* PUBLICATIONCMDS_H */
diff --git a/src/test/regress/expected/publication.out b/src/test/regress/expected/publication.out
index 5de2d64d01..12d0611a7c 100644
--- a/src/test/regress/expected/publication.out
+++ b/src/test/regress/expected/publication.out
@@ -672,6 +672,30 @@ DROP TABLE rf_tbl_abcd_pk;
 DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  Replica identity consists of an unpublished generated column.
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+ERROR:  cannot update table "testpub_gencol"
+DETAIL:  Replica identity consists of an unpublished generated column.
+DROP PUBLICATION pub_gencol;
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/regress/sql/publication.sql b/src/test/regress/sql/publication.sql
index 48e68bcca2..a29587b45d 100644
--- a/src/test/regress/sql/publication.sql
+++ b/src/test/regress/sql/publication.sql
@@ -396,6 +396,31 @@ DROP TABLE rf_tbl_abcd_nopk;
 DROP TABLE rf_tbl_abcd_part_pk;
 -- ======================================================
 
+-- ======================================================
+-- test with generated column
+SET client_min_messages = 'ERROR';
+CREATE TABLE testpub_gencol (a INT, b INT GENERATED ALWAYS AS (a + 1) STORED NOT NULL);
+CREATE UNIQUE INDEX testpub_gencol_idx ON testpub_gencol (b);
+ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
+
+-- error - generated column "b" is not published but part of index set as REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+
+-- error - generated column "b" is not published and REPLICA IDENTITY is set FULL
+ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+
+-- ok - generated column "b" is published and is part of REPLICA IDENTITY
+CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
+UPDATE testpub_gencol SET a = 100 WHERE a = 1;
+DROP PUBLICATION pub_gencol;
+
+DROP TABLE testpub_gencol;
+RESET client_min_messages;
+-- ======================================================
+
 -- fail - duplicate tables are not allowed if that table has any column lists
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_dups FOR TABLE testpub_tbl1 (a), testpub_tbl1 WITH (publish = 'insert');
diff --git a/src/test/subscription/t/100_bugs.pl b/src/test/subscription/t/100_bugs.pl
index cb36ca7b16..794b928f50 100644
--- a/src/test/subscription/t/100_bugs.pl
+++ b/src/test/subscription/t/100_bugs.pl
@@ -377,8 +377,8 @@ $node_publisher->safe_psql('postgres', "DROP PUBLICATION tap_pub_sch");
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');
 
-# The bug was that when the REPLICA IDENTITY FULL is used with dropped or
-# generated columns, we fail to apply updates and deletes
+# The bug was that when the REPLICA IDENTITY FULL is used with dropped
+# we fail to apply updates and deletes
 $node_publisher->rotate_logfile();
 $node_publisher->start();
 
@@ -389,18 +389,14 @@ $node_publisher->safe_psql(
 	'postgres', qq(
 	CREATE TABLE dropped_cols (a int, b_drop int, c int);
 	ALTER TABLE dropped_cols REPLICA IDENTITY FULL;
-	CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
-	ALTER TABLE generated_cols REPLICA IDENTITY FULL;
-	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols, generated_cols;
+	CREATE PUBLICATION pub_dropped_cols FOR TABLE dropped_cols;
 	-- some initial data
 	INSERT INTO dropped_cols VALUES (1, 1, 1);
-	INSERT INTO generated_cols (a, c) VALUES (1, 1);
 ));
 
 $node_subscriber->safe_psql(
 	'postgres', qq(
 	 CREATE TABLE dropped_cols (a int, b_drop int, c int);
-	 CREATE TABLE generated_cols (a int, b_gen int GENERATED ALWAYS AS (5 * a) STORED, c int);
 ));
 
 $publisher_connstr = $node_publisher->connstr . ' dbname=postgres';
@@ -421,7 +417,6 @@ $node_subscriber->safe_psql(
 $node_publisher->safe_psql(
 	'postgres', qq(
 		UPDATE dropped_cols SET a = 100;
-		UPDATE generated_cols SET a = 100;
 ));
 $node_publisher->wait_for_catchup('sub_dropped_cols');
 
@@ -430,11 +425,6 @@ is( $node_subscriber->safe_psql(
 	qq(1),
 	'replication with RI FULL and dropped columns');
 
-is( $node_subscriber->safe_psql(
-		'postgres', "SELECT count(*) FROM generated_cols WHERE a = 100"),
-	qq(1),
-	'replication with RI FULL and generated columns');
-
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');
 
-- 
2.28.0.windows.1

From d57017700a79935b387a4a0f70d89d6227e622de Mon Sep 17 00:00:00 2001
From: Shlok Kyal <shlok.kyal.oss@gmail.com>
Date: Fri, 6 Dec 2024 10:52:44 +0530
Subject: [PATCH v1] Improve error message introduced in commit 87ce27de696

The sentence "Replica identity consists of an unpublished generated
column." implies that the entire replica identity is made up of an
unpublished generated column and nothing else.

This patch fixes it with a better error message.
---
 src/backend/executor/execReplication.c    | 4 ++--
 src/test/regress/expected/publication.out | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/backend/executor/execReplication.c b/src/backend/executor/execReplication.c
index cfdf2eedf4..368a40337a 100644
--- a/src/backend/executor/execReplication.c
+++ b/src/backend/executor/execReplication.c
@@ -825,7 +825,7 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
 				 errmsg("cannot update table \"%s\"",
 						RelationGetRelationName(rel)),
-				 errdetail("Replica identity consists of an unpublished generated column.")));
+				 errdetail("Replica identity must not contain unpublished generated columns.")));
 	else if (cmd == CMD_DELETE && !pubdesc.rf_valid_for_delete)
 		ereport(ERROR,
 				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
@@ -843,7 +843,7 @@ CheckCmdReplicaIdentity(Relation rel, CmdType cmd)
 				(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
 				 errmsg("cannot delete from table \"%s\"",
 						RelationGetRelationName(rel)),
-				 errdetail("Replica identity consists of an unpublished generated column.")));
+				 errdetail("Replica identity must not contain unpublished generated columns.")));
 
 	/* If relation has replica identity we are always good. */
 	if (OidIsValid(RelationGetReplicaIndex(rel)))
diff --git a/src/test/regress/expected/publication.out b/src/test/regress/expected/publication.out
index b44ab007de..c48f11f293 100644
--- a/src/test/regress/expected/publication.out
+++ b/src/test/regress/expected/publication.out
@@ -683,13 +683,13 @@ ALTER TABLE testpub_gencol REPLICA IDENTITY USING index testpub_gencol_idx;
 CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol;
 UPDATE testpub_gencol SET a = 100 WHERE a = 1;
 ERROR:  cannot update table "testpub_gencol"
-DETAIL:  Replica identity consists of an unpublished generated column.
+DETAIL:  Replica identity must not contain unpublished generated columns.
 -- error - generated column "b" must be published explicitly as it is
 -- part of the REPLICA IDENTITY.
 ALTER TABLE testpub_gencol REPLICA IDENTITY FULL;
 UPDATE testpub_gencol SET a = 100 WHERE a = 1;
 ERROR:  cannot update table "testpub_gencol"
-DETAIL:  Replica identity consists of an unpublished generated column.
+DETAIL:  Replica identity must not contain unpublished generated columns.
 DROP PUBLICATION pub_gencol;
 -- ok - generated column "b" is published explicitly
 CREATE PUBLICATION pub_gencol FOR TABLE testpub_gencol with (publish_generated_columns = true);
-- 
2.34.1