Add skip_qualifying_public param to shard_name() to allow qualifying for "public" schema

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,26 @@ 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. The only difference
+-- between below CREATE command and the one in citus--8.0-1.sql is that this one has
+-- ".. OR REPLACE .." clause to be able to replace its existing definition.
+--
+-- 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_qualifying_public);
+CREATE OR REPLACE 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_qualifying_public is set to true by default just for backward compatibility
+DROP FUNCTION pg_catalog.shard_name(object_name regclass, shard_id bigint);
+CREATE OR REPLACE FUNCTION pg_catalog.shard_name(object_name regclass, shard_id bigint, skip_qualifying_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_qualifying_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_qualifying_public is set to true by default just for backward compatibility
+DROP FUNCTION pg_catalog.shard_name(object_name regclass, shard_id bigint);
+CREATE OR REPLACE FUNCTION pg_catalog.shard_name(object_name regclass, shard_id bigint, skip_qualifying_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_qualifying_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/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;