Address reviews

src/backend/distributed/commands/function.c

@@ -1294,6 +1294,14 @@ PostprocessCreateFunctionStmt(Node *node, const char *queryString)
 		return NIL;
 	}
 
+	/*
+	 * This check should have been
+	 * (a) valid for all objects not only for functions
+	 * (b) should check for all types of objects no only for relations.
+	 *
+	 * We do this limited check for now as functions are more likely to be used with
+	 * (a) and (b), and we want to scope it for now
+	 */
 	ObjectAddress *undistributableDependency = GetUndistributableRelationDependency(
 		&functionAddress);
 	if (undistributableDependency != NULL)

src/backend/distributed/metadata/dependency.c

@@ -159,6 +159,8 @@ static bool FollowAllDependencies(ObjectAddressCollector *collector,
 								  DependencyDefinition *definition);
 static void ApplyAddToDependencyList(ObjectAddressCollector *collector,
 									 DependencyDefinition *definition);
+static void ApplyAddAllToDependencyList(ObjectAddressCollector *collector,
+										DependencyDefinition *definition);
 static List * ExpandCitusSupportedTypes(ObjectAddressCollector *collector,
 										ObjectAddress target);
 static ViewDependencyNode * BuildViewDependencyGraph(Oid relationId, HTAB *nodeMap);
