v36-0002-Support-replication-of-generated-column-during-i.patch

From 3962cc7928453df966f1bb1969f70a1194e627aa Mon Sep 17 00:00:00 2001
From: Shubham Khanna <shubham.khanna@fujitsu.com>
Date: Tue, 24 Sep 2024 14:44:44 +0530
Subject: [PATCH v36 2/3] Support replication of generated column during
 initial sync

When 'copy_data' is true, during the initial sync, the data is replicated from
the publisher to the subscriber using the COPY command. The normal COPY
command does not copy generated columns, so when 'publish_generated_columns'
is true, we need to copy using the syntax:
'COPY (SELECT column_name FROM table_name) TO STDOUT'.

Here 'publish_generated_columns' is a PUBLICATION parameter and
'copy_data' is a SUBSCRIPTION parameter.

Summary:

when (publish_generated_columns = true)

* publisher not-generated column => subscriber not-generated column:
This is just normal logical replication (not changed by this patch).

* publisher not-generated column => subscriber generated column: This
will give ERROR.

* publisher generated column => subscriber not-generated column: The
publisher generated column value is copied.

* publisher generated column => subscriber generated column: This
will give ERROR.

when (publish_generated_columns = false)

* publisher not-generated column => subscriber not-generated column:
This is just normal logical replication (not changed by this patch).

* publisher not-generated column => subscriber generated column: This
will give ERROR.

* publisher generated column => subscriber not-generated column:
Publisher generated column is not replicated. The subscriber column
will be filled with the subscriber-side default data.

* publisher generated column => subscriber generated column: Publisher
generated column is not replicated. The subscriber generated column
will be filed with the subscriber-side computed or default data.
---
 src/backend/catalog/pg_subscription.c       |  31 +++
 src/backend/commands/subscriptioncmds.c     |  31 ---
 src/backend/replication/logical/relation.c  |   2 +-
 src/backend/replication/logical/tablesync.c | 197 ++++++++++++++++----
 src/include/catalog/pg_subscription.h       |   4 +
 src/include/replication/logicalrelation.h   |   3 +-
 6 files changed, 199 insertions(+), 69 deletions(-)

diff --git a/src/backend/catalog/pg_subscription.c b/src/backend/catalog/pg_subscription.c
index 9efc9159f2..fcfbf86c0b 100644
--- a/src/backend/catalog/pg_subscription.c
+++ b/src/backend/catalog/pg_subscription.c
@@ -551,3 +551,34 @@ GetSubscriptionRelations(Oid subid, bool not_ready)
 
 	return res;
 }
+
+/*
+ * Add publication names from the list to a string.
+ */
+void
+get_publications_str(List *publications, StringInfo dest, bool quote_literal)
+{
+	ListCell   *lc;
+	bool		first = true;
+
+	Assert(publications != NIL);
+
+	foreach(lc, publications)
+	{
+		char	   *pubname = strVal(lfirst(lc));
+
+		if (first)
+			first = false;
+		else
+			appendStringInfoString(dest, ", ");
+
+		if (quote_literal)
+			appendStringInfoString(dest, quote_literal_cstr(pubname));
+		else
+		{
+			appendStringInfoChar(dest, '"');
+			appendStringInfoString(dest, pubname);
+			appendStringInfoChar(dest, '"');
+		}
+	}
+}
diff --git a/src/backend/commands/subscriptioncmds.c b/src/backend/commands/subscriptioncmds.c
index 02ccc636b8..addf307cb6 100644
--- a/src/backend/commands/subscriptioncmds.c
+++ b/src/backend/commands/subscriptioncmds.c
@@ -439,37 +439,6 @@ parse_subscription_options(ParseState *pstate, List *stmt_options,
 	}
 }
 
-/*
- * Add publication names from the list to a string.
- */
-static void
-get_publications_str(List *publications, StringInfo dest, bool quote_literal)
-{
-	ListCell   *lc;
-	bool		first = true;
-
-	Assert(publications != NIL);
-
-	foreach(lc, publications)
-	{
-		char	   *pubname = strVal(lfirst(lc));
-
-		if (first)
-			first = false;
-		else
-			appendStringInfoString(dest, ", ");
-
-		if (quote_literal)
-			appendStringInfoString(dest, quote_literal_cstr(pubname));
-		else
-		{
-			appendStringInfoChar(dest, '"');
-			appendStringInfoString(dest, pubname);
-			appendStringInfoChar(dest, '"');
-		}
-	}
-}
-
 /*
  * Check that the specified publications are present on the publisher.
  */
