diff --git a/src/backend/distributed/planner/distributed_planner.c b/src/backend/distributed/planner/distributed_planner.c index bd019ee49..125c0236e 100644 --- a/src/backend/distributed/planner/distributed_planner.c +++ b/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(); diff --git a/src/backend/distributed/worker/worker_shard_visibility.c b/src/backend/distributed/worker/worker_shard_visibility.c index 1ad26fc37..c6e16a069 100644 --- a/src/backend/distributed/worker/worker_shard_visibility.c +++ b/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, - ReplaceTableVisibleFunctionWalker, NULL, 0); + return query_tree_walker((Query *) inputNode, ReplaceTableVisibleFunctionWalker, + NULL, 0); } - return expression_tree_mutator(inputNode, ReplaceTableVisibleFunctionWalker, NULL); + return expression_tree_walker(inputNode, ReplaceTableVisibleFunctionWalker, NULL); } diff --git a/src/include/distributed/worker_shard_visibility.h b/src/include/distributed/worker_shard_visibility.h index 7ada241d3..6fa63c9e2 100644 --- a/src/include/distributed/worker_shard_visibility.h +++ b/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 */ diff --git a/src/test/regress/expected/multi_mx_hide_shard_names.out b/src/test/regress/expected/multi_mx_hide_shard_names.out index 94311c805..94c2c7369 100644 --- a/src/test/regress/expected/multi_mx_hide_shard_names.out +++ b/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'; diff --git a/src/test/regress/sql/multi_mx_hide_shard_names.sql b/src/test/regress/sql/multi_mx_hide_shard_names.sql index b90139116..87a3b15d3 100644 --- a/src/test/regress/sql/multi_mx_hide_shard_names.sql +++ b/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';