v43-0001-Support-logical-replication-of-generated-columns.patch

From c83e216a05a3fe008eff32cd8c6d96ecee7a7832 Mon Sep 17 00:00:00 2001
From: Vignesh C <vignesh21@gmail.com>
Date: Fri, 25 Oct 2024 15:23:03 +0530
Subject: [PATCH v43 1/2] Support logical replication of generated columns in
 column list.

Allow logical replication to publish generated columns if they are explicitly
mentioned in the column list.
---
 doc/src/sgml/protocol.sgml                  |  4 +-
 src/backend/catalog/pg_publication.c        | 10 +----
 src/backend/replication/logical/proto.c     | 41 +++++++++++++--------
 src/backend/replication/pgoutput/pgoutput.c | 26 ++++++++++---
 src/test/regress/expected/publication.out   | 21 +++++------
 src/test/regress/sql/publication.sql        |  2 +-
 src/test/subscription/t/031_column_list.pl  | 34 ++++++++++++++++-
 7 files changed, 93 insertions(+), 45 deletions(-)

diff --git a/doc/src/sgml/protocol.sgml b/doc/src/sgml/protocol.sgml
index 057c46f3f5..71b6b2a535 100644
--- a/doc/src/sgml/protocol.sgml
+++ b/doc/src/sgml/protocol.sgml
@@ -6544,7 +6544,7 @@ psql "dbname=postgres replication=database" -c "IDENTIFY_SYSTEM;"
 
      <para>
       Next, the following message part appears for each column included in
-      the publication (except generated columns):
+      the publication:
      </para>
 
      <variablelist>
@@ -7477,7 +7477,7 @@ psql "dbname=postgres replication=database" -c "IDENTIFY_SYSTEM;"
      </variablelist>
 
      <para>
-      Next, one of the following submessages appears for each column (except generated columns):
+      Next, one of the following submessages appears for each column:
 
       <variablelist>
        <varlistentry>
