Only hide shards from client backends and pg bg workers

The aim of hiding shards is to hide shards from client applications.

Certain bg workers (such as pg_cron or Citus maintanince daemon)
should be treated like client applications because users can run
queries from such bg workers. And, these bg workers should follow
the similar application_name checks as client backeends.

Certain other bg workers, such as logical replication or postgres'
parallel workers, should never hide shards. They are internal
operations.

Similarly the other backend types like the walsender or
checkpointer or autovacuum should never hide shards.

(cherry picked from commit 9043a1ed3f)

src/backend/distributed/test/hide_shards.c

@@ -0,0 +1,40 @@
+/*-------------------------------------------------------------------------
+ *
+ * hide_shards.c
+ *
+ * This file contains functions to provide helper UDFs for hiding
+ * shards from the applications.
+ *
+ * Copyright (c) Citus Data, Inc.
+ *
+ *-------------------------------------------------------------------------
+ */
+#include "postgres.h"
+#include "funcapi.h"
+#include "miscadmin.h"
+#include "pgstat.h"
+
+#include "distributed/metadata_utility.h"
+#include "distributed/worker_shard_visibility.h"
+
+
+PG_FUNCTION_INFO_V1(set_backend_type);
+
+/*
+ * set_backend_type is an external API to set the MyBackendType and
+ * re-checks the shard visibility.
+ */
+Datum
+set_backend_type(PG_FUNCTION_ARGS)
+{
+	EnsureSuperUser();
+
+	MyBackendType = PG_GETARG_INT32(0);
+
+	elog(NOTICE, "backend type switched to: %s",
+		 GetBackendTypeDesc(MyBackendType));
+
+	ResetHideShardsDecision();
+
+	PG_RETURN_VOID();
+}

src/backend/distributed/worker/worker_shard_visibility.c

@@ -8,6 +8,7 @@
  */
 
 #include "postgres.h"
+#include "miscadmin.h"
 
 #include "catalog/index.h"
 #include "catalog/namespace.h"
@@ -47,6 +48,7 @@ static HideShardsMode HideShards = CHECK_APPLICATION_NAME;
 
 static bool ShouldHideShards(void);
 static bool ShouldHideShardsInternal(void);
+static bool IsPgBgWorker(void);
 static bool FilterShardsFromPgclass(Node *node, void *context);
 static Node * CreateRelationIsAKnownShardFilter(int pgClassVarno);
 
