Fix comments and add handle INSERTs on pg_class correctly

src/backend/distributed/worker/worker_shard_visibility.c

@@ -54,6 +54,7 @@ static bool ShouldHideShardsInternal(void);
 static bool IsPgBgWorker(void);
 static bool FilterShardsFromPgclass(Node *node, void *context);
 static Node * CreateRelationIsAKnownShardFilter(int pgClassVarno);
+static bool HasRangeTableRef(Node *node, int *varno);
 
 PG_FUNCTION_INFO_V1(citus_table_is_visible);
 PG_FUNCTION_INFO_V1(relation_is_a_known_shard);
@@ -421,8 +422,8 @@ IsPgBgWorker(void)
 
 
 /*
- * FilterShardsFromPgclass adds a NOT relation_is_a_known_shard(oid) filter
+ * FilterShardsFromPgclass adds a "relation_is_a_known_shard(oid) IS NOT TRUE"
- * to the security quals of pg_class RTEs.
+ * filter to the quals of queries that query pg_class.
  */
 static bool
 FilterShardsFromPgclass(Node *node, void *context)
@@ -456,6 +457,17 @@ FilterShardsFromPgclass(Node *node, void *context)
 				continue;
 			}
 
+			/*
+			 * Skip if pg_class is not actually queried. This is possible on
+			 * INSERT statements that insert into pg_class.
+			 */
+			if (!expression_tree_walker((Node *) query->jointree->fromlist,
+										HasRangeTableRef, &varno))
+			{
+				/* the query references pg_class */
+				continue;
+			}
+
 			/* make sure the expression is in the right memory context */
 			MemoryContext originalContext = MemoryContextSwitchTo(queryContext);
 
@@ -485,6 +497,23 @@ FilterShardsFromPgclass(Node *node, void *context)
 }
 
 
+/*
+ * HasRangeTableRef passed to expression_tree_walker to check if a node is a
+ * RangeTblRef of the given varno is present in a fromlist.
+ */
+static bool
+HasRangeTableRef(Node *node, int *varno)
+{
+	if (IsA(node, RangeTblRef))
+	{
+		RangeTblRef *rangeTblRef = (RangeTblRef *) node;
+		return rangeTblRef->rtindex == *varno;
+	}
+
+	return expression_tree_walker(node, HasRangeTableRef, varno);
+}
+
+
 /*
  * CreateRelationIsAKnownShardFilter constructs an expression of the form:
  * pg_catalog.relation_is_a_known_shard(oid) IS NOT TRUE

src/test/regress/expected/multi_mx_hide_shard_names.out

@@ -83,13 +83,52 @@ SELECT relname FROM pg_catalog.pg_class WHERE relnamespace = 'mx_hide_shard_name
  test_table
 (1 row)
 
--- EVen when using subquery and having no existing quals on pg_clcass
+-- Even when using subquery and having no existing quals on pg_clcass
 SELECT relname FROM (SELECT relname, relnamespace FROM pg_catalog.pg_class) AS q WHERE relnamespace = 'mx_hide_shard_names'::regnamespace ORDER BY relname;
   relname
 ---------------------------------------------------------------------
  test_table
 (1 row)
 
+-- Check that inserts into pg_class don't add the filter
+EXPLAIN (COSTS OFF) INSERT INTO pg_class VALUES (1);
+     QUERY PLAN
+---------------------------------------------------------------------
+ Insert on pg_class
+   ->  Result
+(2 rows)
+
+-- Unless it's an INSERT SELECT that queries from pg_class;
+EXPLAIN (COSTS OFF) INSERT INTO pg_class SELECT * FROM pg_class;
+                          QUERY PLAN
+---------------------------------------------------------------------
+ Insert on pg_class
+   ->  Seq Scan on pg_class pg_class_1
+         Filter: (relation_is_a_known_shard(oid) IS NOT TRUE)
+(3 rows)
+
+-- Check that query that psql "\d test_table" does gets optimized to an index
+-- scan
+EXPLAIN (COSTS OFF) SELECT c.oid,
+  n.nspname,
+  c.relname
+FROM pg_catalog.pg_class c
+     LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
+WHERE c.relname OPERATOR(pg_catalog.~) '^(test_table)$' COLLATE pg_catalog.default
+  AND pg_catalog.pg_table_is_visible(c.oid)
+ORDER BY 2, 3;
+                                                                QUERY PLAN
+---------------------------------------------------------------------
+ Sort
+   Sort Key: n.nspname, c.relname
+   ->  Nested Loop Left Join
+         Join Filter: (n.oid = c.relnamespace)
+         ->  Index Scan using pg_class_relname_nsp_index on pg_class c
+               Index Cond: (relname = 'test_table'::text)
+               Filter: ((relname ~ '^(test_table)$'::text) AND (relation_is_a_known_shard(oid) IS NOT TRUE) AND pg_table_is_visible(oid))
+         ->  Seq Scan on pg_namespace n
+(8 rows)
+
 commit prepared 'take-aggressive-lock';
 -- now create an index
 \c - - - :master_port

src/test/regress/sql/multi_mx_hide_shard_names.sql

@@ -50,9 +50,25 @@ prepare transaction 'take-aggressive-lock';
 
 -- shards are hidden when using psql as application_name
 SELECT relname FROM pg_catalog.pg_class WHERE relnamespace = 'mx_hide_shard_names'::regnamespace ORDER BY relname;
--- EVen when using subquery and having no existing quals on pg_clcass
+-- Even when using subquery and having no existing quals on pg_clcass
 SELECT relname FROM (SELECT relname, relnamespace FROM pg_catalog.pg_class) AS q WHERE relnamespace = 'mx_hide_shard_names'::regnamespace ORDER BY relname;
 
+-- Check that inserts into pg_class don't add the filter
+EXPLAIN (COSTS OFF) INSERT INTO pg_class VALUES (1);
+-- Unless it's an INSERT SELECT that queries from pg_class;
+EXPLAIN (COSTS OFF) INSERT INTO pg_class SELECT * FROM pg_class;
+
+-- Check that query that psql "\d test_table" does gets optimized to an index
+-- scan
+EXPLAIN (COSTS OFF) SELECT c.oid,
+  n.nspname,
+  c.relname
+FROM pg_catalog.pg_class c
+     LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
+WHERE c.relname OPERATOR(pg_catalog.~) '^(test_table)$' COLLATE pg_catalog.default
+  AND pg_catalog.pg_table_is_visible(c.oid)
+ORDER BY 2, 3;
+
 commit prepared 'take-aggressive-lock';
 
 -- now create an index