Merge pull request #5869 from citusdata/do_not_hide_shards_from_pg

Only hide shards from client backends or non-pg background workers

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