@@ -334,6 +336,28 @@ ResetHideShardsDecision(void)
 static bool
 ShouldHideShardsInternal(void)
 {
+	if (MyBackendType == B_BG_WORKER)
+	{
+		if (IsPgBgWorker())
+		{
+			/*
+			 * If a background worker belongs to Postgres, we should
+			 * never hide shards. For other background workers, enforce
+			 * the application_name check below.
+			 */
+			return false;
+		}
+	}
+	else if (MyBackendType != B_BACKEND)
+	{
+		/*
+		 * We are aiming only to hide shards from client
+		 * backends or certain background workers(see above),
+		 * not backends like walsender or checkpointer.
+		 */
+		return false;
+	}
+
 	if (IsCitusInternalBackend() || IsRebalancerInternalBackend())
 	{
 		/* we never hide shards from Citus */
@@ -372,6 +396,24 @@ ShouldHideShardsInternal(void)
 }
 
 
+/*
+ * IsPgBgWorker returns true if the current background worker
+ * belongs to Postgres.
+ */
+static bool
+IsPgBgWorker(void)
+{
+	Assert(MyBackendType == B_BG_WORKER);
+
+	if (MyBgworkerEntry)
+	{
+		return strcmp(MyBgworkerEntry->bgw_library_name, "postgres") == 0;
+	}
+
+	return false;
+}
+
+
 /*
  * FilterShardsFromPgclass adds a NOT relation_is_a_known_shard(oid) filter
  * to the security quals of pg_class RTEs.

src/test/regress/expected/multi_mx_hide_shard_names.out

@@ -403,6 +403,86 @@ SELECT * FROM citus_shard_indexes_on_worker WHERE "Schema" = 'CiTuS.TeeN' ORDER
  CiTuS.TeeN | MyTenantIndex | index | postgres | TeeNTabLE.1!?!
 (1 row)
 
+\c - - - :worker_1_port
+-- re-connect to the worker node and show that only
+-- client backends can filter shards
+SET search_path TO "CiTuS.TeeN";
+-- Create the necessary test utility function
+SET citus.enable_metadata_sync TO off;
+CREATE OR REPLACE FUNCTION set_backend_type(backend_type int)
+    RETURNS void
+    LANGUAGE C STRICT
+    AS 'citus';
+RESET citus.enable_metadata_sync;
+-- the shards and indexes do not show up
+SELECT relname FROM pg_catalog.pg_class WHERE relnamespace = 'mx_hide_shard_names'::regnamespace ORDER BY relname;
+       relname
+---------------------------------------------------------------------
+ test_index
+ test_table
+ test_table_102008
+ test_table_2_1130000
+(4 rows)
+
+-- say, we set it to bgworker
+-- the shards and indexes do not show up
+SELECT set_backend_type(4);
+NOTICE:  backend type switched to: background worker
+ set_backend_type
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT relname FROM pg_catalog.pg_class WHERE relnamespace = 'mx_hide_shard_names'::regnamespace ORDER BY relname;
+       relname
+---------------------------------------------------------------------
+ test_index
+ test_table
+ test_table_102008
+ test_table_2_1130000
+(4 rows)
+
+-- or, we set it to walsender
+-- the shards and indexes do show up
+SELECT set_backend_type(9);
+NOTICE:  backend type switched to: walsender
+ set_backend_type
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT relname FROM pg_catalog.pg_class WHERE relnamespace = 'mx_hide_shard_names'::regnamespace ORDER BY relname;
+          relname
+---------------------------------------------------------------------
+ test_index
+ test_index_1130000
+ test_index_1130002
+ test_table
+ test_table_102008
+ test_table_102008_1130004
+ test_table_102008_1130006
+ test_table_1130000
+ test_table_1130002
+ test_table_2_1130000
+(10 rows)
+
+-- but, client backends to see the shards
+SELECT set_backend_type(3);
+NOTICE:  backend type switched to: client backend
+ set_backend_type
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT relname FROM pg_catalog.pg_class WHERE relnamespace = 'mx_hide_shard_names'::regnamespace ORDER BY relname;
+       relname
+---------------------------------------------------------------------
+ test_index
+ test_table
+ test_table_102008
+ test_table_2_1130000
+(4 rows)
+
 -- clean-up
 \c - - - :master_port
 -- show that common psql functions do not show shards

src/test/regress/sql/multi_mx_hide_shard_names.sql

@@ -207,6 +207,38 @@ SELECT * FROM citus_shard_indexes_on_worker WHERE "Schema" = 'CiTuS.TeeN' ORDER
 \d
 \di
 
+
+\c - - - :worker_1_port
+-- re-connect to the worker node and show that only
+-- client backends can filter shards
+SET search_path TO "CiTuS.TeeN";
+
+-- Create the necessary test utility function
+SET citus.enable_metadata_sync TO off;
+CREATE OR REPLACE FUNCTION set_backend_type(backend_type int)
+    RETURNS void
+    LANGUAGE C STRICT
+    AS 'citus';
+RESET citus.enable_metadata_sync;
+
+-- the shards and indexes do not show up
+SELECT relname FROM pg_catalog.pg_class WHERE relnamespace = 'mx_hide_shard_names'::regnamespace ORDER BY relname;
+
+-- say, we set it to bgworker
+-- the shards and indexes do not show up
+SELECT set_backend_type(4);
+SELECT relname FROM pg_catalog.pg_class WHERE relnamespace = 'mx_hide_shard_names'::regnamespace ORDER BY relname;
+
+-- or, we set it to walsender
+-- the shards and indexes do show up
+SELECT set_backend_type(9);
+SELECT relname FROM pg_catalog.pg_class WHERE relnamespace = 'mx_hide_shard_names'::regnamespace ORDER BY relname;
+
+-- but, client backends to see the shards
+SELECT set_backend_type(3);
+SELECT relname FROM pg_catalog.pg_class WHERE relnamespace = 'mx_hide_shard_names'::regnamespace ORDER BY relname;
+
+
 -- clean-up
 \c - - - :master_port