From 8e43c8094d400812878206b6fa208e76c837f1c2 Mon Sep 17 00:00:00 2001
From: Marco Slot <marco.slot@gmail.com>
Date: Tue, 8 Mar 2022 23:27:55 +0100
Subject: [PATCH] Fix CREATE EXTENSION propagation with custom version

---
 src/backend/distributed/commands/extension.c  | 12 +++++-
 .../distributed/deparser/citus_ruleutils.c    | 42 +++++++++++++++++++
 src/include/distributed/citus_ruleutils.h     |  1 +
 .../expected/isolation_extension_commands.out |  4 +-
 4 files changed, 56 insertions(+), 3 deletions(-)

diff --git a/src/backend/distributed/commands/extension.c b/src/backend/distributed/commands/extension.c
index 3aa782c06e8..c1cf0603917 100644
--- a/src/backend/distributed/commands/extension.c
+++ b/src/backend/distributed/commands/extension.c
@@ -769,13 +769,23 @@ RecreateExtensionStmt(Oid extensionOid)
 
 	/* make DefEleme for extensionSchemaName */
 	Node *schemaNameArg = (Node *) makeString(extensionSchemaName);
-
 	DefElem *schemaDefElement = makeDefElem("schema", schemaNameArg, location);
 
 	/* append the schema name DefElem finally */
 	createExtensionStmt->options = lappend(createExtensionStmt->options,
 										   schemaDefElement);
 
+	char *extensionVersion = get_extension_version(extensionOid);
+	if (extensionVersion != NULL)
+	{
+		Node *extensionVersionArg = (Node *) makeString(extensionVersion);
+		DefElem *extensionVersionElement =
+			makeDefElem("new_version", extensionVersionArg, location);
+
+		createExtensionStmt->options = lappend(createExtensionStmt->options,
+											   extensionVersionElement);
+	}
+
 	return (Node *) createExtensionStmt;
 }
 
diff --git a/src/backend/distributed/deparser/citus_ruleutils.c b/src/backend/distributed/deparser/citus_ruleutils.c
index a2002851d5f..3da845d3b76 100644
--- a/src/backend/distributed/deparser/citus_ruleutils.c
+++ b/src/backend/distributed/deparser/citus_ruleutils.c
@@ -123,6 +123,48 @@ pg_get_extensiondef_string(Oid tableRelationId)
 }
 
 
+/*
+ * get_extension_version - given an extension OID, fetch its extversion
+ * or NULL if not found.
+ */
+char *
+get_extension_version(Oid extensionId)
+{
+	char *versionName = NULL;
+
+	Relation relation = table_open(ExtensionRelationId, AccessShareLock);
+
+	ScanKeyData entry[1];
+	ScanKeyInit(&entry[0],
+				Anum_pg_extension_oid,
+				BTEqualStrategyNumber, F_OIDEQ,
+				ObjectIdGetDatum(extensionId));
+
+	SysScanDesc scanDesc = systable_beginscan(relation, ExtensionOidIndexId, true,
+											  NULL, 1, entry);
+
+	HeapTuple tuple = systable_getnext(scanDesc);
+
+	/* We assume that there can be at most one matching tuple */
+	if (HeapTupleIsValid(tuple))
+	{
+		bool isNull = false;
+		Datum versionDatum = heap_getattr(tuple, Anum_pg_extension_extversion,
+										  RelationGetDescr(relation), &isNull);
+		if (!isNull)
+		{
+			versionName = text_to_cstring(DatumGetTextPP(versionDatum));
+		}
+	}
+
+	systable_endscan(scanDesc);
+
+	table_close(relation, AccessShareLock);
+
+	return versionName;
+}
+
+
 /*
  * get_extension_schema - given an extension OID, fetch its extnamespace
  *
diff --git a/src/include/distributed/citus_ruleutils.h b/src/include/distributed/citus_ruleutils.h
index 2396214151a..03d58d031b5 100644
--- a/src/include/distributed/citus_ruleutils.h
+++ b/src/include/distributed/citus_ruleutils.h
@@ -29,6 +29,7 @@
 /* Function declarations for version independent Citus ruleutils wrapper functions */
 extern char * pg_get_extensiondef_string(Oid tableRelationId);
 extern Oid get_extension_schema(Oid ext_oid);
+extern char * get_extension_version(Oid extensionId);
 extern char * pg_get_serverdef_string(Oid tableRelationId);
 extern char * pg_get_sequencedef_string(Oid sequenceRelid);
 extern Form_pg_sequence pg_get_sequencedef(Oid sequenceRelationId);
diff --git a/src/test/regress/expected/isolation_extension_commands.out b/src/test/regress/expected/isolation_extension_commands.out
index 497cf414ec0..028ec21f083 100644
--- a/src/test/regress/expected/isolation_extension_commands.out
+++ b/src/test/regress/expected/isolation_extension_commands.out
@@ -506,7 +506,7 @@ run_command_on_workers
 
 run_command_on_workers
 ---------------------------------------------------------------------
-(localhost,57637,t,1.3)
+(localhost,57637,t,1.1)
 (localhost,57638,t,1.1)
 (2 rows)
 
@@ -589,7 +589,7 @@ run_command_on_workers
 
 run_command_on_workers
 ---------------------------------------------------------------------
-(localhost,57637,t,1.3)
+(localhost,57637,t,1.1)
 (localhost,57638,t,1.2)
 (2 rows)