Change implementation of shard_name UDF to get schema-qualified shard name

src/backend/distributed/Makefile

@@ -11,7 +11,7 @@ EXTVERSIONS = 5.0 5.0-1 5.0-2  \
 	6.0-1 6.0-2 6.0-3 6.0-4 6.0-5 6.0-6 6.0-7 6.0-8 6.0-9 6.0-10 6.0-11 6.0-12 6.0-13 6.0-14 6.0-15 6.0-16 6.0-17 6.0-18 \
 	6.1-1 6.1-2 6.1-3 6.1-4 6.1-5 6.1-6 6.1-7 6.1-8 6.1-9 6.1-10 6.1-11 6.1-12 6.1-13 6.1-14 6.1-15 6.1-16 6.1-17 \
 	6.2-1 6.2-2 6.2-3 6.2-4 \
-	7.0-1
+	7.0-1 7.0-2
 
 # All citus--*.sql files in the source directory
 DATA = $(patsubst $(citus_abs_srcdir)/%.sql,%.sql,$(wildcard $(citus_abs_srcdir)/$(EXTENSION)--*--*.sql))
@@ -141,6 +141,8 @@ $(EXTENSION)--6.2-4.sql: $(EXTENSION)--6.2-3.sql $(EXTENSION)--6.2-3--6.2-4.sql
 	cat $^ > $@
 $(EXTENSION)--7.0-1.sql: $(EXTENSION)--6.2-4.sql $(EXTENSION)--6.2-4--7.0-1.sql
 	cat $^ > $@
+$(EXTENSION)--7.0-2.sql: $(EXTENSION)--7.0-1.sql $(EXTENSION)--7.0-1--7.0-2.sql
+	cat $^ > $@
 
 NO_PGXS = 1
 

src/backend/distributed/citus--7.0-1--7.0-2.sql

@@ -0,0 +1,9 @@
+/* citus--7.0-1--7.0-2.sql */
+
+/* redefine shard_name as STRICT */
+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';

src/backend/distributed/citus.control

@@ -1,6 +1,6 @@
 # Citus extension
 comment = 'Citus distributed database'
-default_version = '7.0-1'
+default_version = '7.0-2'
 module_pathname = '$libdir/citus'
 relocatable = false
 schema = pg_catalog

src/backend/distributed/relay/relay_event_utility.c

@@ -667,31 +667,18 @@ AppendShardIdToName(char **name, uint64 shardId)
 
 /*
  * shard_name() provides a PG function interface to AppendShardNameToId above.
+ * Returns the name of a shard as a quoted schema-qualified identifier.
  */
 Datum
 shard_name(PG_FUNCTION_ARGS)
 {
-	Oid relationId = InvalidOid;
+	Oid relationId = PG_GETARG_OID(0);
-	int64 shardId = 0;
+	int64 shardId = PG_GETARG_INT64(1);
 	char *relationName = NULL;
 
-	/*
+	Oid schemaId = InvalidOid;
-	 * Have to check arguments for NULLness as it can't be declared STRICT
+	char *schemaName = NULL;
-	 * because of min/max arguments, which have to be NULLable for new shards.
+	char *qualifiedName = NULL;
-	 */
-	if (PG_ARGISNULL(0))
-	{
-		ereport(ERROR, (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
-						errmsg("object_name cannot be null")));
-	}
-	if (PG_ARGISNULL(1))
-	{
-		ereport(ERROR, (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
-						errmsg("shard_id cannot be null")));
-	}
-
-	relationId = PG_GETARG_OID(0);
-	shardId = PG_GETARG_INT64(1);
 
 	CheckCitusVersion(ERROR);
 
@@ -717,5 +704,10 @@ shard_name(PG_FUNCTION_ARGS)
 	}
 
 	AppendShardIdToName(&relationName, shardId);
-	PG_RETURN_TEXT_P(cstring_to_text(relationName));
+
+	schemaId = get_rel_namespace(relationId);
+	schemaName = get_namespace_name(schemaId);
+	qualifiedName = quote_qualified_identifier(schemaName, relationName);
+
+	PG_RETURN_TEXT_P(cstring_to_text(qualifiedName));
 }

src/test/regress/expected/isolation_add_node_vs_reference_table_operations.out

@@ -187,7 +187,7 @@ step s2-print-index-count:
 	SELECT 
 		nodeport, success, result 
 	FROM 
-		run_command_on_placements('test_reference_table', 'select count(*) from pg_indexes WHERE tablename = ''%s''')
+		run_command_on_placements('test_reference_table', 'select count(*) from pg_indexes WHERE schemaname || ''.'' || tablename = ''%s''')
 	ORDER BY
 		nodeport;
 
@@ -227,7 +227,7 @@ step s2-print-index-count:
 	SELECT 
 		nodeport, success, result 
 	FROM 
