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

@@ -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;
 }
 

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

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);