diff --git a/src/backend/distributed/test/hide_shards.c b/src/backend/distributed/test/hide_shards.c new file mode 100644 index 000000000..59e738c36 --- /dev/null +++ b/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(); +} diff --git a/src/backend/distributed/worker/worker_shard_visibility.c b/src/backend/distributed/worker/worker_shard_visibility.c index 38fd93eca..da9c87a22 100644 --- a/src/backend/distributed/worker/worker_shard_visibility.c +++ b/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. 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 3bfaef545..c0265b282 100644 --- a/src/test/regress/expected/multi_mx_hide_shard_names.out +++ b/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 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 3d9a7669f..b56329150 100644 --- a/src/test/regress/sql/multi_mx_hide_shard_names.sql +++ b/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