@@ -240,7 +242,10 @@ GetAllSupportedDependenciesForObject(const ObjectAddress *target)
 
 /*
  * GetAllDependenciesForObject returns a list of all the dependent objects of the given
- * object irrespective of whether the dependent object is supported by Citus or not.
+ * object irrespective of whether the dependent object is supported by Citus or not, if
+ * the object can be found as dependency with RecurseObjectDependencies and
+ * ExpandCitusSupportedTypes.
+ *
  * This function will be used to provide meaningful error messages if any dependent
  * object for a given object is not supported. If you want to create dependencies for
  * an object, you probably need to use GetDependenciesForObject().
@@ -254,7 +259,7 @@ GetAllDependenciesForObject(const ObjectAddress *target)
 	RecurseObjectDependencies(*target,
 							  &ExpandCitusSupportedTypes,
 							  &FollowAllDependencies,
-							  &ApplyAddToDependencyList,
+							  &ApplyAddAllToDependencyList,
 							  &collector);
 
 	return collector.dependencyList;
@@ -971,10 +976,12 @@ FollowAllDependencies(ObjectAddressCollector *collector,
 
 
 /*
- * ApplyAddToDependencyList is an apply function for RecurseObjectDependencies that will collect
+ * ApplyAddToDependencyList is an apply function for RecurseObjectDependencies that will
- * all the ObjectAddresses for pg_depend entries to the context. The context here is
+ * collect all the ObjectAddresses for pg_depend entries to the context, except it is
- * assumed to be a (ObjectAddressCollector *) to the location where all ObjectAddresses
+ * extension owned one.
- * will be collected.
+ *
+ * The context here is assumed to be a (ObjectAddressCollector *) to the location where
+ * all ObjectAddresses will be collected.
  */
 static void
 ApplyAddToDependencyList(ObjectAddressCollector *collector,
@@ -998,6 +1005,22 @@ ApplyAddToDependencyList(ObjectAddressCollector *collector,
 }
 
 
+/*
+ * ApplyAddAllToDependencyList is an apply function for RecurseObjectDependencies that will
+ * collect all the ObjectAddresses for pg_depend entries to the context. This one will be
+ * used to provide meaningfull messages if an object has a dependency on extension owned
+ * object. You probably need ApplyAddToDependencyList, if you want to create dependencies
+ * for objects.
+ */
+static void
+ApplyAddAllToDependencyList(ObjectAddressCollector *collector,
+							DependencyDefinition *definition)
+{
+	ObjectAddress address = DependencyDefinitionObjectAddress(definition);
+	CollectObjectAddress(collector, &address);
+}
+
+
 /*
  * ExpandCitusSupportedTypes base on supported types by citus we might want to expand
  * the list of objects to visit in pg_depend.

src/test/regress/expected/function_propagation.out

@@ -283,23 +283,23 @@ SELECT * FROM run_command_on_workers($$SELECT pg_identify_object_as_address(clas
  localhost |    57638 | t       | (function,"{function_propagation_schema,max_of_table}",{})
 (2 rows)
 
+-- Check extension owned table
+CREATE TABLE extension_owned_table(a int);
+CREATE EXTENSION seg;
+ALTER EXTENSION seg ADD TABLE extension_owned_table;
+NOTICE:  Citus does not propagate adding/dropping member objects
+HINT:  You can add/drop the member objects on the workers as well.
+CREATE OR REPLACE FUNCTION func_for_ext_check(param_1 extension_owned_table)
+RETURNS int
+LANGUAGE plpgsql AS
+$$
+BEGIN
+    return 1;
+END;
+$$;
+WARNING:  Citus can't distribute function "func_for_ext_check" having dependency on non-distributed relation "extension_owned_table"
+DETAIL:  Function will be created only locally
+HINT:  To distribute function, distribute dependent relations first. Then, re-create the function
 RESET search_path;
+SET client_min_messages TO WARNING;
 DROP SCHEMA function_propagation_schema CASCADE;
-NOTICE:  drop cascades to 17 other objects
-DETAIL:  drop cascades to type function_propagation_schema.function_prop_type
-drop cascades to function function_propagation_schema.func_1(function_propagation_schema.function_prop_type)
-drop cascades to type function_propagation_schema.function_prop_type_2
-drop cascades to function function_propagation_schema.func_2(integer)
-drop cascades to type function_propagation_schema.function_prop_type_3
-drop cascades to function function_propagation_schema.func_3(integer)
-drop cascades to table function_propagation_schema.function_prop_table
-drop cascades to function function_propagation_schema.func_4(function_propagation_schema.function_prop_table)
-drop cascades to function function_propagation_schema.func_5(integer)
-drop cascades to function function_propagation_schema.func_6(function_propagation_schema.function_prop_table)
-drop cascades to view function_propagation_schema.function_prop_view
-drop cascades to function function_propagation_schema.func_7(function_propagation_schema.function_prop_view)
-drop cascades to function function_propagation_schema.func_8(integer)
-drop cascades to type function_propagation_schema.type_in_transaction
-drop cascades to function function_propagation_schema.func_in_transaction(function_propagation_schema.type_in_transaction)
-drop cascades to table function_propagation_schema.table_in_sql_body
-drop cascades to function function_propagation_schema.max_of_table()

src/test/regress/sql/function_propagation.sql

@@ -167,5 +167,20 @@ SELECT pg_identify_object_as_address(classid, objid, objsubid) from citus.pg_dis
 SELECT pg_identify_object_as_address(classid, objid, objsubid) from citus.pg_dist_object where objid = 'function_propagation_schema.max_of_table'::regproc::oid;
 SELECT * FROM run_command_on_workers($$SELECT pg_identify_object_as_address(classid, objid, objsubid) from citus.pg_dist_object where objid = 'function_propagation_schema.max_of_table'::regproc::oid;$$) ORDER BY 1,2;
 
+-- Check extension owned table
+CREATE TABLE extension_owned_table(a int);
+CREATE EXTENSION seg;
+ALTER EXTENSION seg ADD TABLE extension_owned_table;
+
+CREATE OR REPLACE FUNCTION func_for_ext_check(param_1 extension_owned_table)
+RETURNS int
+LANGUAGE plpgsql AS
+$$
+BEGIN
+    return 1;
+END;
+$$;
+
 RESET search_path;
+SET client_min_messages TO WARNING;
 DROP SCHEMA function_propagation_schema CASCADE;