Fix CREATE EXTENSION propagation with custom version

ryogrid-watch · Mar 9, 2022 · 8e43c80 · 8e43c80
1 parent 7559ad1
commit 8e43c80
Show file tree

Hide file tree

Showing 4 changed files with 56 additions and 3 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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)