Extend multi_mx_call with some of Hadi's suggestions for better test coverage

src/test/regress/expected/multi_mx_call.out

@@ -10,12 +10,36 @@ select create_distributed_table('mx_call_dist_table', 'id');
 (1 row)
 
 insert into mx_call_dist_table values (1),(2),(3),(4),(5);
+create type mx_call_enum as enum ('A', 'S', 'D', 'F');
 CREATE PROCEDURE mx_call_proc(x int, INOUT y int) LANGUAGE plpgsql AS $$
 BEGIN
     y := x + (select case groupid when 0 then 1 else 0 end from pg_dist_local_group);
     y := y + (select sum(id) from mx_call_dist_table);
 END;
 $$;
+CREATE PROCEDURE mx_call_proc_asdf(INOUT x mx_call_enum, INOUT y mx_call_enum) LANGUAGE plpgsql AS $$
+BEGIN
+    y := x;
+    x := (select case groupid when 0 then 'F' else 'S' end from pg_dist_local_group);
+END;
+$$;
+CREATE FUNCTION mx_call_add(int, int) RETURNS int
+    AS 'select $1 + $2;'
+    LANGUAGE SQL
+    IMMUTABLE
+    RETURNS NULL ON NULL INPUT;
+call mx_call_proc(2, 0);
+ y  
+----
+ 18
+(1 row)
+
+call mx_call_proc_asdf('S', 'A');
+ x | y 
+---+---
+ F | S
+(1 row)
+
 select create_distributed_function('mx_call_proc(int,int)');
  create_distributed_function 
 -----------------------------
@@ -26,12 +50,122 @@ update citus.pg_dist_object
 set distribution_argument_index = 1, colocationid = pg_dist_partition.colocationid
 from pg_proc, pg_dist_partition
 where proname = 'mx_call_proc' and oid = objid and pg_dist_partition.logicalrelid = 'mx_call_dist_table'::regclass;
+select create_distributed_function('mx_call_proc_asdf(mx_call_enum,mx_call_enum)');
+ create_distributed_function 
+-----------------------------
+ 
+(1 row)
+
+update citus.pg_dist_object
+set distribution_argument_index = 1, colocationid = pg_dist_partition.colocationid
+from pg_proc, pg_dist_partition
+where proname = 'mx_call_proc_asdf' and oid = objid and pg_dist_partition.logicalrelid = 'mx_call_dist_table'::regclass;
 call mx_call_proc(2, 0);
  y  
 ----
  17
 (1 row)
 
