Fix reshuffle replicated table test cases

For replicated table, we have to examine the data on each segment
to make sure reshuffle successfully.

This commit defines a UDF using plpythonu, in that UDF, the python
code will connect specific segment in utility mode to access the  replicated
table.

src/test/regress/expected/reshuffle.out

@@ -464,8 +464,65 @@ Select count(*) > 0 from r1_reshuffle where gp_segment_id=2;
 
 drop table r1_reshuffle;
 -- Replicated tables
+-- We have to make sure replicated table successfully reshuffled.
+-- Currently we could only connect to the specific segments in utility mode
+-- to see if the data is consistent there. We create some python function here.
+-- The case can only work under the assumption: it's running on 3-segment cluster
+-- in a single machine.
+-- start_ignore
+drop language plpythonu;
+ERROR:  language "plpythonu" does not exist
+-- end_ignore
+create language plpythonu;
+create function update_on_segment(tabname text, segid int, numseg int) returns boolean as
+$$
+import pygresql.pg as pg
+conn = pg.connect(dbname='regression')
+port = conn.query("select port from gp_segment_configuration where content = %d and role = 'p'" % segid).getresult()[0][0]
+conn.close()
+
+conn = pg.connect(dbname='regression', opt='-c gp_session_role=utility', port=port)
+conn.query("set allow_system_table_mods = true")
+conn.query("update gp_distribution_policy set numsegments = %d where localoid = '%s'::regclass" % (numseg, tabname))
+conn.close()
+
+return True
+$$
+LANGUAGE plpythonu;
+create function select_on_segment(sql text, segid int) returns int as
+$$
+import pygresql.pg as pg
+conn = pg.connect(dbname='regression')
+port = conn.query("select port from gp_segment_configuration where content = %d and role = 'p'" % segid).getresult()[0][0]
+conn.close()
+
+conn = pg.connect(dbname='regression', opt='-c gp_session_role=utility', port=port)
+r = conn.query(sql).getresult()
+conn.close()
+
+return r[0][0]
+$$
+LANGUAGE plpythonu;
 Create table r1_reshuffle(a int, b int, c int) distributed replicated;
 update gp_distribution_policy  set numsegments=1 where localoid='r1_reshuffle'::regclass;
+select update_on_segment('r1_reshuffle', 0, 1);
+ update_on_segment 
+-------------------
+ t
+(1 row)
+
+select update_on_segment('r1_reshuffle', 1, 1);
+ update_on_segment 
+-------------------
+ t
+(1 row)
+
+select update_on_segment('r1_reshuffle', 2, 1);
+ update_on_segment 
+-------------------
+ t
+(1 row)
+
 insert into r1_reshuffle select i,i,0 from generate_series(1,100) I;
 Select count(*) from r1_reshuffle;
  count 
@@ -473,16 +530,52 @@ Select count(*) from r1_reshuffle;
    100
 (1 row)
 
+Select select_on_segment('Select count(*) from r1_reshuffle;', 1);
+ select_on_segment 
+-------------------
+                 0
+(1 row)
+
+Select select_on_segment('Select count(*) from r1_reshuffle;', 2);
+ select_on_segment 
+-------------------
+                 0
+(1 row)
+
 Alter table r1_reshuffle set with (reshuffle);
-Select count(*) from r1_reshuffle;
- count 
--------
-   100
+Select select_on_segment('Select count(*) from r1_reshuffle;', 1);
+ select_on_segment 
+-------------------
+               100
+(1 row)
+
+Select select_on_segment('Select count(*) from r1_reshuffle;', 2);
+ select_on_segment 
+-------------------
+               100
 (1 row)
 
 drop table r1_reshuffle;
 Create table r1_reshuffle(a int, b int, c int) distributed replicated;
 update gp_distribution_policy  set numsegments=2 where localoid='r1_reshuffle'::regclass;
+select update_on_segment('r1_reshuffle', 0, 2);
+ update_on_segment 
+-------------------
+ t
+(1 row)
+
+select update_on_segment('r1_reshuffle', 1, 2);
+ update_on_segment 
+-------------------
+ t
+(1 row)
+
+select update_on_segment('r1_reshuffle', 2, 2);
+ update_on_segment 
+-------------------
+ t
+(1 row)
+
 insert into r1_reshuffle select i,i,0 from generate_series(1,100) I;
 Select count(*) from r1_reshuffle;
  count 
