From 40fa3862cebfdfe7d00014262a8cf6a22ca0dcf7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96nder=20Kalac=C4=B1?= <onder@citusdata.com>
Date: Mon, 18 Nov 2019 16:57:10 +0100
Subject: [PATCH] Prevent Citus extension becoming distributed object (#3197)

 Prevent Citus extension being distributed

Because that could prevent doing rolling upgrades, where users may
prefer to upgrade the version on the coordinator but not the workers.

There could be some other edge cases, so I'd prefer to keep Citus
extension outside the picture for now.
---
 src/backend/distributed/metadata/dependency.c | 12 ++++++++
 src/backend/distributed/metadata/distobject.c | 30 ++++++++++++++++++-
 src/include/distributed/metadata/distobject.h |  1 +
 3 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/src/backend/distributed/metadata/dependency.c b/src/backend/distributed/metadata/dependency.c
index a55a5829f..8e68554e4 100644
--- a/src/backend/distributed/metadata/dependency.c
+++ b/src/backend/distributed/metadata/dependency.c
@@ -561,6 +561,12 @@ FollowNewSupportedDependencies(ObjectAddressCollector *collector, Form_pg_depend
 		return false;
 	}
 
+	if (CitusExtensionObject(&address))
+	{
+		/* following citus extension could complicate role management */
+		return false;
+	}
+
 	return true;
 }
 
@@ -611,6 +617,12 @@ FollowAllSupportedDependencies(ObjectAddressCollector *collector, Form_pg_depend
 		return false;
 	}
 
+	if (CitusExtensionObject(&address))
+	{
+		/* following citus extension could complicate role management */
+		return false;
+	}
+
 	return true;
 }
 
diff --git a/src/backend/distributed/metadata/distobject.c b/src/backend/distributed/metadata/distobject.c
index fc029b622..8f9aaa870 100644
--- a/src/backend/distributed/metadata/distobject.c
+++ b/src/backend/distributed/metadata/distobject.c
@@ -19,9 +19,12 @@
 #include "catalog/dependency.h"
 #include "catalog/namespace.h"
 #include "catalog/objectaddress.h"
+#include "catalog/pg_extension_d.h"
 #include "catalog/pg_namespace.h"
 #include "catalog/pg_proc.h"
 #include "catalog/pg_type.h"
+#include "citus_version.h"
+#include "commands/extension.h"
 #include "distributed/metadata/distobject.h"
 #include "distributed/metadata/pg_dist_object.h"
 #include "distributed/metadata_cache.h"
@@ -120,7 +123,7 @@ ObjectExists(const ObjectAddress *address)
 
 /*
  * MarkObjectDistributed marks an object as a distributed object by citus. Marking is done
- * by adding appropriate entries to citus.pg_dist_object
+ * by adding appropriate entries to citus.pg_dist_object.
  */
 void
 MarkObjectDistributed(const ObjectAddress *distAddress)
@@ -150,6 +153,31 @@ MarkObjectDistributed(const ObjectAddress *distAddress)
 }
 
 
+/*
+ * CitusExtensionObject returns true if the objectAddress represents
+ * the Citus extension.
+ */
+bool
+CitusExtensionObject(const ObjectAddress *objectAddress)
+{
+	char *extensionName = false;
+
+	if (objectAddress->classId != ExtensionRelationId)
+	{
+		return false;
+	}
+
+	extensionName = get_extension_name(objectAddress->objectId);
+	if (extensionName != NULL &&
+		strncasecmp(extensionName, CITUS_NAME, NAMEDATALEN) == 0)
+	{
+		return true;
+	}
+
+	return false;
+}
+
+
 /*
  * ExecuteCommandAsSuperuser executes a command via SPI as superuser. Using this
  * function (and in general SPI/SQL with superuser) should be avoided as much as
diff --git a/src/include/distributed/metadata/distobject.h b/src/include/distributed/metadata/distobject.h
index c8bb64829..e8f446997 100644
--- a/src/include/distributed/metadata/distobject.h
+++ b/src/include/distributed/metadata/distobject.h
@@ -17,6 +17,7 @@
 
 
 extern bool ObjectExists(const ObjectAddress *address);
+extern bool CitusExtensionObject(const ObjectAddress *objectAddress);
 extern bool IsObjectDistributed(const ObjectAddress *address);
 extern bool ClusterHasDistributedFunctionWithDistArgument(void);
 extern void MarkObjectDistributed(const ObjectAddress *distAddress);