v6-0003-Add-some-checks-before-using-apply-background-wor.patch

From 4de6226e61f5c27d47ddb2dfc4b127c8d8ce1ace Mon Sep 17 00:00:00 2001
From: wangw <wangw.fnst@fujitsu.com>
Date: Tue, 24 May 2022 21:08:23 +0800
Subject: [PATCH v6 3/4] Add some checks before using apply background worker
 to apply changes.

If any of the following checks are violated, an error will be reported.
1. The unique columns between publisher and subscriber are difference.
2. There is any non-immutable function present in expression in subscriber's
relation. Check from the following 3 items:
    a. The function in triggers;
    b. Column default value expressions and domain constraints;
    c. Constraint expressions.
---
 src/backend/replication/logical/proto.c       |  60 +++++-
 src/backend/replication/logical/relation.c    | 179 +++++++++++++++++-
 src/backend/replication/logical/tablesync.c   |   1 +
 src/backend/replication/logical/worker.c      |  47 +++++
 src/backend/utils/cache/typcache.c            |  17 ++
 src/include/replication/logicalproto.h        |   1 +
 src/include/replication/logicalrelation.h     |  15 ++
 src/include/utils/typcache.h                  |   2 +
 .../subscription/t/022_twophase_cascade.pl    |   7 +
 9 files changed, 324 insertions(+), 5 deletions(-)

diff --git a/src/backend/replication/logical/proto.c b/src/backend/replication/logical/proto.c
index ff8513e2d2..7da0deb830 100644
--- a/src/backend/replication/logical/proto.c
+++ b/src/backend/replication/logical/proto.c
@@ -23,7 +23,8 @@
 /*
  * Protocol message flags.
  */
-#define LOGICALREP_IS_REPLICA_IDENTITY 1
+#define LOGICALREP_IS_REPLICA_IDENTITY		0x0001
+#define LOGICALREP_IS_UNIQUE				0x0002
 
 #define MESSAGE_TRANSACTIONAL (1<<0)
 #define TRUNCATE_CASCADE		(1<<0)
@@ -933,11 +934,53 @@ logicalrep_write_attrs(StringInfo out, Relation rel, Bitmapset *columns)
 	TupleDesc	desc;
 	int			i;
 	uint16		nliveatts = 0;
-	Bitmapset  *idattrs = NULL;
+	Bitmapset  *idattrs = NULL,
+			   *attunique = NULL;
 	bool		replidentfull;
 
 	desc = RelationGetDescr(rel);
 
+	if (rel->rd_rel->relhasindex)
+	{
+		List	   *indexoidlist = RelationGetIndexList(rel);
+		ListCell   *indexoidscan;
+		foreach(indexoidscan, indexoidlist)
+		{
+			Oid			indexoid = lfirst_oid(indexoidscan);
+			Relation	indexRel;
+
+			/* Look up the description for index */
+			indexRel = RelationIdGetRelation(indexoid);
+
+			if (!RelationIsValid(indexRel))
+				elog(ERROR, "could not open relation with OID %u", indexoid);
+
+			if (indexRel->rd_index->indisunique)
+			{
+				int		i;
+				/* Add referenced attributes to idindexattrs */
+				for (i = 0; i < indexRel->rd_index->indnatts; i++)
+				{
+					int			attrnum = indexRel->rd_index->indkey.values[i];
+
+					/*
+					 * We don't include non-key columns into idindexattrs bitmaps. See
+					 * RelationGetIndexAttrBitmap.
+					 */
+					if (attrnum != 0)
+					{
+						if (i < indexRel->rd_index->indnkeyatts &&
+							!bms_is_member(attrnum - FirstLowInvalidHeapAttributeNumber, attunique))
+							attunique = bms_add_member(attunique,
+														  attrnum - FirstLowInvalidHeapAttributeNumber);
+					}
+				}
+			}
+			RelationClose(indexRel);
+		}
+		list_free(indexoidlist);
+	}
+
 	/* send number of live attributes */
 	for (i = 0; i < desc->natts; i++)
 	{
@@ -976,6 +1019,10 @@ logicalrep_write_attrs(StringInfo out, Relation rel, Bitmapset *columns)
 						  idattrs))
 			flags |= LOGICALREP_IS_REPLICA_IDENTITY;
 
+		if (bms_is_member(att->attnum - FirstLowInvalidHeapAttributeNumber,
+						  attunique))
+			flags |= LOGICALREP_IS_UNIQUE;
+
 		pq_sendbyte(out, flags);
 
 		/* attribute name */
