Fixes the bug where undistribute can drop Citus extension

(cherry picked from commit b255706189)

 Conflicts:
	src/backend/distributed/commands/alter_table.c
	src/include/distributed/metadata/dependency.h

src/backend/distributed/commands/create_distributed_table.c

@@ -32,6 +32,7 @@
 #if PG_VERSION_NUM >= 12000
 #include "catalog/pg_proc.h"
 #endif
+#include "catalog/pg_rewrite_d.h"
 #include "catalog/pg_trigger.h"
 #include "commands/defrem.h"
 #include "commands/extension.h"
@@ -127,6 +128,8 @@ static void DoCopyFromLocalTableIntoShards(Relation distributedRelation,
 static void UndistributeTable(Oid relationId);
 static List * GetViewCreationCommandsOfTable(Oid relationId);
 static void ReplaceTable(Oid sourceId, Oid targetId);
+static void ErrorIfUnsupportedCascadeObjects(Oid relationId);
+static bool DoesCascadeDropUnsupportedObject(Oid classId, Oid id, HTAB *nodeMap);
 
 /* exports for SQL callable functions */
 PG_FUNCTION_INFO_V1(master_create_distributed_table);
@@ -1619,6 +1622,8 @@ UndistributeTable(Oid relationId)
 								parentRelationName)));
 	}
 
+	ErrorIfUnsupportedCascadeObjects(relationId);
+
 	List *preLoadCommands = GetPreLoadTableCreationCommands(relationId, true);
 	List *postLoadCommands = GetPostLoadTableCreationCommands(relationId);
 
@@ -1799,3 +1804,91 @@ ReplaceTable(Oid sourceId, Oid targetId)
 						   sourceName, false);
 #endif
 }
+
+
+/*
+ * ErrorIfUnsupportedCascadeObjects gets oid of a relation, finds the objects
+ * that dropping this relation cascades into and errors if there are any extensions
+ * that would be dropped.
+ */
+static void
+ErrorIfUnsupportedCascadeObjects(Oid relationId)
+{
+	HASHCTL info;
+	memset(&info, 0, sizeof(info));
+	info.keysize = sizeof(Oid);
+	info.entrysize = sizeof(Oid);
+	info.hash = oid_hash;
+	uint32 hashFlags = (HASH_ELEM | HASH_FUNCTION);
+	HTAB *nodeMap = hash_create("object dependency map (oid)", 64, &info, hashFlags);
+
+	bool unsupportedObjectInDepGraph =
+		DoesCascadeDropUnsupportedObject(RelationRelationId, relationId, nodeMap);
+
+	if (unsupportedObjectInDepGraph)
+	{
+		ereport(ERROR, (errmsg("cannot alter table because an extension depends on it")));
+	}
+}
+
+
+/*
+ * DoesCascadeDropUnsupportedObject walks through the objects that depend on the
+ * object with object id and returns true if it finds any unsupported objects.
+ *
+ * This function only checks extensions as unsupported objects.
+ *
+ * Extension dependency is different than the rest. If an object depends on an extension
+ * dropping the object would drop the extension too.
+ * So we check with IsObjectAddressOwnedByExtension function.
+ */
+static bool
+DoesCascadeDropUnsupportedObject(Oid classId, Oid objectId, HTAB *nodeMap)
+{
+	bool found = false;
+	hash_search(nodeMap, &objectId, HASH_ENTER, &found);
+
+	if (found)
+	{
+		return false;
+	}
+
+	ObjectAddress objectAddress = { 0 };
+	ObjectAddressSet(objectAddress, classId, objectId);
+
+	if (IsObjectAddressOwnedByExtension(&objectAddress, NULL))
+	{
+		return true;
+	}
+
+	Oid targetObjectClassId = classId;
+	Oid targetObjectId = objectId;
+	List *dependencyTupleList = GetPgDependTuplesForDependingObjects(targetObjectClassId,
+																	 targetObjectId);
+
+	HeapTuple depTup = NULL;
+	foreach_ptr(depTup, dependencyTupleList)
+	{
+		Form_pg_depend pg_depend = (Form_pg_depend) GETSTRUCT(depTup);
+
+		Oid dependingOid = InvalidOid;
+		Oid dependingClassId = InvalidOid;
+
+		if (pg_depend->classid == RewriteRelationId)
+		{
+			dependingOid = GetDependingView(pg_depend);
+			dependingClassId = RelationRelationId;
+		}
+		else
+		{
+			dependingOid = pg_depend->objid;
+			dependingClassId = pg_depend->classid;
+		}
+
+		if (DoesCascadeDropUnsupportedObject(dependingClassId, dependingOid, nodeMap))
+		{
+			return true;
+		}
+	}
+	return false;
+}

