From c235ea36218bf3466397a713dc5e93112349024a Mon Sep 17 00:00:00 2001
From: Petr Jelinek <pjmodos@pjmodos.net>
Date: Sun, 28 Oct 2018 23:22:42 +0100
Subject: [PATCH] Fix REPLICA IDENTITY handling in logical replication protocol

We need column bitmap of replica identity so that we can send the
correct key to the downstream. We used to build all the index
information for this purpose but hat has two flaws:

- The logical decoding does not have proper snapshot for building index
  info for functional indexes.
- It's also not very efficient to build info we will never need

This commit solves both of those by introducing special purpose function
which only builds index info and hence the bitmap for REPLICA IDENTITY
index.
---
 src/backend/replication/logical/proto.c | 58 ++++++++++++++++++++++++-
 1 file changed, 56 insertions(+), 2 deletions(-)

diff --git a/src/backend/replication/logical/proto.c b/src/backend/replication/logical/proto.c
index 94dfee0b24..54855844f6 100644
--- a/src/backend/replication/logical/proto.c
+++ b/src/backend/replication/logical/proto.c
@@ -18,6 +18,7 @@
 #include "libpq/pqformat.h"
 #include "replication/logicalproto.h"
 #include "utils/builtins.h"
+#include "utils/catcache.h"
 #include "utils/lsyscache.h"
 #include "utils/syscache.h"

@@ -501,6 +502,60 @@ logicalrep_read_tuple(StringInfo in, LogicalRepTupleData *tuple)
 	}
 }

+/*
+ * Get a bitmap of identity index attribute numbers.
+ *
+ * We cannot use RelationGetIndexAttrBitmap in output plugin because it would
+ * try to build IndexInfo for all the relation's indexes, which would be a
+ * problem, as the right authorization environment for that has not been set
+ * up (a problem particularly if there are any SQL-language functions.)
+ */
+static Bitmapset *
+get_rel_idattrs(Relation relation)
+{
+	Bitmapset  *idattrs = NULL;
+	Oid			idindex = RelationGetReplicaIndex(relation);
+	int			i;
+
+	/* If we already have the bitmap, just return it. */
+	if (relation->rd_idattr != NULL)
+		return bms_copy(relation->rd_idattr);
+
+	/* Otherwise build it from the identity index if there is one. */
+	if (OidIsValid(idindex))
+	{
+		HeapTuple	tuple;
+		Form_pg_index index;
+		int			natts;
+		MemoryContext oldcxt;
+
+		tuple = SearchSysCache1(INDEXRELID, ObjectIdGetDatum(idindex));
+		if (!tuple)
+			elog(ERROR, "cache lookup failed for index %u", idindex);
+		index = (Form_pg_index) GETSTRUCT(tuple);
+		natts = index->indnatts;
+
+		for (i = 0; i < natts; i++)
+			idattrs = bms_add_member(idattrs, index->indkey.values[i] -
+									 FirstLowInvalidHeapAttributeNumber);
+		ReleaseSysCache(tuple);
+
+		/*
+		 * Save copy of the bitmap in the relcache entry.
+		 * Note this will not make the relcache index info valid for
+		 * RelationGetIndexAttrBitmap, but that does not matter as we don't
+		 * mind if it rebuilds it - it should not anyway as this is supposed
+		 * to be only called in output plugin where
+		 * RelationGetIndexAttrBitmap does not work.
+		 */
+		oldcxt = MemoryContextSwitchTo(CacheMemoryContext);
+		relation->rd_idattr = bms_copy(idattrs);
+		MemoryContextSwitchTo(oldcxt);
+	}
+
+	return idattrs;
+}
+
 /*
  * Write relation attributes to the stream.
  */
@@ -527,8 +582,7 @@ logicalrep_write_attrs(StringInfo out, Relation rel)
 	/* fetch bitmap of REPLICATION IDENTITY attributes */
 	replidentfull = (rel->rd_rel->relreplident == REPLICA_IDENTITY_FULL);
 	if (!replidentfull)
-		idattrs = RelationGetIndexAttrBitmap(rel,
-											 INDEX_ATTR_BITMAP_IDENTITY_KEY);
+		idattrs = get_rel_idattrs(rel);

 	/* send the attributes */
 	for (i = 0; i < desc->natts; i++)
-- 
2.17.1