@@ -490,6 +583,12 @@ Select count(*) from r1_reshuffle;
    100
 (1 row)
 
+Select select_on_segment('Select count(*) from r1_reshuffle;', 2);
+ select_on_segment 
+-------------------
+                 0
+(1 row)
+
 Alter table r1_reshuffle set with (reshuffle);
 Select count(*) from r1_reshuffle;
  count 
@@ -497,6 +596,12 @@ Select count(*) from r1_reshuffle;
    100
 (1 row)
 
+Select select_on_segment('Select count(*) from r1_reshuffle;', 2);
+ select_on_segment 
+-------------------
+               100
+(1 row)
+
 drop table r1_reshuffle;
 -- table with update triggers on distributed key column
 CREATE OR REPLACE FUNCTION trigger_func() RETURNS trigger LANGUAGE plpgsql AS '

src/test/regress/sql/reshuffle.sql

@@ -172,20 +172,71 @@ Select count(*) > 0 from r1_reshuffle where gp_segment_id=2;
 drop table r1_reshuffle;
 
 -- Replicated tables
+-- We have to make sure replicated table successfully reshuffled.
+-- Currently we could only connect to the specific segments in utility mode
+-- to see if the data is consistent there. We create some python function here.
+-- The case can only work under the assumption: it's running on 3-segment cluster
+-- in a single machine.
+-- start_ignore
+drop language plpythonu;
+-- end_ignore
+create language plpythonu;
+create function update_on_segment(tabname text, segid int, numseg int) returns boolean as
+$$
+import pygresql.pg as pg
+conn = pg.connect(dbname='regression')
+port = conn.query("select port from gp_segment_configuration where content = %d and role = 'p'" % segid).getresult()[0][0]
+conn.close()
+
+conn = pg.connect(dbname='regression', opt='-c gp_session_role=utility', port=port)
+conn.query("set allow_system_table_mods = true")
+conn.query("update gp_distribution_policy set numsegments = %d where localoid = '%s'::regclass" % (numseg, tabname))
+conn.close()
+
+return True
+$$
+LANGUAGE plpythonu;
+
+create function select_on_segment(sql text, segid int) returns int as
+$$
+import pygresql.pg as pg
+conn = pg.connect(dbname='regression')
+port = conn.query("select port from gp_segment_configuration where content = %d and role = 'p'" % segid).getresult()[0][0]
+conn.close()
+
+conn = pg.connect(dbname='regression', opt='-c gp_session_role=utility', port=port)
+r = conn.query(sql).getresult()
+conn.close()
+
+return r[0][0]
+$$
+LANGUAGE plpythonu;
+
 Create table r1_reshuffle(a int, b int, c int) distributed replicated;
 update gp_distribution_policy  set numsegments=1 where localoid='r1_reshuffle'::regclass;
+select update_on_segment('r1_reshuffle', 0, 1);
+select update_on_segment('r1_reshuffle', 1, 1);
+select update_on_segment('r1_reshuffle', 2, 1);
 insert into r1_reshuffle select i,i,0 from generate_series(1,100) I;
 Select count(*) from r1_reshuffle;
+Select select_on_segment('Select count(*) from r1_reshuffle;', 1);
+Select select_on_segment('Select count(*) from r1_reshuffle;', 2);
 Alter table r1_reshuffle set with (reshuffle);
-Select count(*) from r1_reshuffle;
+Select select_on_segment('Select count(*) from r1_reshuffle;', 1);
+Select select_on_segment('Select count(*) from r1_reshuffle;', 2);
 drop table r1_reshuffle;
 
 Create table r1_reshuffle(a int, b int, c int) distributed replicated;
 update gp_distribution_policy  set numsegments=2 where localoid='r1_reshuffle'::regclass;
+select update_on_segment('r1_reshuffle', 0, 2);
+select update_on_segment('r1_reshuffle', 1, 2);
+select update_on_segment('r1_reshuffle', 2, 2);
 insert into r1_reshuffle select i,i,0 from generate_series(1,100) I;
 Select count(*) from r1_reshuffle;
+Select select_on_segment('Select count(*) from r1_reshuffle;', 2);
 Alter table r1_reshuffle set with (reshuffle);
 Select count(*) from r1_reshuffle;
+Select select_on_segment('Select count(*) from r1_reshuffle;', 2);
 drop table r1_reshuffle;
 
 -- table with update triggers on distributed key column