@@ -1001,7 +1048,8 @@ logicalrep_read_attrs(StringInfo in, LogicalRepRelation *rel)
 	int			natts;
 	char	  **attnames;
 	Oid		   *atttyps;
-	Bitmapset  *attkeys = NULL;
+	Bitmapset  *attkeys = NULL,
+			   *attunique = NULL;
 
 	natts = pq_getmsgint(in, 2);
 	attnames = palloc(natts * sizeof(char *));
@@ -1012,11 +1060,14 @@ logicalrep_read_attrs(StringInfo in, LogicalRepRelation *rel)
 	{
 		uint8		flags;
 
-		/* Check for replica identity column */
+		/* Check for replica identity and unique column */
 		flags = pq_getmsgbyte(in);
 		if (flags & LOGICALREP_IS_REPLICA_IDENTITY)
 			attkeys = bms_add_member(attkeys, i);
 
+		if (flags & LOGICALREP_IS_UNIQUE)
+			attunique = bms_add_member(attunique, i);
+
 		/* attribute name */
 		attnames[i] = pstrdup(pq_getmsgstring(in));
 
@@ -1030,6 +1081,7 @@ logicalrep_read_attrs(StringInfo in, LogicalRepRelation *rel)
 	rel->attnames = attnames;
 	rel->atttyps = atttyps;
 	rel->attkeys = attkeys;
+	rel->attunique = attunique;
 	rel->natts = natts;
 }
 
diff --git a/src/backend/replication/logical/relation.c b/src/backend/replication/logical/relation.c
index 80fb561a9a..d038f7dc7a 100644
--- a/src/backend/replication/logical/relation.c
+++ b/src/backend/replication/logical/relation.c
@@ -19,12 +19,18 @@
 
 #include "access/table.h"
 #include "catalog/namespace.h"
+#include "catalog/pg_proc.h"
 #include "catalog/pg_subscription_rel.h"
 #include "executor/executor.h"
 #include "nodes/makefuncs.h"
+#include "optimizer/optimizer.h"
 #include "replication/logicalrelation.h"
 #include "replication/worker_internal.h"
+#include "rewrite/rewriteHandler.h"
 #include "utils/inval.h"
+#include "utils/lsyscache.h"
+#include "utils/syscache.h"
+#include "utils/typcache.h"
 
 
 static MemoryContext LogicalRepRelMapContext = NULL;
@@ -91,6 +97,23 @@ logicalrep_relmap_invalidate_cb(Datum arg, Oid reloid)
 	}
 }
 
+/*
+ * Reset the flag volatility of all existing entry in the relation map cache.
+ */
+static void
+logicalrep_relmap_reset_volatility_cb(Datum arg, int cacheid, uint32 hashvalue)
+{
+	HASH_SEQ_STATUS hash_seq;
+	LogicalRepRelMapEntry *entry;
+
+	if (LogicalRepRelMap == NULL)
+		return;
+
+	hash_seq_init(&hash_seq, LogicalRepRelMap);
+	while ((entry = hash_seq_search(&hash_seq)) != NULL)
+		entry->volatility = FUNCTION_UNKNOWN;
+}
+
 /*
  * Initialize the relation map cache.
  */
@@ -116,6 +139,9 @@ logicalrep_relmap_init(void)
 	/* Watch for invalidation events. */
 	CacheRegisterRelcacheCallback(logicalrep_relmap_invalidate_cb,
 								  (Datum) 0);
