Merge pull request #2386 from citusdata/improve_walker

Use tree walker instead of mutator in relation visibility

src/backend/distributed/planner/distributed_planner.c

@@ -134,7 +134,7 @@ distributed_planner(Query *parse, int cursorOptions, ParamListInfo boundParams)
 	 * Make sure that we hide shard names on the Citus MX worker nodes. See comments in
 	 * ReplaceTableVisibleFunction() for the details.
 	 */
-	parse = (Query *) ReplaceTableVisibleFunction((Node *) parse);
+	ReplaceTableVisibleFunction((Node *) parse);
 
 	/* create a restriction context and put it at the end if context list */
 	plannerRestrictionContext = CreateAndPushPlannerRestrictionContext();

src/backend/distributed/worker/worker_shard_visibility.c

@@ -24,7 +24,7 @@
 bool OverrideTableVisibility = true;
 
 static bool RelationIsAKnownShard(Oid shardRelationId);
-static Node * ReplaceTableVisibleFunctionWalker(Node *inputNode);
+static bool ReplaceTableVisibleFunctionWalker(Node *inputNode);
 
 PG_FUNCTION_INFO_V1(citus_table_is_visible);
 PG_FUNCTION_INFO_V1(relation_is_a_known_shard);
@@ -189,16 +189,16 @@ RelationIsAKnownShard(Oid shardRelationId)
  * The replace functionality can be enabled/disable via a GUC. This function also
  * ensures that the extension is loaded and the version is compatible.
  */
-Node *
+void
 ReplaceTableVisibleFunction(Node *inputNode)
 {
 	if (!OverrideTableVisibility ||
 		!CitusHasBeenLoaded() || !CheckCitusVersion(DEBUG2))
 	{
-		return inputNode;
+		return;
 	}
 
-	return ReplaceTableVisibleFunctionWalker(inputNode);
+	ReplaceTableVisibleFunctionWalker(inputNode);
 }
 
 
@@ -210,13 +210,19 @@ ReplaceTableVisibleFunction(Node *inputNode)
  * Note that the only difference between the functions is that
  * the latter filters the tables that are known to be shards on
  * Citus MX worker (data) nodes.
+ *
+ * Note that although the function mutates the input node, we
+ * prefer to use query_tree_walker/expression_tree_walker over
+ * their mutator equivalents. This is safe because we ensure that
+ * the replaced function has the exact same input/output values with
+ * its precedent.
  */
-static Node *
+static bool
 ReplaceTableVisibleFunctionWalker(Node *inputNode)
 {
 	if (inputNode == NULL)
 	{
-		return NULL;
+		return false;
 	}
 
 	if (IsA(inputNode, FuncExpr))
@@ -235,14 +241,16 @@ ReplaceTableVisibleFunctionWalker(Node *inputNode)
 			 */
 			functionToProcess->funcid = CitusTableVisibleFuncId();
 
-			return (Node *) functionToProcess;
+			/* although not very likely, we could have nested calls to pg_table_is_visible */
+			return expression_tree_walker(inputNode, ReplaceTableVisibleFunctionWalker,
+										  NULL);
 		}
 	}
 	else if (IsA(inputNode, Query))
 	{
-		return (Node *) query_tree_mutator((Query *) inputNode,
+		return query_tree_walker((Query *) inputNode, ReplaceTableVisibleFunctionWalker,
-										   ReplaceTableVisibleFunctionWalker, NULL, 0);
+								 NULL, 0);
 	}
 
-	return expression_tree_mutator(inputNode, ReplaceTableVisibleFunctionWalker, NULL);
+	return expression_tree_walker(inputNode, ReplaceTableVisibleFunctionWalker, NULL);
 }

src/include/distributed/worker_shard_visibility.h

@@ -16,7 +16,7 @@
 extern bool OverrideTableVisibility;
 
 
-extern Node * ReplaceTableVisibleFunction(Node *inputNode);
+extern void ReplaceTableVisibleFunction(Node *inputNode);
 
 
 #endif /* WORKER_SHARD_VISIBILITY_H */

src/test/regress/expected/multi_mx_hide_shard_names.out

@@ -72,6 +72,23 @@ SELECT * FROM citus_shard_indexes_on_worker ORDER BY 2;
 --------+------+------+-------+-------
 (0 rows)
 
+-- also show that nested calls to pg_table_is_visible works fine
+-- if both of the calls to the pg_table_is_visible haven't been
+-- replaced, we would get 0 rows in the output
+SELECT 
+	pg_table_is_visible((SELECT 
+								"t1"."Name"::regclass 
+						 FROM 
+						 	citus_shards_on_worker as t1 
+						 WHERE 
+						 	NOT pg_table_is_visible("t1"."Name"::regclass) 
+						 LIMIT 
+						 	1));
+ pg_table_is_visible 
+---------------------
+ f
+(1 row)
+
 -- now create an index
 \c - - - :master_port
 SET search_path TO 'mx_hide_shard_names';

src/test/regress/sql/multi_mx_hide_shard_names.sql

@@ -43,6 +43,19 @@ SET search_path TO 'mx_hide_shard_names';
 SELECT * FROM citus_shards_on_worker ORDER BY 2;
 SELECT * FROM citus_shard_indexes_on_worker ORDER BY 2;
 
+-- also show that nested calls to pg_table_is_visible works fine
+-- if both of the calls to the pg_table_is_visible haven't been
+-- replaced, we would get 0 rows in the output
+SELECT 
+	pg_table_is_visible((SELECT 
+								"t1"."Name"::regclass 
+						 FROM 
+						 	citus_shards_on_worker as t1 
+						 WHERE 
+						 	NOT pg_table_is_visible("t1"."Name"::regclass) 
+						 LIMIT 
+						 	1));
+
 -- now create an index
 \c - - - :master_port
 SET search_path TO 'mx_hide_shard_names';