Add skip_qualify_public param to shard_name() to allow qualifying for "public" schema (#8014)

DESCRIPTION: Adds skip_qualify_public param to `shard_name()` UDF to
allow qualifying for "public" schema when needed.

src/backend/distributed/relay/relay_event_utility.c

@@ -962,6 +962,7 @@ shard_name(PG_FUNCTION_ARGS)
 
 	Oid relationId = PG_GETARG_OID(0);
 	int64 shardId = PG_GETARG_INT64(1);
+	bool skipQualifyPublic = PG_GETARG_BOOL(2);
 
 	char *qualifiedName = NULL;
 
@@ -991,7 +992,7 @@ shard_name(PG_FUNCTION_ARGS)
 	Oid schemaId = get_rel_namespace(relationId);
 	char *schemaName = get_namespace_name(schemaId);
 
-	if (strncmp(schemaName, "public", NAMEDATALEN) == 0)
+	if (skipQualifyPublic && strncmp(schemaName, "public", NAMEDATALEN) == 0)
 	{
 		qualifiedName = (char *) quote_identifier(relationName);
 	}

src/backend/distributed/sql/citus--13.0-1--13.1-1.sql

@@ -51,3 +51,10 @@ DROP VIEW IF EXISTS pg_catalog.citus_lock_waits;
 #include "udfs/citus_stat_counters/13.1-1.sql"
 #include "udfs/citus_stat_counters_reset/13.1-1.sql"
 #include "udfs/citus_nodes/13.1-1.sql"
+
+-- Since shard_name/13.1-1.sql first drops the function and then creates it, we first
+-- need to drop citus_shards view since that view depends on this function. And immediately
+-- after creating the function, we recreate citus_shards view again.
+DROP VIEW pg_catalog.citus_shards;
+#include "udfs/shard_name/13.1-1.sql"
+#include "udfs/citus_shards/12.0-1.sql"

src/backend/distributed/sql/downgrades/citus--13.1-1--13.0-1.sql

@@ -46,3 +46,24 @@ DROP VIEW pg_catalog.citus_stat_counters;
 DROP FUNCTION pg_catalog.citus_stat_counters(oid);
 DROP FUNCTION pg_catalog.citus_stat_counters_reset(oid);
 DROP VIEW IF EXISTS pg_catalog.citus_nodes;
+
+-- Definition of shard_name() prior to this release doesn't have a separate SQL file
+-- because it's quite an old UDF that its prior definition(s) was(were) squashed into
+-- citus--8.0-1.sql. For this reason, to downgrade it, here we directly execute its old
+-- definition instead of including it from such a separate file.
+--
+-- And before dropping and creating the function, we also need to drop citus_shards view
+-- since it depends on it. And immediately after creating the function, we recreate
+-- citus_shards view again.
+
+DROP VIEW pg_catalog.citus_shards;
+
+DROP FUNCTION pg_catalog.shard_name(object_name regclass, shard_id bigint, skip_qualify_public boolean);
+CREATE FUNCTION pg_catalog.shard_name(object_name regclass, shard_id bigint)
+    RETURNS text
+    LANGUAGE C STABLE STRICT
+    AS 'MODULE_PATHNAME', $$shard_name$$;
+COMMENT ON FUNCTION pg_catalog.shard_name(object_name regclass, shard_id bigint)
+    IS 'returns schema-qualified, shard-extended identifier of object name';
+
+#include "../udfs/citus_shards/12.0-1.sql"

src/backend/distributed/sql/udfs/shard_name/13.1-1.sql

@@ -0,0 +1,8 @@
+-- skip_qualify_public is set to true by default just for backward compatibility
+DROP FUNCTION pg_catalog.shard_name(object_name regclass, shard_id bigint);
+CREATE FUNCTION pg_catalog.shard_name(object_name regclass, shard_id bigint, skip_qualify_public boolean DEFAULT true)
+    RETURNS text
+    LANGUAGE C STABLE STRICT
+    AS 'MODULE_PATHNAME', $$shard_name$$;
+COMMENT ON FUNCTION pg_catalog.shard_name(object_name regclass, shard_id bigint, skip_qualify_public boolean)
+    IS 'returns schema-qualified, shard-extended identifier of object name';

src/backend/distributed/sql/udfs/shard_name/latest.sql

@@ -0,0 +1,8 @@
+-- skip_qualify_public is set to true by default just for backward compatibility
+DROP FUNCTION pg_catalog.shard_name(object_name regclass, shard_id bigint);
+CREATE FUNCTION pg_catalog.shard_name(object_name regclass, shard_id bigint, skip_qualify_public boolean DEFAULT true)
+    RETURNS text
+    LANGUAGE C STABLE STRICT
+    AS 'MODULE_PATHNAME', $$shard_name$$;
+COMMENT ON FUNCTION pg_catalog.shard_name(object_name regclass, shard_id bigint, skip_qualify_public boolean)
+    IS 'returns schema-qualified, shard-extended identifier of object name';

src/test/regress/expected/citus_shards.out

@@ -33,5 +33,33 @@ SELECT * FROM citus_shards;
  t1             | 99456903 | citus_shards.t1_99456903             | distributed      |        456900 | localhost |    57638 |       8192
 (8 rows)
 
+SET search_path TO public;
+CREATE TABLE t3 (i int);
+SELECT citus_add_local_table_to_metadata('t3');
+ citus_add_local_table_to_metadata
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT shard_name('t3', shardid) FROM pg_dist_shard WHERE logicalrelid = 't3'::regclass;
+ shard_name
+---------------------------------------------------------------------
+ t3_99456908
+(1 row)
+
+SELECT shard_name('t3', shardid, true) FROM pg_dist_shard WHERE logicalrelid = 't3'::regclass;
+ shard_name
+---------------------------------------------------------------------
+ t3_99456908
+(1 row)
+
+SELECT shard_name('t3', shardid, false) FROM pg_dist_shard WHERE logicalrelid = 't3'::regclass;
+     shard_name
+---------------------------------------------------------------------
+ public.t3_99456908
+(1 row)
+
+DROP TABLE t3;
+SET search_path TO citus_shards;
 SET client_min_messages TO WARNING;
 DROP SCHEMA citus_shards CASCADE;

src/test/regress/expected/multi_extension.out

@@ -1455,6 +1455,7 @@ SELECT * FROM multi_extension.print_extension_changes();
                         previous_object                         |                                        current_object
 ---------------------------------------------------------------------
  function citus_unmark_object_distributed(oid,oid,integer) void |
+ function shard_name(regclass,bigint) text                      |
                                                                 | function citus_internal.acquire_citus_advisory_object_class_lock(integer,cstring) void
                                                                 | function citus_internal.add_colocation_metadata(integer,integer,integer,regtype,oid) void
                                                                 | function citus_internal.add_object_metadata(text,text[],text[],integer,integer,boolean) void
@@ -1483,9 +1484,10 @@ SELECT * FROM multi_extension.print_extension_changes();
                                                                 | function citus_stat_counters(oid) SETOF record
                                                                 | function citus_stat_counters_reset(oid) void
                                                                 | function citus_unmark_object_distributed(oid,oid,integer,boolean) void
+                                                                | function shard_name(regclass,bigint,boolean) text
                                                                 | view citus_nodes
                                                                 | view citus_stat_counters
-(31 rows)
+(33 rows)
 
 DROP TABLE multi_extension.prev_objects, multi_extension.extension_diff;
 -- show running version

src/test/regress/expected/upgrade_list_citus_objects.out

@@ -289,7 +289,7 @@ ORDER BY 1;
  function run_command_on_placements(regclass,text,boolean)
  function run_command_on_shards(regclass,text,boolean)
  function run_command_on_workers(text,boolean)
- function shard_name(regclass,bigint)
+ function shard_name(regclass,bigint,boolean)
  function start_metadata_sync_to_all_nodes()
  function start_metadata_sync_to_node(text,integer)
  function stop_metadata_sync_to_node(text,integer,boolean)

src/test/regress/sql/citus_shards.sql

@@ -13,5 +13,16 @@ INSERT INTO t1 SELECT generate_series(1, 100);
 INSERT INTO "t with space" SELECT generate_series(1, 1000);
 SELECT * FROM citus_shards;
 
+SET search_path TO public;
+CREATE TABLE t3 (i int);
+SELECT citus_add_local_table_to_metadata('t3');
+
+SELECT shard_name('t3', shardid) FROM pg_dist_shard WHERE logicalrelid = 't3'::regclass;
+SELECT shard_name('t3', shardid, true) FROM pg_dist_shard WHERE logicalrelid = 't3'::regclass;
+SELECT shard_name('t3', shardid, false) FROM pg_dist_shard WHERE logicalrelid = 't3'::regclass;
+
+DROP TABLE t3;
+SET search_path TO citus_shards;
+
 SET client_min_messages TO WARNING;
 DROP SCHEMA citus_shards CASCADE;