+call mx_call_proc_asdf('S', 'A');
+ x | y 
+---+---
+ S | S
+(1 row)
+
+set client_min_messages to DEBUG2;
+begin;
+select sum(id) from mx_call_dist_table;
+DEBUG:  Router planner cannot handle multi-shard select queries
+ sum 
+-----
+  15
+(1 row)
+
+call mx_call_proc(2, 0);
+DEBUG:  cannot push down CALL in multi-statement transaction
+DEBUG:  Router planner cannot handle multi-shard select queries
+CONTEXT:  SQL statement "SELECT y + (select sum(id) from mx_call_dist_table)"
+PL/pgSQL function mx_call_proc(integer,integer) line 4 at assignment
+DEBUG:  Router planner cannot handle multi-shard select queries
+CONTEXT:  SQL statement "SELECT y + (select sum(id) from mx_call_dist_table)"
+PL/pgSQL function mx_call_proc(integer,integer) line 4 at assignment
+DEBUG:  generating subplan 5_1 for subquery SELECT sum(id) AS sum FROM public.mx_call_dist_table
+CONTEXT:  SQL statement "SELECT y + (select sum(id) from mx_call_dist_table)"
+PL/pgSQL function mx_call_proc(integer,integer) line 4 at assignment
+DEBUG:  Plan 5 query after replacing subqueries and CTEs: SELECT (3 OPERATOR(pg_catalog.+) (SELECT intermediate_result.sum FROM read_intermediate_result('5_1'::text, 'binary'::citus_copy_format) intermediate_result(sum bigint)))
+CONTEXT:  SQL statement "SELECT y + (select sum(id) from mx_call_dist_table)"
+PL/pgSQL function mx_call_proc(integer,integer) line 4 at assignment
+DEBUG:  Creating router plan
+CONTEXT:  SQL statement "SELECT y + (select sum(id) from mx_call_dist_table)"
+PL/pgSQL function mx_call_proc(integer,integer) line 4 at assignment
+DEBUG:  Plan is router executable
+CONTEXT:  SQL statement "SELECT y + (select sum(id) from mx_call_dist_table)"
+PL/pgSQL function mx_call_proc(integer,integer) line 4 at assignment
+ y  
+----
+ 18
+(1 row)
+
+commit;
+update citus.pg_dist_object
+set distribution_argument_index = -1, colocationid = pg_dist_partition.colocationid
+from pg_proc, pg_dist_partition
+where proname = 'mx_call_proc' and oid = objid and pg_dist_partition.logicalrelid = 'mx_call_dist_table'::regclass;
+call mx_call_proc(2, 0);
+DEBUG:  cannot push down invalid distribution_argument_index
+DEBUG:  Router planner cannot handle multi-shard select queries
+CONTEXT:  SQL statement "SELECT y + (select sum(id) from mx_call_dist_table)"
+PL/pgSQL function mx_call_proc(integer,integer) line 4 at assignment
+DEBUG:  Router planner cannot handle multi-shard select queries
+CONTEXT:  SQL statement "SELECT y + (select sum(id) from mx_call_dist_table)"
+PL/pgSQL function mx_call_proc(integer,integer) line 4 at assignment
+DEBUG:  generating subplan 7_1 for subquery SELECT sum(id) AS sum FROM public.mx_call_dist_table
+CONTEXT:  SQL statement "SELECT y + (select sum(id) from mx_call_dist_table)"
+PL/pgSQL function mx_call_proc(integer,integer) line 4 at assignment
+DEBUG:  Plan 7 query after replacing subqueries and CTEs: SELECT (3 OPERATOR(pg_catalog.+) (SELECT intermediate_result.sum FROM read_intermediate_result('7_1'::text, 'binary'::citus_copy_format) intermediate_result(sum bigint)))
+CONTEXT:  SQL statement "SELECT y + (select sum(id) from mx_call_dist_table)"
+PL/pgSQL function mx_call_proc(integer,integer) line 4 at assignment
+DEBUG:  Creating router plan
+CONTEXT:  SQL statement "SELECT y + (select sum(id) from mx_call_dist_table)"
+PL/pgSQL function mx_call_proc(integer,integer) line 4 at assignment
+DEBUG:  Plan is router executable
+CONTEXT:  SQL statement "SELECT y + (select sum(id) from mx_call_dist_table)"
+PL/pgSQL function mx_call_proc(integer,integer) line 4 at assignment
+ y  
+----
+ 18
+(1 row)
+
+update citus.pg_dist_object
+set distribution_argument_index = 1, colocationid = pg_dist_partition.colocationid
+from pg_proc, pg_dist_partition
+where proname = 'mx_call_proc' and oid = objid and pg_dist_partition.logicalrelid = 'mx_call_dist_table'::regclass;
+call mx_call_proc(2, mx_call_add(3, 4));
+DEBUG:  cannot push down non-constant argument value
+DEBUG:  Router planner cannot handle multi-shard select queries
+CONTEXT:  SQL statement "SELECT y + (select sum(id) from mx_call_dist_table)"
+PL/pgSQL function mx_call_proc(integer,integer) line 4 at assignment
+DEBUG:  Router planner cannot handle multi-shard select queries
+CONTEXT:  SQL statement "SELECT y + (select sum(id) from mx_call_dist_table)"
+PL/pgSQL function mx_call_proc(integer,integer) line 4 at assignment
+DEBUG:  generating subplan 9_1 for subquery SELECT sum(id) AS sum FROM public.mx_call_dist_table
+CONTEXT:  SQL statement "SELECT y + (select sum(id) from mx_call_dist_table)"
+PL/pgSQL function mx_call_proc(integer,integer) line 4 at assignment
+DEBUG:  Plan 9 query after replacing subqueries and CTEs: SELECT (3 OPERATOR(pg_catalog.+) (SELECT intermediate_result.sum FROM read_intermediate_result('9_1'::text, 'binary'::citus_copy_format) intermediate_result(sum bigint)))
+CONTEXT:  SQL statement "SELECT y + (select sum(id) from mx_call_dist_table)"
+PL/pgSQL function mx_call_proc(integer,integer) line 4 at assignment
+DEBUG:  Creating router plan
+CONTEXT:  SQL statement "SELECT y + (select sum(id) from mx_call_dist_table)"
+PL/pgSQL function mx_call_proc(integer,integer) line 4 at assignment
+DEBUG:  Plan is router executable
+CONTEXT:  SQL statement "SELECT y + (select sum(id) from mx_call_dist_table)"
+PL/pgSQL function mx_call_proc(integer,integer) line 4 at assignment
+ y  
+----
+ 18
+(1 row)
+
+reset client_min_messages;
 DROP TABLE mx_call_dist_table;
 DROP PROCEDURE mx_call_proc;
 reset citus.shard_replication_factor;