+	CacheRegisterSyscacheCallback(PROCOID,
+								  logicalrep_relmap_reset_volatility_cb,
+								  (Datum) 0);
 }
 
 /*
@@ -142,6 +168,7 @@ logicalrep_relmap_free_entry(LogicalRepRelMapEntry *entry)
 		pfree(remoterel->atttyps);
 	}
 	bms_free(remoterel->attkeys);
+	bms_free(remoterel->attunique);
 
 	if (entry->attrmap)
 		pfree(entry->attrmap);
@@ -190,6 +217,7 @@ logicalrep_relmap_update(LogicalRepRelation *remoterel)
 	}
 	entry->remoterel.replident = remoterel->replident;
 	entry->remoterel.attkeys = bms_copy(remoterel->attkeys);
+	entry->remoterel.attunique = bms_copy(remoterel->attunique);
 	MemoryContextSwitchTo(oldctx);
 }
 
@@ -307,7 +335,8 @@ logicalrep_rel_open(LogicalRepRelId remoteid, LOCKMODE lockmode)
 	if (!entry->localrelvalid)
 	{
 		Oid			relid;
-		Bitmapset  *idkey;
+		Bitmapset  *idkey,
+				   *ukey;
 		TupleDesc	desc;
 		MemoryContext oldctx;
 		int			i;
@@ -415,6 +444,153 @@ logicalrep_rel_open(LogicalRepRelId remoteid, LOCKMODE lockmode)
 			}
 		}
 
+		/*
+		 * Check whether the unique index of publisher and subscriber are
+		 * consistent.
+		 */
+		entry->sameunique = true;
+		ukey = RelationGetIndexAttrBitmap(entry->localrel,
+										  INDEX_ATTR_BITMAP_KEY);
+
+		if (ukey)
+		{
+			i = -1;
+			while ((i = bms_next_member(ukey, i)) >= 0)
+			{
+				int			attnum = i + FirstLowInvalidHeapAttributeNumber;
+
+				attnum = AttrNumberGetAttrOffset(attnum);
+
+				if (entry->attrmap->attnums[attnum] < 0 ||
+					!bms_is_member(entry->attrmap->attnums[attnum], remoterel->attunique))
+				{
+					entry->sameunique = false;
+					break;
+				}
+			}
+		}
+
+		if (entry->volatility == FUNCTION_UNKNOWN)
+		{
+			/*
+			 * Check from the following points:
+			 *
+			 * a. The function in triggers;
+			 * b. Column default value expressions and domain constraints;
+			 * c. Constraint expressions.
+			 */
+			bool	nonimmutable = false;
+
+			/* Check the trigger functions. */
+			if (!nonimmutable)
+			{
+				int	i;
+
+				if (entry->localrel->trigdesc != NULL)
+					for (i = 0; i < entry->localrel->trigdesc->numtriggers; i++)
+					{
+						Trigger		*trig = entry->localrel->trigdesc->triggers + i;
+
+						if (func_volatile(trig->tgfoid) != PROVOLATILE_IMMUTABLE)
+						{
+							nonimmutable = true;
+							break;
+						}
+					}
+			}
+
+			/* Check the columns. */
+			if (!nonimmutable)
+			{
+				TupleDesc	tupdesc = RelationGetDescr(entry->localrel);
+				int			attnum;
+
+				for (attnum = 0; attnum < tupdesc->natts; attnum++)
+				{
+					Form_pg_attribute att = TupleDescAttr(tupdesc, attnum);
+
+					/* break if we set nonimmutable in last check for domain. */
+					if (nonimmutable)
+						break;
+
+					/* We don't need info for dropped or generated attributes */
+					if (att->attisdropped || att->attgenerated)
+						continue;
+
+					/* We don't need to check columns that only exist on the subscriber */
+					if (entry->attrmap->attnums[attnum] < 0)
+						continue;
+
+					if (att->atthasdef)
+					{
+						Node	   *defaultexpr;
+
+						defaultexpr = build_column_default(entry->localrel, attnum + 1);
+						if (contain_mutable_functions(defaultexpr))
+						{
+							nonimmutable = true;
+							break;
+						}
+					}
+
+					/*
+					 * If the column is of a DOMAIN type, determine whether that
+					 * domain has any CHECK expressions that are not immutable.
+					 */
+					if (get_typtype(att->atttypid) == TYPTYPE_DOMAIN)
+					{
+						List		   *domain_constraints;
+						ListCell	   *lc;
+
+						domain_constraints = GetDomainConstraints(att->atttypid);
+
+						foreach(lc, domain_constraints)
+						{
+							DomainConstraintState *con = (DomainConstraintState *) lfirst(lc);
+
+							if (con->check_expr && contain_mutable_functions((Node *) con->check_expr))
+							{
+								nonimmutable = true;
+								break;
+							}
+						}
+					}
+				}
+			}
+
+			/* Check the constraints. */
+			if (!nonimmutable)
+			{
+				TupleDesc		tupdesc = RelationGetDescr(entry->localrel);
+
+				if (tupdesc->constr)
+				{
+					ConstrCheck	   *check = tupdesc->constr->check;
+					int				i;
+
+					/*
+					 * Determine if there are any CHECK constraints which contains
+					 * non-immutable function.
+					 */
+					for (i = 0; i < tupdesc->constr->num_check; i++)
+					{
+						Expr	   *check_expr = stringToNode(check[i].ccbin);
+
+						if (contain_mutable_functions((Node *) check_expr))
+						{
+							nonimmutable = true;
+							break;
+						}
+					}
+				}
+			}
+
+			if (nonimmutable)
+				entry->volatility = FUNCTION_NONIMMUTABLE;
+			else
+				entry->volatility = FUNCTION_IMMUTABLE;
+		}
+
 		entry->localrelvalid = true;
 	}
 