-		run_command_on_placements('test_reference_table', 'select count(*) from pg_indexes WHERE tablename = ''%s''')
+		run_command_on_placements('test_reference_table', 'select count(*) from pg_indexes WHERE schemaname || ''.'' || tablename = ''%s''')
 	ORDER BY
 		nodeport;
 
@@ -409,7 +409,7 @@ step s2-print-index-count:
 	SELECT 
 		nodeport, success, result 
 	FROM 
-		run_command_on_placements('test_reference_table', 'select count(*) from pg_indexes WHERE tablename = ''%s''')
+		run_command_on_placements('test_reference_table', 'select count(*) from pg_indexes WHERE schemaname || ''.'' || tablename = ''%s''')
 	ORDER BY
 		nodeport;
 
@@ -446,7 +446,7 @@ step s2-print-index-count:
 	SELECT 
 		nodeport, success, result 
 	FROM 
-		run_command_on_placements('test_reference_table', 'select count(*) from pg_indexes WHERE tablename = ''%s''')
+		run_command_on_placements('test_reference_table', 'select count(*) from pg_indexes WHERE schemaname || ''.'' || tablename = ''%s''')
 	ORDER BY
 		nodeport;
 

src/test/regress/expected/isolation_add_node_vs_reference_table_operations_0.out

@@ -187,7 +187,7 @@ step s2-print-index-count:
 	SELECT 
 		nodeport, success, result 
 	FROM 
-		run_command_on_placements('test_reference_table', 'select count(*) from pg_indexes WHERE tablename = ''%s''')
+		run_command_on_placements('test_reference_table', 'select count(*) from pg_indexes WHERE schemaname || ''.'' || tablename = ''%s''')
 	ORDER BY
 		nodeport;
 
@@ -227,7 +227,7 @@ step s2-print-index-count:
 	SELECT 
 		nodeport, success, result 
 	FROM 
-		run_command_on_placements('test_reference_table', 'select count(*) from pg_indexes WHERE tablename = ''%s''')
+		run_command_on_placements('test_reference_table', 'select count(*) from pg_indexes WHERE schemaname || ''.'' || tablename = ''%s''')
 	ORDER BY
 		nodeport;
 
@@ -408,7 +408,7 @@ step s2-print-index-count:
 	SELECT 
 		nodeport, success, result 
 	FROM 
-		run_command_on_placements('test_reference_table', 'select count(*) from pg_indexes WHERE tablename = ''%s''')
+		run_command_on_placements('test_reference_table', 'select count(*) from pg_indexes WHERE schemaname || ''.'' || tablename = ''%s''')
 	ORDER BY
 		nodeport;
 
@@ -445,7 +445,7 @@ step s2-print-index-count:
 	SELECT 
 		nodeport, success, result 
 	FROM 
-		run_command_on_placements('test_reference_table', 'select count(*) from pg_indexes WHERE tablename = ''%s''')
+		run_command_on_placements('test_reference_table', 'select count(*) from pg_indexes WHERE schemaname || ''.'' || tablename = ''%s''')
 	ORDER BY
 		nodeport;
 

src/test/regress/expected/isolation_copy_placement_vs_modification.out

@@ -245,7 +245,7 @@ step s2-print-index-count:
 	SELECT 
 		nodeport, success, result 
 	FROM 
-		run_command_on_placements('test_table', 'select count(*) from pg_indexes WHERE tablename = ''%s''')
+		run_command_on_placements('test_table', 'select count(*) from pg_indexes WHERE schemaname || ''.'' || tablename = ''%s''')
 	ORDER BY
 		nodeport;
 
@@ -486,7 +486,7 @@ step s2-print-index-count:
 	SELECT 
 		nodeport, success, result 
 	FROM 
-		run_command_on_placements('test_table', 'select count(*) from pg_indexes WHERE tablename = ''%s''')
+		run_command_on_placements('test_table', 'select count(*) from pg_indexes WHERE schemaname || ''.'' || tablename = ''%s''')
 	ORDER BY
 		nodeport;
 

src/test/regress/expected/multi_extension.out

@@ -111,6 +111,7 @@ ALTER EXTENSION citus UPDATE TO '6.2-2';
 ALTER EXTENSION citus UPDATE TO '6.2-3';
 ALTER EXTENSION citus UPDATE TO '6.2-4';
 ALTER EXTENSION citus UPDATE TO '7.0-1';
+ALTER EXTENSION citus UPDATE TO '7.0-2';
 -- show running version
 SHOW citus.version;
  citus.version 

src/test/regress/expected/multi_name_lengths.out

