From e454c38afe989cd89608cc838bc6bfb804e46b29 Mon Sep 17 00:00:00 2001
From: Peter Eisentraut <peter@eisentraut.org>
Date: Thu, 5 Mar 2020 10:11:13 +0100
Subject: [PATCH] Preserve replica identity index across ALTER TABLE rewrite

---
 src/backend/commands/tablecmds.c              | 46 +++++++++++++++++++
 src/backend/utils/cache/lsyscache.c           | 23 ++++++++++
 src/include/utils/lsyscache.h                 |  1 +
 .../regress/expected/replica_identity.out     | 46 +++++++++++++++++++
 src/test/regress/sql/replica_identity.sql     | 21 +++++++++
 5 files changed, 137 insertions(+)

diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 9c6ca7ff9c..380c52e8df 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -168,6 +168,7 @@ typedef struct AlteredTableInfo
 	List	   *changedConstraintDefs;	/* string definitions of same */
 	List	   *changedIndexOids;		/* OIDs of indexes to rebuild */
 	List	   *changedIndexDefs;		/* string definitions of same */
+	char	   *replicaIdentityIndex;	/* index to reset as REPLICA IDENTITY */
 } AlteredTableInfo;
 
 /* Struct describing one new constraint to check in Phase 3 scan */
@@ -8596,6 +8597,22 @@ ATExecAlterColumnType(AlteredTableInfo *tab, Relation rel,
 	return address;
 }
 
+/*
+ * Subroutine for ATExecAlterColumnType: remember that a replica identity
+ * needs to be reset.
+ */
+static void
+RememberReplicaIdentityForRebuilding(Oid indoid, AlteredTableInfo *tab)
+{
+	if (!get_index_is_replident(indoid))
+		return;
+
+	if (tab->replicaIdentityIndex)
+		elog(ERROR, "relation %u has multiple indexes marked as replica identity", tab->relid);
+
+	tab->replicaIdentityIndex = get_rel_name(indoid);
+}
+
 /*
  * Subroutine for ATExecAlterColumnType: remember that a constraint needs
  * to be rebuilt (which we might already know).
@@ -8615,6 +8632,7 @@ RememberConstraintForRebuilding(Oid conoid, AlteredTableInfo *tab,
 	{
 		/* OK, capture the constraint's existing definition string */
 		char	   *defstring = pg_get_constraintdef_command(conoid);
+		Oid			indoid;
 
 		/*
 		 * Put NORMAL dependencies at the front of the list and AUTO
@@ -8630,6 +8648,10 @@ RememberConstraintForRebuilding(Oid conoid, AlteredTableInfo *tab,
 												   tab->changedConstraintOids);
 			tab->changedConstraintDefs = lcons(defstring,
 											   tab->changedConstraintDefs);
+
+			indoid = get_constraint_index(conoid);
+			if (OidIsValid(indoid))
+				RememberReplicaIdentityForRebuilding(indoid, tab);
 		}
 		else
 		{
@@ -8637,6 +8659,10 @@ RememberConstraintForRebuilding(Oid conoid, AlteredTableInfo *tab,
 													 conoid);
 			tab->changedConstraintDefs = lappend(tab->changedConstraintDefs,
 												 defstring);
+
+			indoid = get_constraint_index(conoid);
+			if (OidIsValid(indoid))
+				RememberReplicaIdentityForRebuilding(indoid, tab);
 		}
 	}
 }
@@ -8681,6 +8707,8 @@ RememberIndexForRebuilding(Oid indoid, AlteredTableInfo *tab)
 												indoid);
 			tab->changedIndexDefs = lappend(tab->changedIndexDefs,
 											defstring);
+
+			RememberReplicaIdentityForRebuilding(indoid, tab);
 		}
 	}
 }
@@ -8874,6 +8902,24 @@ ATPostAlterTypeCleanup(List **wqueue, AlteredTableInfo *tab, LOCKMODE lockmode)
 							 wqueue, lockmode, tab->rewrite);
 	}
 
+	/*
+	 * Queue up command to restore replica identity index marking
+	 */
+	if (tab->replicaIdentityIndex)
+	{
+		AlterTableCmd *cmd = makeNode(AlterTableCmd);
+		ReplicaIdentityStmt *subcmd = makeNode(ReplicaIdentityStmt);
+
+		subcmd->identity_type = REPLICA_IDENTITY_INDEX;
+		subcmd->name = tab->replicaIdentityIndex;
+		cmd->subtype = AT_ReplicaIdentity;
+		cmd->def = (Node *) subcmd;
+
+		/* do it after indexes and constraints */
+		tab->subcmds[AT_PASS_OLD_CONSTR] =
+			lappend(tab->subcmds[AT_PASS_OLD_CONSTR], cmd);
+	}
+
 	/*
 	 * Now we can drop the existing constraints and indexes --- constraints
 	 * first, since some of them might depend on the indexes.  In fact, we
diff --git a/src/backend/utils/cache/lsyscache.c b/src/backend/utils/cache/lsyscache.c
index e39c1d45cd..c49efd5e54 100644
--- a/src/backend/utils/cache/lsyscache.c
+++ b/src/backend/utils/cache/lsyscache.c
@@ -3088,3 +3088,26 @@ get_range_collation(Oid rangeOid)
 	else
 		return InvalidOid;
 }
+
+/*
+ * get_index_is_replident
+ *
+ * Returns indisreplident field.
+ */
+bool
+get_index_is_replident(Oid index_oid)
+{
+	HeapTuple		tuple;
+	Form_pg_index	rd_index;
+	bool			result;
+
+	tuple = SearchSysCache1(INDEXRELID, ObjectIdGetDatum(index_oid));
+	if (!HeapTupleIsValid(tuple))
+		return false;
+
+	rd_index = (Form_pg_index) GETSTRUCT(tuple);
+	result = rd_index->indisreplident;
+	ReleaseSysCache(tuple);
+
+	return result;
+}
diff --git a/src/include/utils/lsyscache.h b/src/include/utils/lsyscache.h
index b4757a3466..54a4dc8200 100644
--- a/src/include/utils/lsyscache.h
+++ b/src/include/utils/lsyscache.h
@@ -160,6 +160,7 @@ extern char *get_namespace_name(Oid nspid);
 extern char *get_namespace_name_or_temp(Oid nspid);
 extern Oid	get_range_subtype(Oid rangeOid);
 extern Oid	get_range_collation(Oid rangeOid);