src/test/regress/expected/multi_mx_call_0.out

@@ -10,6 +10,7 @@ select create_distributed_table('mx_call_dist_table', 'id');
 (1 row)
 
 insert into mx_call_dist_table values (1),(2),(3),(4),(5);
+create type mx_call_enum as enum ('A', 'S', 'D', 'F');
 CREATE PROCEDURE mx_call_proc(x int, INOUT y int) LANGUAGE plpgsql AS $$
 BEGIN
     y := x + (select case groupid when 0 then 1 else 0 end from pg_dist_local_group);
@@ -19,6 +20,28 @@ $$;
 ERROR:  syntax error at or near "PROCEDURE"
 LINE 1: CREATE PROCEDURE mx_call_proc(x int, INOUT y int) LANGUAGE p...
                ^
+CREATE PROCEDURE mx_call_proc_asdf(INOUT x mx_call_enum, INOUT y mx_call_enum) LANGUAGE plpgsql AS $$
+BEGIN
+    y := x;
+    x := (select case groupid when 0 then 'F' else 'S' end from pg_dist_local_group);
+END;
+$$;
+ERROR:  syntax error at or near "PROCEDURE"
+LINE 1: CREATE PROCEDURE mx_call_proc_asdf(INOUT x mx_call_enum, INO...
+               ^
+CREATE FUNCTION mx_call_add(int, int) RETURNS int
+    AS 'select $1 + $2;'
+    LANGUAGE SQL
+    IMMUTABLE
+    RETURNS NULL ON NULL INPUT;
+call mx_call_proc(2, 0);
+ERROR:  syntax error at or near "call"
+LINE 1: call mx_call_proc(2, 0);
+        ^
+call mx_call_proc_asdf('S', 'A');
+ERROR:  syntax error at or near "call"
+LINE 1: call mx_call_proc_asdf('S', 'A');
+        ^
 select create_distributed_function('mx_call_proc(int,int)');
 ERROR:  function "mx_call_proc(int,int)" does not exist
 LINE 1: select create_distributed_function('mx_call_proc(int,int)');
@@ -27,10 +50,53 @@ update citus.pg_dist_object
 set distribution_argument_index = 1, colocationid = pg_dist_partition.colocationid
 from pg_proc, pg_dist_partition
 where proname = 'mx_call_proc' and oid = objid and pg_dist_partition.logicalrelid = 'mx_call_dist_table'::regclass;
+select create_distributed_function('mx_call_proc_asdf(mx_call_enum,mx_call_enum)');
+ERROR:  function "mx_call_proc_asdf(mx_call_enum,mx_call_enum)" does not exist
+LINE 1: select create_distributed_function('mx_call_proc_asdf(mx_cal...
+                                           ^
+update citus.pg_dist_object
+set distribution_argument_index = 1, colocationid = pg_dist_partition.colocationid
+from pg_proc, pg_dist_partition
+where proname = 'mx_call_proc_asdf' and oid = objid and pg_dist_partition.logicalrelid = 'mx_call_dist_table'::regclass;
 call mx_call_proc(2, 0);
 ERROR:  syntax error at or near "call"
 LINE 1: call mx_call_proc(2, 0);
         ^
+call mx_call_proc_asdf('S', 'A');
+ERROR:  syntax error at or near "call"
+LINE 1: call mx_call_proc_asdf('S', 'A');
+        ^
+set client_min_messages to DEBUG2;
+begin;
+select sum(id) from mx_call_dist_table;
+DEBUG:  Router planner cannot handle multi-shard select queries
+ sum 
+-----
+  15
+(1 row)
+
+call mx_call_proc(2, 0);
+ERROR:  syntax error at or near "call"
+LINE 1: call mx_call_proc(2, 0);
+        ^
+commit;
+update citus.pg_dist_object
+set distribution_argument_index = -1, colocationid = pg_dist_partition.colocationid
+from pg_proc, pg_dist_partition
+where proname = 'mx_call_proc' and oid = objid and pg_dist_partition.logicalrelid = 'mx_call_dist_table'::regclass;
+call mx_call_proc(2, 0);
+ERROR:  syntax error at or near "call"
+LINE 1: call mx_call_proc(2, 0);
+        ^
+update citus.pg_dist_object
+set distribution_argument_index = 1, colocationid = pg_dist_partition.colocationid
+from pg_proc, pg_dist_partition
+where proname = 'mx_call_proc' and oid = objid and pg_dist_partition.logicalrelid = 'mx_call_dist_table'::regclass;
+call mx_call_proc(2, mx_call_add(3, 4));
+ERROR:  syntax error at or near "call"
+LINE 1: call mx_call_proc(2, mx_call_add(3, 4));
+        ^
+reset client_min_messages;
 DROP TABLE mx_call_dist_table;
 DROP PROCEDURE mx_call_proc;
 ERROR:  syntax error at or near "PROCEDURE"

src/test/regress/sql/multi_mx_call.sql

@@ -9,6 +9,8 @@ CREATE TABLE mx_call_dist_table(id int);
 select create_distributed_table('mx_call_dist_table', 'id');
 insert into mx_call_dist_table values (1),(2),(3),(4),(5);
 
+create type mx_call_enum as enum ('A', 'S', 'D', 'F');
+
 CREATE PROCEDURE mx_call_proc(x int, INOUT y int) LANGUAGE plpgsql AS $$
 BEGIN
     y := x + (select case groupid when 0 then 1 else 0 end from pg_dist_local_group);
@@ -16,14 +18,56 @@ BEGIN
 END;
 $$;
 
+CREATE PROCEDURE mx_call_proc_asdf(INOUT x mx_call_enum, INOUT y mx_call_enum) LANGUAGE plpgsql AS $$
+BEGIN
+    y := x;
+    x := (select case groupid when 0 then 'F' else 'S' end from pg_dist_local_group);
+END;
+$$;
+
+CREATE FUNCTION mx_call_add(int, int) RETURNS int
+    AS 'select $1 + $2;'
+    LANGUAGE SQL
+    IMMUTABLE
+    RETURNS NULL ON NULL INPUT;
+
+call mx_call_proc(2, 0);
+call mx_call_proc_asdf('S', 'A');
+
 select create_distributed_function('mx_call_proc(int,int)');
 update citus.pg_dist_object
 set distribution_argument_index = 1, colocationid = pg_dist_partition.colocationid
 from pg_proc, pg_dist_partition
 where proname = 'mx_call_proc' and oid = objid and pg_dist_partition.logicalrelid = 'mx_call_dist_table'::regclass;
 
-call mx_call_proc(2, 0);
+select create_distributed_function('mx_call_proc_asdf(mx_call_enum,mx_call_enum)');
+update citus.pg_dist_object
+set distribution_argument_index = 1, colocationid = pg_dist_partition.colocationid
+from pg_proc, pg_dist_partition
+where proname = 'mx_call_proc_asdf' and oid = objid and pg_dist_partition.logicalrelid = 'mx_call_dist_table'::regclass;
 
+call mx_call_proc(2, 0);
+call mx_call_proc_asdf('S', 'A');
+
+set client_min_messages to DEBUG2;
+begin;
+select sum(id) from mx_call_dist_table;
+call mx_call_proc(2, 0);
+commit;
+
+update citus.pg_dist_object
+set distribution_argument_index = -1, colocationid = pg_dist_partition.colocationid
+from pg_proc, pg_dist_partition
+where proname = 'mx_call_proc' and oid = objid and pg_dist_partition.logicalrelid = 'mx_call_dist_table'::regclass;
+call mx_call_proc(2, 0);
+update citus.pg_dist_object
+set distribution_argument_index = 1, colocationid = pg_dist_partition.colocationid
+from pg_proc, pg_dist_partition
+where proname = 'mx_call_proc' and oid = objid and pg_dist_partition.logicalrelid = 'mx_call_dist_table'::regclass;
+
+call mx_call_proc(2, mx_call_add(3, 4));
+
+reset client_min_messages;
 DROP TABLE mx_call_dist_table;
 DROP PROCEDURE mx_call_proc;
 reset citus.shard_replication_factor;