@@ -31,17 +31,25 @@ SELECT master_create_worker_shards('too_long_12345678901234567890123456789012345
 \c - - - :master_port
 -- Verify that the UDF works and rejects bad arguments.
 SELECT shard_name(NULL, 666666);
-ERROR:  object_name cannot be null
+ shard_name 
+------------
+ 
+(1 row)
+
 SELECT shard_name(0, 666666);
 ERROR:  object_name does not reference a valid relation
 SELECT shard_name('too_long_12345678901234567890123456789012345678901234567890'::regclass, 666666);
                                shard_name                               
------------------------------------------------------------------
+------------------------------------------------------------------------
- too_long_12345678901234567890123456789012345678_e0119164_666666
+ public.too_long_12345678901234567890123456789012345678_e0119164_666666
 (1 row)
 
 SELECT shard_name('too_long_12345678901234567890123456789012345678901234567890'::regclass, NULL);
-ERROR:  shard_id cannot be null
+ shard_name 
+------------
+ 
+(1 row)
+
 SELECT shard_name('too_long_12345678901234567890123456789012345678901234567890'::regclass, -21);
 ERROR:  shard_id cannot be zero or negative value
 DROP TABLE too_long_12345678901234567890123456789012345678901234567890 CASCADE;
@@ -304,6 +312,15 @@ SELECT master_create_worker_shards(U&'elephant_!0441!043B!043E!043D!0441!043B!04
  
 (1 row)
 
+-- Verify that quoting is used in shard_name
+SELECT shard_name(U&'elephant_!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D' UESCAPE '!'::regclass, min(shardid))
+FROM pg_dist_shard
+WHERE logicalrelid = U&'elephant_!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D' UESCAPE '!'::regclass;
+                        shard_name                        
+----------------------------------------------------------
+ public."elephant_слонслонслонсло_c8b737c2_2250000000002"
+(1 row)
+
 \c - - - :worker_1_port
 \dt public.elephant_*
                               List of relations
@@ -322,6 +339,32 @@ SELECT master_create_worker_shards(U&'elephant_!0441!043B!043E!043D!0441!043B!04
 (2 rows)
 
 \c - - - :master_port
+-- Verify that shard_name UDF supports schemas
+CREATE SCHEMA multi_name_lengths;
+CREATE TABLE multi_name_lengths.too_long_12345678901234567890123456789012345678901234567890 (
+        col1 integer not null,
+        col2 integer not null);
+SELECT master_create_distributed_table('multi_name_lengths.too_long_12345678901234567890123456789012345678901234567890', 'col1', 'hash');
+ master_create_distributed_table 
+---------------------------------
+ 
+(1 row)
+
+SELECT master_create_worker_shards('multi_name_lengths.too_long_12345678901234567890123456789012345678901234567890', 2, 1);
+ master_create_worker_shards 
+-----------------------------
+ 
+(1 row)
+
+SELECT shard_name('multi_name_lengths.too_long_12345678901234567890123456789012345678901234567890'::regclass, min(shardid))
+FROM pg_dist_shard
+WHERE logicalrelid = 'multi_name_lengths.too_long_12345678901234567890123456789012345678901234567890'::regclass;
+                                     shard_name                                     
+------------------------------------------------------------------------------------
+ multi_name_lengths.too_long_1234567890123456789012345678901_e0119164_2250000000004
+(1 row)
+
+DROP TABLE multi_name_lengths.too_long_12345678901234567890123456789012345678901234567890;
 -- Clean up.
 DROP TABLE name_lengths CASCADE;
 DROP TABLE U&"elephant_!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D" UESCAPE '!' CASCADE;

src/test/regress/expected/multi_schema_support.out

@@ -1095,3 +1095,32 @@ SELECT run_command_on_coordinator_and_workers('DROP USER "test-user"');
 (1 row)
 
 DROP FUNCTION run_command_on_coordinator_and_workers(p_sql text);
+-- test run_command_on_* UDFs with schema
+CREATE SCHEMA run_test_schema;
+CREATE TABLE run_test_schema.test_table(id int);
+SELECT create_distributed_table('run_test_schema.test_table','id');
+ create_distributed_table 
+--------------------------
+ 
+(1 row)
+
+-- randomly insert data to evaluate below UDFs better
+INSERT INTO run_test_schema.test_table VALUES(1);
+INSERT INTO run_test_schema.test_table VALUES(7);
+INSERT INTO run_test_schema.test_table VALUES(9);
+-- try UDFs which call shard_name as a subroutine
+SELECT sum(result::int) FROM run_command_on_placements('run_test_schema.test_table','SELECT pg_table_size(''%s'')');
+  sum  
+-------
+ 49152
+(1 row)
+
+SELECT sum(result::int) FROM run_command_on_shards('run_test_schema.test_table','SELECT pg_table_size(''%s'')');
+  sum  
+-------
+ 24576
+(1 row)
+
+-- Clean up the created schema
+DROP SCHEMA run_test_schema CASCADE;
+NOTICE:  drop cascades to table run_test_schema.test_table

src/test/regress/specs/isolation_add_node_vs_reference_table_operations.spec

@@ -86,7 +86,7 @@ step "s2-print-index-count"
 	SELECT 
 		nodeport, success, result 
 	FROM 
-		run_command_on_placements('test_reference_table', 'select count(*) from pg_indexes WHERE tablename = ''%s''')
+		run_command_on_placements('test_reference_table', 'select count(*) from pg_indexes WHERE schemaname || ''.'' || tablename = ''%s''')
 	ORDER BY
 		nodeport;
 }

src/test/regress/specs/isolation_copy_placement_vs_modification.spec

@@ -104,7 +104,7 @@ step "s2-print-index-count"
 	SELECT 
 		nodeport, success, result 
 	FROM 
-		run_command_on_placements('test_table', 'select count(*) from pg_indexes WHERE tablename = ''%s''')
+		run_command_on_placements('test_table', 'select count(*) from pg_indexes WHERE schemaname || ''.'' || tablename = ''%s''')
 	ORDER BY
 		nodeport;
 }

src/test/regress/sql/multi_extension.sql

@@ -111,6 +111,7 @@ ALTER EXTENSION citus UPDATE TO '6.2-2';
 ALTER EXTENSION citus UPDATE TO '6.2-3';
 ALTER EXTENSION citus UPDATE TO '6.2-4';
 ALTER EXTENSION citus UPDATE TO '7.0-1';
+ALTER EXTENSION citus UPDATE TO '7.0-2';
 
 -- show running version
 SHOW citus.version;

src/test/regress/sql/multi_name_lengths.sql

@@ -150,11 +150,31 @@ CREATE TABLE U&"elephant_!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E
 SELECT master_create_distributed_table(U&'elephant_!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D' UESCAPE '!', 'col1', 'hash');
 SELECT master_create_worker_shards(U&'elephant_!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D' UESCAPE '!', '2', '2');
 
+-- Verify that quoting is used in shard_name
+SELECT shard_name(U&'elephant_!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D' UESCAPE '!'::regclass, min(shardid))
+FROM pg_dist_shard
+WHERE logicalrelid = U&'elephant_!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D' UESCAPE '!'::regclass;
+
 \c - - - :worker_1_port
 \dt public.elephant_*
 \di public.elephant_*
 \c - - - :master_port
 
+-- Verify that shard_name UDF supports schemas
+CREATE SCHEMA multi_name_lengths;
+CREATE TABLE multi_name_lengths.too_long_12345678901234567890123456789012345678901234567890 (
+        col1 integer not null,
+        col2 integer not null);
+SELECT master_create_distributed_table('multi_name_lengths.too_long_12345678901234567890123456789012345678901234567890', 'col1', 'hash');
+SELECT master_create_worker_shards('multi_name_lengths.too_long_12345678901234567890123456789012345678901234567890', 2, 1);
+
+SELECT shard_name('multi_name_lengths.too_long_12345678901234567890123456789012345678901234567890'::regclass, min(shardid))
+FROM pg_dist_shard
+WHERE logicalrelid = 'multi_name_lengths.too_long_12345678901234567890123456789012345678901234567890'::regclass;
+
+
+DROP TABLE multi_name_lengths.too_long_12345678901234567890123456789012345678901234567890;
+
 -- Clean up.
 DROP TABLE name_lengths CASCADE;
 DROP TABLE U&"elephant_!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D!0441!043B!043E!043D" UESCAPE '!' CASCADE;

src/test/regress/sql/multi_schema_support.sql

@@ -765,3 +765,20 @@ SELECT run_command_on_workers('DROP OWNED BY "test-user" CASCADE');
 SELECT run_command_on_coordinator_and_workers('DROP USER "test-user"');
 
 DROP FUNCTION run_command_on_coordinator_and_workers(p_sql text);
+
+-- test run_command_on_* UDFs with schema
+CREATE SCHEMA run_test_schema;
+CREATE TABLE run_test_schema.test_table(id int);
+SELECT create_distributed_table('run_test_schema.test_table','id');
+
+-- randomly insert data to evaluate below UDFs better
+INSERT INTO run_test_schema.test_table VALUES(1);
+INSERT INTO run_test_schema.test_table VALUES(7);
+INSERT INTO run_test_schema.test_table VALUES(9);
+
+-- try UDFs which call shard_name as a subroutine
+SELECT sum(result::int) FROM run_command_on_placements('run_test_schema.test_table','SELECT pg_table_size(''%s'')');
+SELECT sum(result::int) FROM run_command_on_shards('run_test_schema.test_table','SELECT pg_table_size(''%s'')');
+
+-- Clean up the created schema
+DROP SCHEMA run_test_schema CASCADE;