+extern bool	get_index_is_replident(Oid index_oid);
 
 #define type_is_array(typid)  (get_element_type(typid) != InvalidOid)
 /* type_is_array_domain accepts both plain arrays and domains over arrays */
diff --git a/src/test/regress/expected/replica_identity.out b/src/test/regress/expected/replica_identity.out
index 27e2b67b8d..1071af9247 100644
--- a/src/test/regress/expected/replica_identity.out
+++ b/src/test/regress/expected/replica_identity.out
@@ -187,5 +187,51 @@ SELECT relreplident FROM pg_class WHERE oid = 'test_replica_identity'::regclass;
  n
 (1 row)
 
+---
+-- Test that ALTER TABLE rewrite preserves nondefault replica identity
+---
+-- constraint variant
+CREATE TABLE test_replica_identity2 (id int UNIQUE NOT NULL);
+ALTER TABLE test_replica_identity2 REPLICA IDENTITY USING INDEX test_replica_identity2_id_key;
+\d test_replica_identity2
+Table "public.test_replica_identity2"
+ Column |  Type   | Modifiers 
+--------+---------+-----------
+ id     | integer | not null
+Indexes:
+    "test_replica_identity2_id_key" UNIQUE CONSTRAINT, btree (id) REPLICA IDENTITY
+
+ALTER TABLE test_replica_identity2 ALTER COLUMN id TYPE bigint;
+\d test_replica_identity2
+Table "public.test_replica_identity2"
+ Column |  Type  | Modifiers 
+--------+--------+-----------
+ id     | bigint | not null
+Indexes:
+    "test_replica_identity2_id_key" UNIQUE CONSTRAINT, btree (id) REPLICA IDENTITY
+
+-- straight index variant
+CREATE TABLE test_replica_identity3 (id int NOT NULL);
+CREATE UNIQUE INDEX test_replica_identity3_id_key ON test_replica_identity3 (id);
+ALTER TABLE test_replica_identity3 REPLICA IDENTITY USING INDEX test_replica_identity3_id_key;
+\d test_replica_identity3
+Table "public.test_replica_identity3"
+ Column |  Type   | Modifiers 
+--------+---------+-----------
+ id     | integer | not null
+Indexes:
+    "test_replica_identity3_id_key" UNIQUE, btree (id) REPLICA IDENTITY
+
+ALTER TABLE test_replica_identity3 ALTER COLUMN id TYPE bigint;
+\d test_replica_identity3
+Table "public.test_replica_identity3"
+ Column |  Type  | Modifiers 
+--------+--------+-----------
+ id     | bigint | not null
+Indexes:
+    "test_replica_identity3_id_key" UNIQUE, btree (id) REPLICA IDENTITY
+
 DROP TABLE test_replica_identity;
+DROP TABLE test_replica_identity2;
+DROP TABLE test_replica_identity3;
 DROP TABLE test_replica_identity_othertable;
diff --git a/src/test/regress/sql/replica_identity.sql b/src/test/regress/sql/replica_identity.sql
index 3d2171c733..5277ff39cd 100644
--- a/src/test/regress/sql/replica_identity.sql
+++ b/src/test/regress/sql/replica_identity.sql
@@ -79,5 +79,26 @@ SELECT relreplident FROM pg_class WHERE oid = 'test_replica_identity'::regclass;
 ALTER TABLE test_replica_identity REPLICA IDENTITY NOTHING;
 SELECT relreplident FROM pg_class WHERE oid = 'test_replica_identity'::regclass;
 
+---
+-- Test that ALTER TABLE rewrite preserves nondefault replica identity
+---
+
+-- constraint variant
+CREATE TABLE test_replica_identity2 (id int UNIQUE NOT NULL);
+ALTER TABLE test_replica_identity2 REPLICA IDENTITY USING INDEX test_replica_identity2_id_key;
+\d test_replica_identity2
+ALTER TABLE test_replica_identity2 ALTER COLUMN id TYPE bigint;
+\d test_replica_identity2
+
+-- straight index variant
+CREATE TABLE test_replica_identity3 (id int NOT NULL);
+CREATE UNIQUE INDEX test_replica_identity3_id_key ON test_replica_identity3 (id);
+ALTER TABLE test_replica_identity3 REPLICA IDENTITY USING INDEX test_replica_identity3_id_key;
+\d test_replica_identity3
+ALTER TABLE test_replica_identity3 ALTER COLUMN id TYPE bigint;
+\d test_replica_identity3
+
 DROP TABLE test_replica_identity;
+DROP TABLE test_replica_identity2;
+DROP TABLE test_replica_identity3;
 DROP TABLE test_replica_identity_othertable;
-- 
2.20.1