@@ -570,6 +746,7 @@ logicalrep_partition_open(LogicalRepRelMapEntry *root,
 	}
 	entry->remoterel.replident = remoterel->replident;
 	entry->remoterel.attkeys = bms_copy(remoterel->attkeys);
+	entry->remoterel.attunique = bms_copy(remoterel->attunique);
 
 	entry->localrel = partrel;
 	entry->localreloid = partOid;
diff --git a/src/backend/replication/logical/tablesync.c b/src/backend/replication/logical/tablesync.c
index 5eb60327b6..4885eaa078 100644
--- a/src/backend/replication/logical/tablesync.c
+++ b/src/backend/replication/logical/tablesync.c
@@ -880,6 +880,7 @@ fetch_remote_table_info(char *nspname, char *relname,
 	lrel->attnames = palloc0(MaxTupleAttributeNumber * sizeof(char *));
 	lrel->atttyps = palloc0(MaxTupleAttributeNumber * sizeof(Oid));
 	lrel->attkeys = NULL;
+	lrel->attunique = NULL;
 
 	/*
 	 * Store the columns as a list of names.  Ignore those that are not
diff --git a/src/backend/replication/logical/worker.c b/src/backend/replication/logical/worker.c
index fd47c48524..62c3a0d8da 100644
--- a/src/backend/replication/logical/worker.c
+++ b/src/backend/replication/logical/worker.c
@@ -287,6 +287,7 @@ static void apply_bgworker_send_data(WorkerState *wstate, Size nbytes,
 static void apply_bgworker_free(WorkerState *wstate);
 static void apply_bgworker_check_status(void);
 static void apply_bgworker_set_state(char state);
+static void apply_bgworker_relation_check(LogicalRepRelMapEntry *rel);
 
 typedef struct FlushPosition
 {
@@ -2127,6 +2128,8 @@ apply_handle_insert(StringInfo s)
 	/* Set relation for error callback */
 	apply_error_callback_arg.rel = rel;
 
+	apply_bgworker_relation_check(rel);
+
 	/* Initialize the executor state. */
 	edata = create_edata_for_relation(rel);
 	estate = edata->estate;
@@ -2263,6 +2266,8 @@ apply_handle_update(StringInfo s)
 	/* Check if we can do the update. */
 	check_relation_updatable(rel);
 
+	apply_bgworker_relation_check(rel);
+
 	/* Initialize the executor state. */
 	edata = create_edata_for_relation(rel);
 	estate = edata->estate;
@@ -2431,6 +2436,8 @@ apply_handle_delete(StringInfo s)
 	/* Check if we can do the delete. */
 	check_relation_updatable(rel);
 
+	apply_bgworker_relation_check(rel);
+
 	/* Initialize the executor state. */
 	edata = create_edata_for_relation(rel);
 	estate = edata->estate;
@@ -5119,3 +5126,43 @@ apply_bgworker_set_state(char state)
 	MyParallelState->state = state;
 	SpinLockRelease(&MyParallelState->mutex);
 }
+
+/*
+ * Check if changes on this logical replication relation  can be applied by
+ * apply background worker.
+ */
+static void
+apply_bgworker_relation_check(LogicalRepRelMapEntry *rel)
+{
+	/* Check only we are in apply bgworker. */
+	if (!am_apply_bgworker())
+		return;
+
+	/*
+	 * FIXME: Checks on partitioned tables are currently not supported.
+	 * Adding this check requires additional fixes for other issues.
+	 * I will add this later, after the related issues are fixed.
+	 */
+	if (rel->localrel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE)
+		return;
+
+	/*
+	 * If any unique index exist, check that they are same as remoterel.
+	 */
+	if (!rel->sameunique)
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("cannot replicate relation with different unique index"),
+				 errhint("Please change the streaming option to 'on' instead of 'apply'.")));
+
+	/*
+	 * Check if there is any non-immutable function present in expression in
+	 * this relation.
+	 */
+	if (rel->volatility == FUNCTION_NONIMMUTABLE)
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("cannot replicate relation. There is at least one non-immutable function"),
+				 errhint("Please change the streaming option to 'on' instead of 'apply'.")));
+
+}
diff --git a/src/backend/utils/cache/typcache.c b/src/backend/utils/cache/typcache.c
index 808f9ebd0d..b248899d82 100644
--- a/src/backend/utils/cache/typcache.c
+++ b/src/backend/utils/cache/typcache.c
@@ -2540,6 +2540,23 @@ compare_values_of_enum(TypeCacheEntry *tcache, Oid arg1, Oid arg2)
 		return 0;
 }
 