diff --git a/src/backend/replication/logical/relation.c b/src/backend/replication/logical/relation.c
index f139e7b01e..338b083696 100644
--- a/src/backend/replication/logical/relation.c
+++ b/src/backend/replication/logical/relation.c
@@ -205,7 +205,7 @@ logicalrep_relmap_update(LogicalRepRelation *remoterel)
  *
  * Returns -1 if not found.
  */
-static int
+int
 logicalrep_rel_att_by_name(LogicalRepRelation *remoterel, const char *attname)
 {
 	int			i;
diff --git a/src/backend/replication/logical/tablesync.c b/src/backend/replication/logical/tablesync.c
index e03e761392..22ebe40336 100644
--- a/src/backend/replication/logical/tablesync.c
+++ b/src/backend/replication/logical/tablesync.c
@@ -118,6 +118,7 @@
 #include "utils/builtins.h"
 #include "utils/lsyscache.h"
 #include "utils/memutils.h"
+#include "utils/rel.h"
 #include "utils/rls.h"
 #include "utils/snapmgr.h"
 #include "utils/syscache.h"
@@ -693,20 +694,72 @@ process_syncing_tables(XLogRecPtr current_lsn)
 
 /*
  * Create list of columns for COPY based on logical relation mapping.
+ * Exclude columns that are subscription table generated columns.
  */
 static List *
-make_copy_attnamelist(LogicalRepRelMapEntry *rel)
+make_copy_attnamelist(LogicalRepRelMapEntry *rel, bool *remotegenlist)
 {
 	List	   *attnamelist = NIL;
-	int			i;
+	bool	   *localgenlist;
+	TupleDesc	desc;
 
-	for (i = 0; i < rel->remoterel.natts; i++)
+	desc = RelationGetDescr(rel->localrel);
+
+	/*
+	 * localgenlist stores if a generated column on remoterel has a matching
+	 * name corresponding to a generated column on localrel.
+	 */
+	localgenlist = palloc0(rel->remoterel.natts * sizeof(bool));
+
+	/*
+	 * This loop checks for generated columns of the subscription table.
+	 */
+	for (int i = 0; i < desc->natts; i++)
 	{
-		attnamelist = lappend(attnamelist,
-							  makeString(rel->remoterel.attnames[i]));
+		int			remote_attnum;
+		Form_pg_attribute attr = TupleDescAttr(desc, i);
+
+		if (!attr->attgenerated)
+			continue;
+
+		remote_attnum = logicalrep_rel_att_by_name(&rel->remoterel,
+												   NameStr(attr->attname));
+
+		if (remote_attnum >= 0)
+		{
+			/*
+			 * Check if the subscription table generated column has same name
+			 * as a non-generated column in the corresponding publication
+			 * table.
+			 */
+			if (!remotegenlist[remote_attnum])
+				ereport(ERROR,
+						(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+						 errmsg("logical replication target relation \"%s.%s\" has a generated column \"%s\" "
+								"but corresponding column on source relation is not a generated column",
+								rel->remoterel.nspname, rel->remoterel.relname, NameStr(attr->attname))));
+
+			/*
+			 * 'localgenlist' records that this is a generated column in the
+			 * subscription table. Later, we use this information to skip
+			 * adding this column to the column list for COPY.
+			 */
+			localgenlist[remote_attnum] = true;
+		}
 	}
 
+	/*
+	 * Construct column list for COPY, excluding columns that are subscription
+	 * table generated columns.
+	 */
+	for (int i = 0; i < rel->remoterel.natts; i++)
+	{
+		if (!localgenlist[i])
+			attnamelist = lappend(attnamelist,
+								  makeString(rel->remoterel.attnames[i]));
+	}
 
+	pfree(localgenlist);
 	return attnamelist;
 }
 
@@ -791,19 +844,21 @@ copy_read_data(void *outbuf, int minread, int maxread)
  * qualifications to be used in the COPY command.
  */
 static void
-fetch_remote_table_info(char *nspname, char *relname,
+fetch_remote_table_info(char *nspname, char *relname, bool **remotegenlist_res,
 						LogicalRepRelation *lrel, List **qual)
 {
 	WalRcvExecResult *res;
 	StringInfoData cmd;
 	TupleTableSlot *slot;
 	Oid			tableRow[] = {OIDOID, CHAROID, CHAROID};
-	Oid			attrRow[] = {INT2OID, TEXTOID, OIDOID, BOOLOID};
+	Oid			attrRow[] = {INT2OID, TEXTOID, OIDOID, BOOLOID, BOOLOID};
 	Oid			qualRow[] = {TEXTOID};
 	bool		isnull;
+	bool	   *remotegenlist;
+	bool		has_pub_with_pubgencols = false;
 	int			natt;
-	ListCell   *lc;
 	Bitmapset  *included_cols = NULL;
+	int			server_version = walrcv_server_version(LogRepWorkerWalRcvConn);
 
 	lrel->nspname = nspname;
 	lrel->relname = relname;
@@ -846,30 +901,24 @@ fetch_remote_table_info(char *nspname, char *relname,
 
 
 	/*
-	 * Get column lists for each relation.
+	 * Get column lists for each relation, and check if any of the
+	 * publications have the 'publish_generated_columns' parameter enabled.
 	 *
 	 * We need to do this before fetching info about column names and types,
 	 * so that we can skip columns that should not be replicated.
 	 */
-	if (walrcv_server_version(LogRepWorkerWalRcvConn) >= 150000)
+	if (server_version >= 150000)
 	{
 		WalRcvExecResult *pubres;
 		TupleTableSlot *tslot;
 		Oid			attrsRow[] = {INT2VECTOROID};
-		StringInfoData pub_names;
-
-		initStringInfo(&pub_names);
-		foreach(lc, MySubscription->publications)
-		{
-			if (foreach_current_index(lc) > 0)
-				appendStringInfoString(&pub_names, ", ");
-			appendStringInfoString(&pub_names, quote_literal_cstr(strVal(lfirst(lc))));
-		}
+		StringInfo	pub_names = makeStringInfo();
 
 		/*
 		 * Fetch info about column lists for the relation (from all the
 		 * publications).
 		 */
+		get_publications_str(MySubscription->publications, pub_names, true);
 		resetStringInfo(&cmd);
 		appendStringInfo(&cmd,
 						 "SELECT DISTINCT"
@@ -881,7 +930,7 @@ fetch_remote_table_info(char *nspname, char *relname,
 						 " WHERE gpt.relid = %u AND c.oid = gpt.relid"
 						 "   AND p.pubname IN ( %s )",
 						 lrel->remoteid,
-						 pub_names.data);
+						 pub_names->data);
 
 		pubres = walrcv_exec(LogRepWorkerWalRcvConn, cmd.data,
 							 lengthof(attrsRow), attrsRow);
@@ -937,7 +986,44 @@ fetch_remote_table_info(char *nspname, char *relname,
 
 		walrcv_clear_result(pubres);
 
-		pfree(pub_names.data);
+		/*
+		 * Check if any of the publications have the
+		 * 'publish_generated_columns' parameter enabled.
+		 */
+		if (server_version >= 180000)
+		{
+			WalRcvExecResult *gencolres;
+			Oid			gencolsRow[] = {BOOLOID};
+
+			resetStringInfo(&cmd);
+			appendStringInfo(&cmd,
+							 "SELECT count(*) > 0 FROM pg_catalog.pg_publication "
+							 "WHERE pubname IN ( %s ) AND pubgencols = 't'",
+							 pub_names->data);
+
+			gencolres = walrcv_exec(LogRepWorkerWalRcvConn, cmd.data,
+									lengthof(gencolsRow), gencolsRow);
+			if (gencolres->status != WALRCV_OK_TUPLES)
+				ereport(ERROR,
+						errcode(ERRCODE_CONNECTION_FAILURE),
+						errmsg("could not fetch generated column publication information from publication list: %s",
+							   pub_names->data));
+
+			tslot = MakeSingleTupleTableSlot(gencolres->tupledesc, &TTSOpsMinimalTuple);
+			if (!tuplestore_gettupleslot(gencolres->tuplestore, true, false, tslot))
+				ereport(ERROR,
+						errcode(ERRCODE_UNDEFINED_OBJECT),
+						errmsg("failed to fetch tuple for generated column publication information from publication list: %s",
+							   pub_names->data));
+
+			has_pub_with_pubgencols = DatumGetBool(slot_getattr(tslot, 1, &isnull));
+			Assert(!isnull);
+
+			ExecClearTuple(tslot);
+			walrcv_clear_result(gencolres);
+		}
+
+		pfree(pub_names->data);
 	}
 
 	/*
@@ -948,20 +1034,28 @@ fetch_remote_table_info(char *nspname, char *relname,
 					 "SELECT a.attnum,"
 					 "       a.attname,"
 					 "       a.atttypid,"
-					 "       a.attnum = ANY(i.indkey)"
+					 "       a.attnum = ANY(i.indkey)");
+
+	if (server_version >= 180000)
+		appendStringInfo(&cmd, ", a.attgenerated != ''");
+
+	appendStringInfo(&cmd,
 					 "  FROM pg_catalog.pg_attribute a"
 					 "  LEFT JOIN pg_catalog.pg_index i"
 					 "       ON (i.indexrelid = pg_get_replica_identity_index(%u))"
 					 " WHERE a.attnum > 0::pg_catalog.int2"
-					 "   AND NOT a.attisdropped %s"
+					 "   AND NOT a.attisdropped", lrel->remoteid);
+
+	if (!has_pub_with_pubgencols)
+		appendStringInfo(&cmd, " AND a.attgenerated = ''");
+
+	appendStringInfo(&cmd,
 					 "   AND a.attrelid = %u"
 					 " ORDER BY a.attnum",
-					 lrel->remoteid,
-					 (walrcv_server_version(LogRepWorkerWalRcvConn) >= 120000 ?
-					  "AND a.attgenerated = ''" : ""),
 					 lrel->remoteid);
+
 	res = walrcv_exec(LogRepWorkerWalRcvConn, cmd.data,
-					  lengthof(attrRow), attrRow);
+					  server_version >= 180000 ? lengthof(attrRow) : lengthof(attrRow) - 1, attrRow);
 
 	if (res->status != WALRCV_OK_TUPLES)
 		ereport(ERROR,
@@ -973,6 +1067,7 @@ fetch_remote_table_info(char *nspname, char *relname,
 	lrel->attnames = palloc0(MaxTupleAttributeNumber * sizeof(char *));
 	lrel->atttyps = palloc0(MaxTupleAttributeNumber * sizeof(Oid));
 	lrel->attkeys = NULL;
+	remotegenlist = palloc0(MaxTupleAttributeNumber * sizeof(bool));
 
 	/*
 	 * Store the columns as a list of names.  Ignore those that are not
@@ -1005,6 +1100,9 @@ fetch_remote_table_info(char *nspname, char *relname,
 		if (DatumGetBool(slot_getattr(slot, 4, &isnull)))
 			lrel->attkeys = bms_add_member(lrel->attkeys, natt);
 
+		if (server_version >= 120000)
+			remotegenlist[natt] = DatumGetBool(slot_getattr(slot, 5, &isnull));
+
 		/* Should never happen. */
 		if (++natt >= MaxTupleAttributeNumber)
 			elog(ERROR, "too many columns in remote table \"%s.%s\"",
@@ -1015,7 +1113,7 @@ fetch_remote_table_info(char *nspname, char *relname,
 	ExecDropSingleTupleTableSlot(slot);
 
 	lrel->natts = natt;
-
+	*remotegenlist_res = remotegenlist;
 	walrcv_clear_result(res);
 
 	/*
@@ -1037,7 +1135,7 @@ fetch_remote_table_info(char *nspname, char *relname,
 	 * 3) one of the subscribed publications is declared as TABLES IN SCHEMA
 	 * that includes this relation
 	 */
-	if (walrcv_server_version(LogRepWorkerWalRcvConn) >= 150000)
+	if (server_version >= 150000)
 	{
 		StringInfoData pub_names;
 
@@ -1123,10 +1221,13 @@ copy_table(Relation rel)
 	List	   *attnamelist;
 	ParseState *pstate;
 	List	   *options = NIL;
+	bool	   *remotegenlist;
+	bool		gencol_copy_needed = false;
 
 	/* Get the publisher relation info. */
 	fetch_remote_table_info(get_namespace_name(RelationGetNamespace(rel)),
-							RelationGetRelationName(rel), &lrel, &qual);
+							RelationGetRelationName(rel), &remotegenlist,
+							&lrel, &qual);
 
 	/* Put the relation into relmap. */
 	logicalrep_relmap_update(&lrel);
@@ -1135,11 +1236,29 @@ copy_table(Relation rel)
 	relmapentry = logicalrep_rel_open(lrel.remoteid, NoLock);
 	Assert(rel == relmapentry->localrel);
 
+	attnamelist = make_copy_attnamelist(relmapentry, remotegenlist);
+
 	/* Start copy on the publisher. */
 	initStringInfo(&cmd);
 
-	/* Regular table with no row filter */
-	if (lrel.relkind == RELKIND_RELATION && qual == NIL)
+	/*
+	 * Check if the remote table has any generated columns that should be
+	 * copied.
+	 */
+	for (int i = 0; i < relmapentry->remoterel.natts; i++)
+	{
+		if (remotegenlist[i])
+		{
+			gencol_copy_needed = true;
+			break;
+		}
+	}
+
+	/*
+	 * Regular table with no row filter and copy of generated columns is not
+	 * necessary.
+	 */
+	if (lrel.relkind == RELKIND_RELATION && qual == NIL && !gencol_copy_needed)
 	{
 		appendStringInfo(&cmd, "COPY %s",
 						 quote_qualified_identifier(lrel.nspname, lrel.relname));
@@ -1173,13 +1292,20 @@ copy_table(Relation rel)
 		 * (SELECT ...), but we can't just do SELECT * because we need to not
 		 * copy generated columns. For tables with any row filters, build a
 		 * SELECT query with OR'ed row filters for COPY.
+		 *
+		 * We also need to use this same COPY (SELECT ...) syntax when
+		 * 'publish_generated_columns' is specified as true and the remote
+		 * table has generated columns, because copy of generated columns is
+		 * not supported by the normal COPY.
 		 */
+		int			i = 0;
+
 		appendStringInfoString(&cmd, "COPY (SELECT ");
-		for (int i = 0; i < lrel.natts; i++)
+		foreach_node(String, att_name, attnamelist)
 		{
-			appendStringInfoString(&cmd, quote_identifier(lrel.attnames[i]));
-			if (i < lrel.natts - 1)
+			if (i++)
 				appendStringInfoString(&cmd, ", ");
+			appendStringInfoString(&cmd, quote_identifier(strVal(att_name)));
 		}
 
 		appendStringInfoString(&cmd, " FROM ");
@@ -1237,7 +1363,6 @@ copy_table(Relation rel)
 	(void) addRangeTableEntryForRelation(pstate, rel, AccessShareLock,
 										 NULL, false, false);
 
-	attnamelist = make_copy_attnamelist(relmapentry);
 	cstate = BeginCopyFrom(pstate, rel, NULL, NULL, false, copy_read_data, attnamelist, options);
 
 	/* Do the copy */
diff --git a/src/include/catalog/pg_subscription.h b/src/include/catalog/pg_subscription.h
index 0aa14ec4a2..158b444275 100644
--- a/src/include/catalog/pg_subscription.h
+++ b/src/include/catalog/pg_subscription.h
@@ -20,6 +20,7 @@
 #include "access/xlogdefs.h"
 #include "catalog/genbki.h"
 #include "catalog/pg_subscription_d.h"
+#include "lib/stringinfo.h"
 
 #include "nodes/pg_list.h"
 
@@ -180,4 +181,7 @@ extern void DisableSubscription(Oid subid);
 
 extern int	CountDBSubscriptions(Oid dbid);
 
+extern void get_publications_str(List *publications, StringInfo dest,
+								 bool quote_literal);
+
 #endif							/* PG_SUBSCRIPTION_H */
diff --git a/src/include/replication/logicalrelation.h b/src/include/replication/logicalrelation.h
index e687b40a56..8cdb7affbf 100644
--- a/src/include/replication/logicalrelation.h
+++ b/src/include/replication/logicalrelation.h
@@ -41,7 +41,8 @@ typedef struct LogicalRepRelMapEntry
 
 extern void logicalrep_relmap_update(LogicalRepRelation *remoterel);
 extern void logicalrep_partmap_reset_relmap(LogicalRepRelation *remoterel);
-
+extern int	logicalrep_rel_att_by_name(LogicalRepRelation *remoterel,
+									   const char *attname);
 extern LogicalRepRelMapEntry *logicalrep_rel_open(LogicalRepRelId remoteid,
 												  LOCKMODE lockmode);
 extern LogicalRepRelMapEntry *logicalrep_partition_open(LogicalRepRelMapEntry *root,
-- 
2.41.0.windows.3