src/backend/distributed/metadata/dependency.c

@@ -154,7 +154,6 @@ static void ApplyAddToDependencyList(ObjectAddressCollector *collector,
 static List * ExpandCitusSupportedTypes(ObjectAddressCollector *collector,
 										ObjectAddress target);
 static ViewDependencyNode * BuildViewDependencyGraph(Oid relationId, HTAB *nodeMap);
-static Oid GetDependingView(Form_pg_depend pg_depend);
 
 
 /*

src/include/distributed/metadata/dependency.h

@@ -15,6 +15,7 @@
 #include "postgres.h"
 
 #include "catalog/objectaddress.h"
+#include "catalog/pg_depend.h"
 #include "nodes/pg_list.h"
 
 extern List * GetUniqueDependenciesList(List *objectAddressesList);
@@ -24,5 +25,6 @@ extern bool SupportedDependencyByCitus(const ObjectAddress *address);
 extern List * GetPgDependTuplesForDependingObjects(Oid targetObjectClassId,
 												   Oid targetObjectId);
 extern List * GetDependingViews(Oid relationId);
+extern Oid GetDependingView(Form_pg_depend pg_depend);
 
 #endif /* CITUS_DEPENDENCY_H */

src/test/regress/expected/undistribute_table.out

@@ -383,9 +383,99 @@ SELECT * FROM another_schema.undis_view3 ORDER BY 1, 2;
  6 | 9
 (3 rows)
 
+-- test the drop in undistribute_table cannot cascade to an extension
+CREATE TABLE extension_table (a INT);
+SELECT create_distributed_table('extension_table', 'a');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE VIEW extension_view AS SELECT * FROM extension_table;
+ALTER EXTENSION plpgsql ADD VIEW extension_view;
+NOTICE:  Citus does not propagate adding/dropping member objects
+HINT:  You can add/drop the member objects on the workers as well.
+SELECT undistribute_table ('extension_table');
+ERROR:  cannot alter table because an extension depends on it
+-- test the drop that doesn't cascade to an extension
+CREATE TABLE dist_type_table (a INT, b citus.distribution_type);
+SELECT create_distributed_table('dist_type_table', 'a');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT undistribute_table('dist_type_table');
+NOTICE:  creating a new local table for undistribute_table.dist_type_table
+NOTICE:  Moving the data of undistribute_table.dist_type_table
+NOTICE:  Dropping the old undistribute_table.dist_type_table
+NOTICE:  Renaming the new table to undistribute_table.dist_type_table
+ undistribute_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- test CREATE RULE with ON SELECT
+CREATE TABLE rule_table_1 (a INT);
+CREATE TABLE rule_table_2 (a INT);
+SELECT create_distributed_table('rule_table_2', 'a');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE RULE "_RETURN" AS ON SELECT TO rule_table_1 DO INSTEAD SELECT * FROM rule_table_2;
+-- the CREATE RULE turns rule_table_1 into a view
+ALTER EXTENSION plpgsql ADD VIEW rule_table_1;
+NOTICE:  Citus does not propagate adding/dropping member objects
+HINT:  You can add/drop the member objects on the workers as well.
+SELECT undistribute_table('rule_table_2');
+ERROR:  cannot alter table because an extension depends on it
+-- test CREATE RULE without ON SELECT
+CREATE TABLE rule_table_3 (a INT);
+CREATE TABLE rule_table_4 (a INT);
+SELECT create_distributed_table('rule_table_4', 'a');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE RULE "rule_1" AS ON INSERT TO rule_table_3 DO INSTEAD SELECT * FROM rule_table_4;
+ALTER EXTENSION plpgsql ADD TABLE rule_table_3;
+NOTICE:  Citus does not propagate adding/dropping member objects
+HINT:  You can add/drop the member objects on the workers as well.
+SELECT undistribute_table('rule_table_4');
+NOTICE:  creating a new local table for undistribute_table.rule_table_4
+NOTICE:  Moving the data of undistribute_table.rule_table_4
+NOTICE:  Dropping the old undistribute_table.rule_table_4
+NOTICE:  drop cascades to rule rule_1 on table rule_table_3
+CONTEXT:  SQL statement "DROP TABLE undistribute_table.rule_table_4 CASCADE"
+NOTICE:  Renaming the new table to undistribute_table.rule_table_4
+ undistribute_table
+---------------------------------------------------------------------
+
+(1 row)
+
+ALTER EXTENSION plpgsql DROP VIEW extension_view;
+NOTICE:  Citus does not propagate adding/dropping member objects
+HINT:  You can add/drop the member objects on the workers as well.
+ALTER EXTENSION plpgsql DROP VIEW rule_table_1;
+NOTICE:  Citus does not propagate adding/dropping member objects
+HINT:  You can add/drop the member objects on the workers as well.
+ALTER EXTENSION plpgsql DROP TABLE rule_table_3;
+NOTICE:  Citus does not propagate adding/dropping member objects
+HINT:  You can add/drop the member objects on the workers as well.
 DROP TABLE view_table CASCADE;
 NOTICE:  drop cascades to 3 other objects
 DETAIL:  drop cascades to view undis_view1
 drop cascades to view undis_view2
 drop cascades to view another_schema.undis_view3
 DROP SCHEMA undistribute_table, another_schema CASCADE;
+NOTICE:  drop cascades to 7 other objects
+DETAIL:  drop cascades to table extension_table
+drop cascades to view extension_view
+drop cascades to table dist_type_table
+drop cascades to table rule_table_2
+drop cascades to view rule_table_1
+drop cascades to table rule_table_3
+drop cascades to table rule_table_4

src/test/regress/sql/undistribute_table.sql

@@ -128,6 +128,47 @@ SELECT undistribute_table('view_table');
 SELECT schemaname, viewname, viewowner, definition FROM pg_views WHERE viewname LIKE 'undis\_view%' ORDER BY viewname;
 SELECT * FROM another_schema.undis_view3 ORDER BY 1, 2;
 
+-- test the drop in undistribute_table cannot cascade to an extension
+CREATE TABLE extension_table (a INT);
+SELECT create_distributed_table('extension_table', 'a');
+CREATE VIEW extension_view AS SELECT * FROM extension_table;
+ALTER EXTENSION plpgsql ADD VIEW extension_view;
+
+SELECT undistribute_table ('extension_table');
+
+-- test the drop that doesn't cascade to an extension
+CREATE TABLE dist_type_table (a INT, b citus.distribution_type);
+SELECT create_distributed_table('dist_type_table', 'a');
+
+SELECT undistribute_table('dist_type_table');
+
+-- test CREATE RULE with ON SELECT
+CREATE TABLE rule_table_1 (a INT);
+CREATE TABLE rule_table_2 (a INT);
+SELECT create_distributed_table('rule_table_2', 'a');
+
+CREATE RULE "_RETURN" AS ON SELECT TO rule_table_1 DO INSTEAD SELECT * FROM rule_table_2;
+
+-- the CREATE RULE turns rule_table_1 into a view
+ALTER EXTENSION plpgsql ADD VIEW rule_table_1;
+
+SELECT undistribute_table('rule_table_2');
+
+-- test CREATE RULE without ON SELECT
+CREATE TABLE rule_table_3 (a INT);
+CREATE TABLE rule_table_4 (a INT);
+SELECT create_distributed_table('rule_table_4', 'a');
+
+CREATE RULE "rule_1" AS ON INSERT TO rule_table_3 DO INSTEAD SELECT * FROM rule_table_4;
+
+ALTER EXTENSION plpgsql ADD TABLE rule_table_3;
+
+SELECT undistribute_table('rule_table_4');
+
+ALTER EXTENSION plpgsql DROP VIEW extension_view;
+ALTER EXTENSION plpgsql DROP VIEW rule_table_1;
+ALTER EXTENSION plpgsql DROP TABLE rule_table_3;
+
 DROP TABLE view_table CASCADE;
 
 DROP SCHEMA undistribute_table, another_schema CASCADE;