+/*
+ * GetDomainConstraints --- get DomainConstraintState list of specified domain type
+ */
+List *
+GetDomainConstraints(Oid type_id)
+{
+	TypeCacheEntry *typentry;
+	List		   *constraints = NIL;
+
+	typentry = lookup_type_cache(type_id, TYPECACHE_DOMAIN_CONSTR_INFO);
+
+	if(typentry->domainData != NULL)
+		constraints = typentry->domainData->constraints;
+
+	return constraints;
+}
+
 /*
  * Load (or re-load) the enumData member of the typcache entry.
  */
diff --git a/src/include/replication/logicalproto.h b/src/include/replication/logicalproto.h
index a771ab8ff3..a88535820b 100644
--- a/src/include/replication/logicalproto.h
+++ b/src/include/replication/logicalproto.h
@@ -108,6 +108,7 @@ typedef struct LogicalRepRelation
 	char		replident;		/* replica identity */
 	char		relkind;		/* remote relation kind */
 	Bitmapset  *attkeys;		/* Bitmap of key columns */
+	Bitmapset  *attunique;		/* Bitmap of unique columns */
 } LogicalRepRelation;
 
 /* Type mapping info */
diff --git a/src/include/replication/logicalrelation.h b/src/include/replication/logicalrelation.h
index 7bf8cd22bd..9f568e9c00 100644
--- a/src/include/replication/logicalrelation.h
+++ b/src/include/replication/logicalrelation.h
@@ -15,6 +15,17 @@
 #include "access/attmap.h"
 #include "replication/logicalproto.h"
 
+/*
+ *	States to determine volatility of the function in expressions in one
+ *	relation.
+ */
+typedef enum RelFuncVolatility
+{
+	FUNCTION_UNKNOWN = 0,	/* initializing  */
+	FUNCTION_IMMUTABLE,		/* all functions are immutable function */
+	FUNCTION_NONIMMUTABLE	/* at least one non-immutable function */
+} RelFuncVolatility;
+
 typedef struct LogicalRepRelMapEntry
 {
 	LogicalRepRelation remoterel;	/* key is remoterel.remoteid */
@@ -31,6 +42,10 @@ typedef struct LogicalRepRelMapEntry
 	Relation	localrel;		/* relcache entry (NULL when closed) */
 	AttrMap    *attrmap;		/* map of local attributes to remote ones */
 	bool		updatable;		/* Can apply updates/deletes? */
+	bool		sameunique;		/* Is the unique column of local and remote
+								   consistent? */
+	RelFuncVolatility	volatility;		/* all functions in localrel are
+								   immutable function? */
 
 	/* Sync state. */
 	char		state;
diff --git a/src/include/utils/typcache.h b/src/include/utils/typcache.h
index 431ad7f1b3..ed7c2e7f48 100644
--- a/src/include/utils/typcache.h
+++ b/src/include/utils/typcache.h
@@ -199,6 +199,8 @@ extern uint64 assign_record_type_identifier(Oid type_id, int32 typmod);
 
 extern int	compare_values_of_enum(TypeCacheEntry *tcache, Oid arg1, Oid arg2);
 
+extern List *GetDomainConstraints(Oid type_id);
+
 extern size_t SharedRecordTypmodRegistryEstimate(void);
 
 extern void SharedRecordTypmodRegistryInit(SharedRecordTypmodRegistry *,
diff --git a/src/test/subscription/t/022_twophase_cascade.pl b/src/test/subscription/t/022_twophase_cascade.pl
index 663808acb9..b317144230 100644
--- a/src/test/subscription/t/022_twophase_cascade.pl
+++ b/src/test/subscription/t/022_twophase_cascade.pl
@@ -38,6 +38,13 @@ sub test_streaming
 		ALTER SUBSCRIPTION tap_sub_C
 		SET (streaming = $streaming_mode)");
 
+	if ($streaming_mode eq 'apply')
+	{
+		$node_C->safe_psql(
+		'postgres', "
+		ALTER TABLE test_tab ALTER c DROP DEFAULT");
+	}
+
 	# Wait for subscribers to finish initialization
 
 	$node_A->poll_query_until(
-- 
2.23.0.windows.1