diff --git a/src/backend/catalog/pg_publication.c b/src/backend/catalog/pg_publication.c
index 7e5e357fd9..17a6093d06 100644
--- a/src/backend/catalog/pg_publication.c
+++ b/src/backend/catalog/pg_publication.c
@@ -500,8 +500,7 @@ publication_add_relation(Oid pubid, PublicationRelInfo *pri,
  * pub_collist_validate
  *		Process and validate the 'columns' list and ensure the columns are all
  *		valid to use for a publication.  Checks for and raises an ERROR for
- * 		any; unknown columns, system columns, duplicate columns or generated
- *		columns.
+ * 		any unknown columns, system columns, or duplicate columns.
  *
  * Looks up each column's attnum and returns a 0-based Bitmapset of the
  * corresponding attnums.
@@ -511,7 +510,6 @@ pub_collist_validate(Relation targetrel, List *columns)
 {
 	Bitmapset  *set = NULL;
 	ListCell   *lc;
-	TupleDesc	tupdesc = RelationGetDescr(targetrel);
 
 	foreach(lc, columns)
 	{
@@ -530,12 +528,6 @@ pub_collist_validate(Relation targetrel, List *columns)
 					errmsg("cannot use system column \"%s\" in publication column list",
 						   colname));
 
-		if (TupleDescAttr(tupdesc, attnum - 1)->attgenerated)
-			ereport(ERROR,
-					errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
-					errmsg("cannot use generated column \"%s\" in publication column list",
-						   colname));
-
 		if (bms_is_member(attnum, set))
 			ereport(ERROR,
 					errcode(ERRCODE_DUPLICATE_OBJECT),
diff --git a/src/backend/replication/logical/proto.c b/src/backend/replication/logical/proto.c
index 980f6e2741..cfc810a71a 100644
--- a/src/backend/replication/logical/proto.c
+++ b/src/backend/replication/logical/proto.c
@@ -52,6 +52,27 @@ column_in_column_list(int attnum, Bitmapset *columns)
 	return (columns == NULL || bms_is_member(attnum, columns));
 }
 
+/*
+ * Check if the column should be published.
+ */
+static bool
+should_publish_column(Form_pg_attribute att, Bitmapset *columns)
+{
+	if (att->attisdropped)
+		return false;
+
+	/*
+	 * Skip publishing generated columns if they are not included in the
+	 * column list.
+	 */
+	if (att->attgenerated && !columns)
+		return false;
+
+	if (!column_in_column_list(att->attnum, columns))
+		return false;
+
+	return true;
+}
 
 /*
  * Write BEGIN to the output stream.
@@ -781,10 +802,7 @@ logicalrep_write_tuple(StringInfo out, Relation rel, TupleTableSlot *slot,
 	{
 		Form_pg_attribute att = TupleDescAttr(desc, i);
 
-		if (att->attisdropped || att->attgenerated)
-			continue;
-
-		if (!column_in_column_list(att->attnum, columns))
+		if (!should_publish_column(att, columns))
 			continue;
 
 		nliveatts++;
@@ -802,10 +820,7 @@ logicalrep_write_tuple(StringInfo out, Relation rel, TupleTableSlot *slot,
 		Form_pg_type typclass;
 		Form_pg_attribute att = TupleDescAttr(desc, i);
 
-		if (att->attisdropped || att->attgenerated)
-			continue;
-
-		if (!column_in_column_list(att->attnum, columns))
+		if (!should_publish_column(att, columns))
 			continue;
 
 		if (isnull[i])
@@ -938,10 +953,7 @@ logicalrep_write_attrs(StringInfo out, Relation rel, Bitmapset *columns)
 	{
 		Form_pg_attribute att = TupleDescAttr(desc, i);
 
-		if (att->attisdropped || att->attgenerated)
-			continue;
-
-		if (!column_in_column_list(att->attnum, columns))
+		if (!should_publish_column(att, columns))
 			continue;
 
 		nliveatts++;
@@ -959,10 +971,7 @@ logicalrep_write_attrs(StringInfo out, Relation rel, Bitmapset *columns)
 		Form_pg_attribute att = TupleDescAttr(desc, i);
 		uint8		flags = 0;
 
-		if (att->attisdropped || att->attgenerated)
-			continue;
-
-		if (!column_in_column_list(att->attnum, columns))
+		if (!should_publish_column(att, columns))
 			continue;
 
 		/* REPLICA IDENTITY FULL means all columns are sent as part of key. */
diff --git a/src/backend/replication/pgoutput/pgoutput.c b/src/backend/replication/pgoutput/pgoutput.c
index 00e7024563..d59a8f5032 100644
--- a/src/backend/replication/pgoutput/pgoutput.c
+++ b/src/backend/replication/pgoutput/pgoutput.c
@@ -766,12 +766,19 @@ send_relation_and_attrs(Relation relation, TransactionId xid,
 	{
 		Form_pg_attribute att = TupleDescAttr(desc, i);
 
-		if (att->attisdropped || att->attgenerated)
+		if (att->attisdropped)
 			continue;
 
 		if (att->atttypid < FirstGenbkiObjectId)
 			continue;
 
+		/*
+		 * Skip publishing generated columns if they are not included in the
+		 * column list.
+		 */
+		if (att->attgenerated && !columns)
+			continue;
+
 		/* Skip this attribute if it's not present in the column list */
 		if (columns != NULL && !bms_is_member(att->attnum, columns))
 			continue;
@@ -1074,6 +1081,7 @@ pgoutput_column_list_init(PGOutputData *data, List *publications,
 					int			i;
 					int			nliveatts = 0;
 					TupleDesc	desc = RelationGetDescr(relation);
+					bool		gencolpresent = false;
 
 					pgoutput_ensure_entry_cxt(data, entry);
 
@@ -1085,17 +1093,25 @@ pgoutput_column_list_init(PGOutputData *data, List *publications,
 					{
 						Form_pg_attribute att = TupleDescAttr(desc, i);
 
-						if (att->attisdropped || att->attgenerated)
+						if (att->attisdropped)
 							continue;
 
+						if (att->attgenerated)
+						{
+							if (bms_is_member(att->attnum, cols))
+								gencolpresent = true;
+
+							continue;
+						}
+
 						nliveatts++;
 					}
 
 					/*
-					 * If column list includes all the columns of the table,
-					 * set it to NULL.
+					 * If column list includes all the columns of the table
+					 * and there are no generated columns, set it to NULL.
 					 */
-					if (bms_num_members(cols) == nliveatts)
+					if (bms_num_members(cols) == nliveatts && !gencolpresent)
 					{
 						bms_free(cols);
 						cols = NULL;
diff --git a/src/test/regress/expected/publication.out b/src/test/regress/expected/publication.out
index 660245ed0c..c248c2d717 100644
--- a/src/test/regress/expected/publication.out
+++ b/src/test/regress/expected/publication.out
@@ -687,43 +687,42 @@ UPDATE testpub_tbl5 SET a = 1;
 ERROR:  cannot update table "testpub_tbl5"
 DETAIL:  Column list used by the publication does not cover the replica identity.
 ALTER PUBLICATION testpub_fortable DROP TABLE testpub_tbl5;
--- error: generated column "d" can't be in list
+-- ok: generated column "d" can be in the list too
 ALTER PUBLICATION testpub_fortable ADD TABLE testpub_tbl5 (a, d);
-ERROR:  cannot use generated column "d" in publication column list
 -- error: system attributes "ctid" not allowed in column list
 ALTER PUBLICATION testpub_fortable ADD TABLE testpub_tbl5 (a, ctid);
-ERROR:  cannot use system column "ctid" in publication column list
+ERROR:  relation "testpub_tbl5" is already member of publication "testpub_fortable"
 ALTER PUBLICATION testpub_fortable SET TABLE testpub_tbl1 (id, ctid);
 ERROR:  cannot use system column "ctid" in publication column list
 -- error: duplicates not allowed in column list
 ALTER PUBLICATION testpub_fortable ADD TABLE testpub_tbl5 (a, a);
-ERROR:  duplicate column "a" in publication column list
+ERROR:  relation "testpub_tbl5" is already member of publication "testpub_fortable"
 ALTER PUBLICATION testpub_fortable SET TABLE testpub_tbl5 (a, a);
 ERROR:  duplicate column "a" in publication column list
 -- ok
 ALTER PUBLICATION testpub_fortable ADD TABLE testpub_tbl5 (a, c);
+ERROR:  relation "testpub_tbl5" is already member of publication "testpub_fortable"
 ALTER TABLE testpub_tbl5 DROP COLUMN c;		-- no dice
-ERROR:  cannot drop column c of table testpub_tbl5 because other objects depend on it
-DETAIL:  publication of table testpub_tbl5 in publication testpub_fortable depends on column c of table testpub_tbl5
-HINT:  Use DROP ... CASCADE to drop the dependent objects too.
 -- ok: for insert-only publication, any column list is acceptable
 ALTER PUBLICATION testpub_fortable_insert ADD TABLE testpub_tbl5 (b, c);
+ERROR:  column "c" of relation "testpub_tbl5" does not exist
 /* not all replica identities are good enough */
 CREATE UNIQUE INDEX testpub_tbl5_b_key ON testpub_tbl5 (b, c);
+ERROR:  column "c" does not exist
 ALTER TABLE testpub_tbl5 ALTER b SET NOT NULL, ALTER c SET NOT NULL;
+ERROR:  column "c" of relation "testpub_tbl5" does not exist
 ALTER TABLE testpub_tbl5 REPLICA IDENTITY USING INDEX testpub_tbl5_b_key;
+ERROR:  index "testpub_tbl5_b_key" for table "testpub_tbl5" does not exist
 -- error: replica identity (b,c) is not covered by column list (a, c)
 UPDATE testpub_tbl5 SET a = 1;
-ERROR:  cannot update table "testpub_tbl5"
-DETAIL:  Column list used by the publication does not cover the replica identity.
 ALTER PUBLICATION testpub_fortable DROP TABLE testpub_tbl5;
 -- error: change the replica identity to "b", and column list to (a, c)
 -- then update fails, because (a, c) does not cover replica identity
 ALTER TABLE testpub_tbl5 REPLICA IDENTITY USING INDEX testpub_tbl5_b_key;
+ERROR:  index "testpub_tbl5_b_key" for table "testpub_tbl5" does not exist
 ALTER PUBLICATION testpub_fortable ADD TABLE testpub_tbl5 (a, c);
+ERROR:  column "c" of relation "testpub_tbl5" does not exist
 UPDATE testpub_tbl5 SET a = 1;
-ERROR:  cannot update table "testpub_tbl5"
-DETAIL:  Column list used by the publication does not cover the replica identity.
 /* But if upd/del are not published, it works OK */
 SET client_min_messages = 'ERROR';
 CREATE PUBLICATION testpub_table_ins WITH (publish = 'insert, truncate');
diff --git a/src/test/regress/sql/publication.sql b/src/test/regress/sql/publication.sql
index f68a5b5986..9feb8442f2 100644
--- a/src/test/regress/sql/publication.sql
+++ b/src/test/regress/sql/publication.sql
@@ -413,7 +413,7 @@ ALTER PUBLICATION testpub_fortable ADD TABLE testpub_tbl5 (a, x);
 ALTER PUBLICATION testpub_fortable ADD TABLE testpub_tbl5 (b, c);
 UPDATE testpub_tbl5 SET a = 1;
 ALTER PUBLICATION testpub_fortable DROP TABLE testpub_tbl5;
--- error: generated column "d" can't be in list
+-- ok: generated column "d" can be in the list too
 ALTER PUBLICATION testpub_fortable ADD TABLE testpub_tbl5 (a, d);
 -- error: system attributes "ctid" not allowed in column list
 ALTER PUBLICATION testpub_fortable ADD TABLE testpub_tbl5 (a, ctid);
diff --git a/src/test/subscription/t/031_column_list.pl b/src/test/subscription/t/031_column_list.pl
index 9a97fa5020..a0dc9bf3e6 100644
--- a/src/test/subscription/t/031_column_list.pl
+++ b/src/test/subscription/t/031_column_list.pl
@@ -1202,7 +1202,7 @@ $result = $node_publisher->safe_psql(
 is( $result, qq(t
 t), 'check the number of columns in the old tuple');
 
-# TEST: Generated and dropped columns are not considered for the column list.
+# TEST: Dropped columns are not considered for the column list.
 # So, the publication having a column list except for those columns and a
 # publication without any column (aka all columns as part of the columns
 # list) are considered to have the same column list.
@@ -1275,6 +1275,38 @@ ok( $stderr =~
 	  qr/cannot use different column lists for table "public.test_mix_1" in different publications/,
 	'different column lists detected');
 
+# TEST: Generated columns are considered for the column list.
+$node_publisher->safe_psql(
+	'postgres', qq(
+	CREATE TABLE test_gen (a int PRIMARY KEY, b int GENERATED ALWAYS AS (a + 1) STORED);
+
+	CREATE PUBLICATION pub_gen FOR TABLE test_gen (a, b);
+));
+
+$node_subscriber->safe_psql(
+	'postgres', qq(
+	CREATE TABLE test_gen (a int PRIMARY KEY, b int);
+));
+
+$node_subscriber->safe_psql(
+	'postgres', qq(
+	CREATE SUBSCRIPTION sub_gen CONNECTION '$publisher_connstr' PUBLICATION pub_gen;
+));
+
+$node_subscriber->wait_for_subscription_sync;
+
+$node_publisher->safe_psql(
+	'postgres', qq(
+	INSERT INTO test_gen VALUES (1);
+));
+
+$node_publisher->wait_for_catchup('sub_gen');
+
+is( $node_subscriber->safe_psql(
+		'postgres', "SELECT * FROM test_gen ORDER BY a"),
+	qq(1|2),
+	'replication with generated columns in column list');
+
 # TEST: If the column list is changed after creating the subscription, we
 # should catch the error reported by walsender.
 
-- 
2.41.0.windows.3