src/test/isolation2/sql/setup.sql - cloudberry - Git at Google

 create or replace language plpython3u;

 -- Helper function, to call either __gp_aoseg, or gp_aocsseg, depending
 -- on whether the table is row- or column-oriented. This allows us to
 -- run the same test queries on both.
 --
 -- The Python utility that runs this doesn't know about dollar-quoting,
 -- and thinks that a ';' at end of line ends the command. The /* in func */
 -- comments at the end of each line thwarts that.
 CREATE OR REPLACE FUNCTION gp_ao_or_aocs_seg(rel regclass,
   segment_id OUT integer,
   segno OUT integer,
   tupcount OUT bigint,
   modcount OUT bigint,
   formatversion OUT smallint,
   state OUT smallint)
 RETURNS SETOF record as $$
 declare
   amname_var text;	/* in func */
 begin	/* in func */
   select amname into amname_var from pg_class c, pg_am am where c.relam = am.oid and c.oid = rel; /* in func */
   if amname_var = 'ao_column' then	/* in func */
     for segment_id, segno, tupcount, modcount, formatversion, state in SELECT DISTINCT x.segment_id, x.segno, x.tupcount, x.modcount, x.formatversion, x.state FROM gp_toolkit.__gp_aocsseg(rel) x loop	/* in func */
       return next;	/* in func */
     end loop;	/* in func */
   elsif amname_var = 'ao_row' then	/* in func */
     for segment_id, segno, tupcount, modcount, formatversion, state in SELECT x.segment_id, x.segno, x.tupcount, x.modcount, x.formatversion, x.state FROM gp_toolkit.__gp_aoseg(rel) x loop	/* in func */
       return next;	/* in func */
     end loop;	/* in func */
   else	/* in func */
     raise '% is not an AO_ROW or AO_COLUMN table', rel::text;	/* in func */
   end if;	/* in func */
 end;	/* in func */
 $$ LANGUAGE plpgsql;

 -- Show locks in master and in segments. Because the number of segments
 -- in the cluster depends on configuration, we print only summary information
 -- of the locks in segments. If a relation is locked only on one segment,
 -- we print that as a special case, but otherwise we just print "n segments",
 -- meaning the relation is locked on more than one segment.
 create or replace view locktest_master as
 select coalesce(
   case when relname like 'pg_toast%index' then 'toast index'
        when relname like 'pg_toast%' then 'toast table'
        when relname like 'pg_aoseg%' then 'aoseg table'
        when relname like 'pg_aovisimap%index' then 'aovisimap index'
        when relname like 'pg_aovisimap%' then 'aovisimap table'
        else relname end, 'dropped table'),
   mode,
   locktype,
   'master'::text as node
 from pg_locks l
 left outer join pg_class c on ((l.locktype = 'append-only segment file' and l.relation = c.relfilenode) or (l.locktype != 'append-only segment file' and l.relation = c.oid)),
 pg_database d
 where relation is not null
 and l.database = d.oid
 and (relname <> 'gp_fault_strategy' and relname != 'locktest_master' or relname is NULL)
 and d.datname = current_database()
 and l.gp_segment_id = -1
 group by l.gp_segment_id, relation, relname, locktype, mode
 order by 1, 3, 2;

 create or replace view locktest_segments_dist as
 select relname,
   mode,
   locktype,
   l.gp_segment_id as node,
   relation
 from pg_locks l
 left outer join pg_class c on ((l.locktype = 'append-only segment file' and l.relation = c.relfilenode) or (l.locktype != 'append-only segment file' and l.relation = c.oid)),
 pg_database d
 where relation is not null
 and l.database = d.oid
 and (relname <> 'gp_fault_strategy' and relname != 'locktest_segments_dist' or relname is NULL)
 and d.datname = current_database()
 and l.gp_segment_id > -1
 group by l.gp_segment_id, relation, relname, locktype, mode;

 create or replace view locktest_segments as
 SELECT coalesce(
   case when relname like 'pg_toast%index' then 'toast index'
        when relname like 'pg_toast%' then 'toast table'
        when relname like 'pg_aoseg%' then 'aoseg table'
        when relname like 'pg_aovisimap%index' then 'aovisimap index'
        when relname like 'pg_aovisimap%' then 'aovisimap table'
        else relname end, 'dropped table'),
   mode,
   locktype,
   case when count(*) = 1 then '1 segment'
        else 'n segments' end as node
   FROM gp_dist_random('locktest_segments_dist')
   group by relname, relation, mode, locktype;

 -- Helper function
 CREATE or REPLACE FUNCTION wait_until_waiting_for_required_lock (rel_name text, lmode text, segment_id integer) /*in func*/
 RETURNS bool AS
 $$
 declare
 retries int; /* in func */
 begin /* in func */
   retries := 1200; /* in func */
   loop /* in func */
     if (select not granted from pg_locks l where granted='f' and l.relation::regclass = rel_name::regclass and l.mode=lmode and l.gp_segment_id=segment_id) then /* in func */
       return true; /* in func */
     end if; /* in func */
     if retries <= 0 then /* in func */
       return false; /* in func */
     end if; /* in func */
     perform pg_sleep(0.1); /* in func */
     retries := retries - 1; /* in func */
   end loop; /* in func */
 end; /* in func */
 $$ language plpgsql;


 --
 -- pg_ctl:
 --   datadir: data directory of process to target with `pg_ctl`
 --   command: commands valid for `pg_ctl`
 --   command_mode: modes valid for `pg_ctl -m`
 --
 create or replace function pg_ctl(datadir text, command text, command_mode text default 'immediate')
 returns text as $$
     class PgCtlError(Exception):
         def __init__(self, errmsg):
             self.errmsg = errmsg
         def __str__(self):
             return repr(self.errmsg)

     import subprocess
     if command == 'promote':
         cmd = 'pg_ctl promote -D %s' % datadir
     elif command in ('stop', 'restart'):
         cmd = 'pg_ctl -l postmaster.log -D %s ' % datadir
         cmd = cmd + '-w -t 600 -m %s %s' % (command_mode, command)
     else:
         return 'Invalid command input'

     proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                             shell=True)
     stdout, stderr = proc.communicate()

     if proc.returncode == 0:
         return 'OK'
     else:
         raise PgCtlError(stdout.decode()+'|'+stderr.decode())
 $$ language plpython3u;

 --
 -- pg_ctl_start:
 --
 -- Start a specific greenplum segment
 --
 -- intentionally separate from pg_ctl() because it needs more information
 --
 --   datadir: data directory of process to target with `pg_ctl`
 --   port: which port the server should start on
 --
 create or replace function pg_ctl_start(datadir text, port bigint)
 returns text as $$
     import subprocess
     cmd = 'pg_ctl -l postmaster.log -D %s ' % datadir
     opts = '-p %d' % (port)
     opts = opts + ' -c gp_role=execute'
     cmd = cmd + '-o "%s" start' % opts
     return subprocess.check_output(cmd, stderr=subprocess.STDOUT, shell=True).decode().replace('.', '')
 $$ language plpython3u;

 --
 -- restart_primary_segments_containing_data_for(table_name text):
 --     table_name: the table containing data whose segment that needs a restart
 --
 -- Note: this does an immediate restart, which forces recovery
 --
 create or replace function restart_primary_segments_containing_data_for(table_name text) returns setof integer as $$
 declare
 	segment_id integer; /* in func */
 begin
 	for segment_id in select * from primary_segments_containing_data_for(table_name)
 	loop
 		perform pg_ctl(
       (select get_data_directory_for(segment_id)),
       'restart',
       'immediate'
     ); /* in func */
 	end loop; /* in func */
 end; /* in func */
 $$ language plpgsql;

 --
 -- clean_restart_primary_segments_containing_data_for(table_name text):
 --     table_name: the table containing data whose segment that needs a restart
 --
 -- Note: this does a fast restart, which does not require recovery
 --
 create or replace function clean_restart_primary_segments_containing_data_for(table_name text) returns setof integer as $$
 declare
 	segment_id integer; /* in func */
 begin
 	for segment_id in select * from primary_segments_containing_data_for(table_name)
 	loop
 		perform pg_ctl(
       (select get_data_directory_for(segment_id)),
       'restart',
       'fast'
     ); /* in func */
 	end loop; /* in func */
 end; /* in func */
 $$ language plpgsql;

 create or replace function primary_segments_containing_data_for(table_name text) returns setof integer as $$
 begin
 	return query execute 'select distinct gp_segment_id from ' || table_name; /* in func */
 end; /* in func */
 $$ language plpgsql;


 create or replace function get_data_directory_for(segment_number int, segment_role text default 'p') returns text as $$
 BEGIN
 	return (
 		select datadir
 		from gp_segment_configuration
 		where role=segment_role and
 		content=segment_number
 	); /* in func */
 END; /* in func */
 $$ language plpgsql;

 create or replace function master() returns setof gp_segment_configuration as $$
 	select * from gp_segment_configuration where role='p' and content=-1; /* in func */
 $$ language sql;

 create or replace function wait_until_segment_synchronized(segment_number int) returns text as $$
 begin
 	for i in 1..6000 loop
 		if (select count(*) = 0 from gp_segment_configuration where content = segment_number and mode != 's') then
 			return 'OK'; /* in func */
 		end if; /* in func */
 		perform pg_sleep(0.1); /* in func */
 		perform gp_request_fts_probe_scan(); /* in func */
 	end loop; /* in func */
 	return 'Fail'; /* in func */
 end; /* in func */
 $$ language plpgsql;

 create or replace function wait_until_all_segments_synchronized() returns text as $$
 begin
 	for i in 1..6000 loop
 		if (select count(*) = 0 from gp_segment_configuration where content != -1 and mode != 's') then
 			return 'OK'; /* in func */
 		end if; /* in func */
 		perform pg_sleep(0.1); /* in func */
 		perform gp_request_fts_probe_scan(); /* in func */
  	end loop; /* in func */
 	return 'Fail'; /* in func */
 end; /* in func */
 $$ language plpgsql;

 CREATE OR REPLACE FUNCTION is_query_waiting_for_syncrep(iterations int, check_query text) RETURNS bool AS $$
     for i in range(iterations):
         results = plpy.execute("SELECT gp_execution_segment() AS content, query, wait_event\
                                 FROM gp_dist_random('pg_stat_activity')\
                                 WHERE gp_execution_segment() = 1 AND\
                                 query = '%s' AND\
                                 wait_event = 'SyncRep'" % check_query )
         if results:
             return True
     return False
 $$ LANGUAGE plpython3u VOLATILE;

 create or replace function wait_for_replication_replay (segid int, retries int) returns bool as
 $$
 declare
 	i int; /* in func */
 	result bool; /* in func */
 begin
 	i := 0; /* in func */
 	-- Wait until the mirror/standby has replayed up to flush location
 	loop
 		SELECT flush_lsn = replay_lsn INTO result from gp_stat_replication where gp_segment_id = segid; /* in func */
 		if result then
 			return true; /* in func */
 		end if; /* in func */

 		if i >= retries then
 		   return false; /* in func */
 		end if; /* in func */
 		perform pg_sleep(0.1); /* in func */
 		perform pg_stat_clear_snapshot(); /* in func */
 		i := i + 1; /* in func */
 	end loop; /* in func */
 end; /* in func */
 $$ language plpgsql;

 create or replace function wait_until_standby_in_state(targetstate text)
 returns text as $$
 declare
    replstate text; /* in func */
    i int; /* in func */
 begin
    i := 0; /* in func */
    while i < 1200 loop
       select state into replstate from pg_stat_replication; /* in func */
       if replstate = targetstate then
           return replstate; /* in func */
       end if; /* in func */
       perform pg_sleep(0.1); /* in func */
       perform pg_stat_clear_snapshot(); /* in func */
       i := i + 1; /* in func */
    end loop; /* in func */
    return replstate; /* in func */
 end; /* in func */
 $$ language plpgsql;

 --
 -- pg_basebackup:
 --   host: host of the gpdb segment to back up
 --   port: port of the gpdb segment to back up
 --   slotname: desired slot name to create and associate with backup
 --   datadir: destination data directory of the backup
 --   forceoverwrite: overwrite the destination directory if it exists already
 --   xlog_method: (stream/fetch) how to obtain XLOG segment files from source
 --
 -- usage: `select pg_basebackup('somehost', 12345, 'some_slot_name', '/some/destination/data/directory')`
 --
 create or replace function pg_basebackup(host text, dbid int, port bigint, create_slot boolean, slotname text, datadir text, force_overwrite boolean, xlog_method text) returns text as $$
     import subprocess
     import os
     cmd = 'pg_basebackup --no-sync --checkpoint=fast -h %s -p %d -R -D %s --target-gp-dbid %d' % (host, port, datadir, dbid)

     if create_slot:
         cmd += ' --create-slot'

     if slotname is not None:
         cmd += ' --slot %s' % (slotname)

     if force_overwrite:
         cmd += ' --force-overwrite'

     if xlog_method == 'stream':
         cmd += ' --wal-method stream'
     elif xlog_method == 'fetch':
         cmd += ' --wal-method fetch'
     else:
         plpy.error('invalid xlog method')

     cmd += ' --no-verify-checksums'

     try:
         # Unset PGAPPNAME so that the pg_stat_replication.application_name is not affected
         if os.getenv('PGAPPNAME') is not None:
             os.environ.pop('PGAPPNAME')
         results = subprocess.check_output(cmd, stderr=subprocess.STDOUT, shell=True).replace(b'.', b'').decode()
     except subprocess.CalledProcessError as e:
         results = str(e) + "\ncommand output: " + (e.output.decode())

     return results
 $$ language plpython3u;

 create or replace function count_of_items_in_directory(user_path text) returns text as $$
        import subprocess
        cmd = 'ls {user_path}'.format(user_path=user_path)
        results = subprocess.check_output(cmd, stderr=subprocess.STDOUT, shell=True).replace(b'.', b'').decode()
        return len([result for result in results.splitlines() if result != ''])
 $$ language plpython3u;

 create or replace function count_of_items_in_database_directory(user_path text, database_oid oid) returns int as $$
        import subprocess
        import os
        directory = os.path.join(user_path, str(database_oid))
        cmd = 'ls ' + directory
        results = subprocess.check_output(cmd, stderr=subprocess.STDOUT, shell=True).replace(b'.', b'').decode()
        return len([result for result in results.splitlines() if result != ''])
 $$ language plpython3u;

 create or replace function validate_tablespace_symlink(datadir text, tablespacedir text, dbid int, tablespace_oid oid) returns boolean as $$
     import os
     return os.readlink('%s/pg_tblspc/%d' % (datadir, tablespace_oid)) == ('%s/%d' % (tablespacedir, dbid))
 $$ language plpython3u;

 -- This function is used to loop until master shutsdown, to make sure
 -- next command executed is only after restart and doesn't go through
 -- while PANIC is still being processed by master, as master continues
 -- to accept connections for a while despite undergoing PANIC.
 CREATE OR REPLACE FUNCTION wait_till_master_shutsdown()
 RETURNS void AS
 $$
   DECLARE
     i int; /* in func */
   BEGIN
     i := 0; /* in func */
     while i < 120 loop
       i := i + 1; /* in func */
       PERFORM pg_sleep(.5); /* in func */
     end loop; /* in func */
   END; /* in func */
 $$ LANGUAGE plpgsql;

 -- start_ignore
 set statement_timeout='180s';
 -- end_ignore
 -- Helper function that ensures stats collector receives stat from the latest operation.
 create or replace function wait_until_dead_tup_change_to(relid oid, stat_val_expected bigint)
     returns text as $$
 declare
     stat_val int; /* in func */
     i int; /* in func */
 begin
     i := 0; /* in func */
     while i < 1200 loop
             select pg_stat_get_dead_tuples(relid) into stat_val; /* in func */
             if stat_val = stat_val_expected then /* in func */
                 return 'OK'; /* in func */
             end if; /* in func */
             perform pg_sleep(0.1); /* in func */
             perform pg_stat_clear_snapshot(); /* in func */
             i := i + 1; /* in func */
         end loop; /* in func */
     return 'Fail'; /* in func */
 end; /* in func */
 $$ language plpgsql;

 -- Helper function that ensures mirror of the specified contentid is down.
 create or replace function wait_for_mirror_down(contentid smallint, timeout_sec integer) returns bool as
 $$
 declare i int; /* in func */
 begin /* in func */
     i := 0; /* in func */
     loop /* in func */
         perform gp_request_fts_probe_scan(); /* in func */
         if (select count(1) from gp_segment_configuration where role='m' and content=$1 and status='d') = 1 then /* in func */
             return true; /* in func */
         end if; /* in func */
         if i >= 2 * $2 then /* in func */
             return false; /* in func */
         end if; /* in func */
         perform pg_sleep(0.5); /* in func */
         i = i + 1; /* in func */
     end loop; /* in func */
 end; /* in func */
 $$ language plpgsql;

 -- Helper function that ensures stats collector receives stat from the latest operation.
 create or replace function wait_until_vacuum_count_change_to(relid oid, stat_val_expected bigint)
     returns text as $$
 declare
     stat_val int; /* in func */
     i int; /* in func */
 begin
     i := 0; /* in func */
     while i < 1200 loop
             select pg_stat_get_vacuum_count(relid) into stat_val; /* in func */
             if stat_val = stat_val_expected then /* in func */
                 return 'OK'; /* in func */
             end if; /* in func */
             perform pg_sleep(0.1); /* in func */
             perform pg_stat_clear_snapshot(); /* in func */
             i := i + 1; /* in func */
         end loop; /* in func */
     return 'Fail'; /* in func */
 end; /* in func */
 $$ language plpgsql;

 -- Helper function to get the number of blocks in a relation.
 CREATE OR REPLACE FUNCTION nblocks(rel regclass) RETURNS int AS $$ /* in func */
 BEGIN /* in func */
 RETURN pg_relation_size(rel) / current_setting('block_size')::int; /* in func */
 END; $$ /* in func */
     LANGUAGE PLPGSQL;

 -- Helper function to populate logical heap pages in a certain block sequence.
 -- Can be used for both heap and AO/CO tables. The target block sequence into
 -- which we insert the pages depends on the session which is inserting the data.
 -- This is currently meant to be used with a single column integer table.
 --
 -- Sample usage: SELECT populate_pages('foo', 1, tid '(33554435,0)')
 -- This will insert tuples with value=1 into a single QE such that logical
 -- heap blocks [33554432, 33554434] will be full and 33554435 will have only
 -- 1 tuple.
 --
 -- Note: while using this with AO/CO tables, please account for how the block
 -- sequences start/end based on the concurrency level (see AOSegmentGet_startHeapBlock())
 CREATE OR REPLACE FUNCTION populate_pages(relname text, value int, upto tid) RETURNS VOID AS $$ /* in func */
 DECLARE curtid tid; /* in func */
 BEGIN /* in func */
 LOOP /* in func */
 EXECUTE format('INSERT INTO %I VALUES($1) RETURNING ctid', relname) INTO curtid USING value; /* in func */
 EXIT WHEN curtid > upto; /* in func */
 END LOOP; /* in func */
 END; $$ /* in func */
     LANGUAGE PLPGSQL;
	create or replace language plpython3u;

	-- Helper function, to call either __gp_aoseg, or gp_aocsseg, depending
	-- on whether the table is row- or column-oriented. This allows us to
	-- run the same test queries on both.
	--
	-- The Python utility that runs this doesn't know about dollar-quoting,
	-- and thinks that a ';' at end of line ends the command. The /* in func */
	-- comments at the end of each line thwarts that.
	CREATE OR REPLACE FUNCTION gp_ao_or_aocs_seg(rel regclass,
	segment_id OUT integer,
	segno OUT integer,
	tupcount OUT bigint,
	modcount OUT bigint,
	formatversion OUT smallint,
	state OUT smallint)
	RETURNS SETOF record as $$
	declare
	amname_var text; /* in func */
	begin /* in func */
	select amname into amname_var from pg_class c, pg_am am where c.relam = am.oid and c.oid = rel; /* in func */
	if amname_var = 'ao_column' then /* in func */
	for segment_id, segno, tupcount, modcount, formatversion, state in SELECT DISTINCT x.segment_id, x.segno, x.tupcount, x.modcount, x.formatversion, x.state FROM gp_toolkit.__gp_aocsseg(rel) x loop /* in func */
	return next; /* in func */
	end loop; /* in func */
	elsif amname_var = 'ao_row' then /* in func */
	for segment_id, segno, tupcount, modcount, formatversion, state in SELECT x.segment_id, x.segno, x.tupcount, x.modcount, x.formatversion, x.state FROM gp_toolkit.__gp_aoseg(rel) x loop /* in func */
	return next; /* in func */
	end loop; /* in func */
	else /* in func */
	raise '% is not an AO_ROW or AO_COLUMN table', rel::text; /* in func */
	end if; /* in func */
	end; /* in func */
	$$ LANGUAGE plpgsql;

	-- Show locks in master and in segments. Because the number of segments
	-- in the cluster depends on configuration, we print only summary information
	-- of the locks in segments. If a relation is locked only on one segment,
	-- we print that as a special case, but otherwise we just print "n segments",
	-- meaning the relation is locked on more than one segment.
	create or replace view locktest_master as
	select coalesce(
	case when relname like 'pg_toast%index' then 'toast index'
	when relname like 'pg_toast%' then 'toast table'
	when relname like 'pg_aoseg%' then 'aoseg table'
	when relname like 'pg_aovisimap%index' then 'aovisimap index'
	when relname like 'pg_aovisimap%' then 'aovisimap table'
	else relname end, 'dropped table'),
	mode,
	locktype,
	'master'::text as node
	from pg_locks l
	left outer join pg_class c on ((l.locktype = 'append-only segment file' and l.relation = c.relfilenode) or (l.locktype != 'append-only segment file' and l.relation = c.oid)),
	pg_database d
	where relation is not null
	and l.database = d.oid
	and (relname <> 'gp_fault_strategy' and relname != 'locktest_master' or relname is NULL)
	and d.datname = current_database()
	and l.gp_segment_id = -1
	group by l.gp_segment_id, relation, relname, locktype, mode
	order by 1, 3, 2;

	create or replace view locktest_segments_dist as
	select relname,
	mode,
	locktype,
	l.gp_segment_id as node,
	relation
	from pg_locks l
	left outer join pg_class c on ((l.locktype = 'append-only segment file' and l.relation = c.relfilenode) or (l.locktype != 'append-only segment file' and l.relation = c.oid)),
	pg_database d
	where relation is not null
	and l.database = d.oid
	and (relname <> 'gp_fault_strategy' and relname != 'locktest_segments_dist' or relname is NULL)
	and d.datname = current_database()
	and l.gp_segment_id > -1
	group by l.gp_segment_id, relation, relname, locktype, mode;

	create or replace view locktest_segments as
	SELECT coalesce(
	case when relname like 'pg_toast%index' then 'toast index'
	when relname like 'pg_toast%' then 'toast table'
	when relname like 'pg_aoseg%' then 'aoseg table'
	when relname like 'pg_aovisimap%index' then 'aovisimap index'
	when relname like 'pg_aovisimap%' then 'aovisimap table'
	else relname end, 'dropped table'),
	mode,
	locktype,
	case when count(*) = 1 then '1 segment'
	else 'n segments' end as node
	FROM gp_dist_random('locktest_segments_dist')
	group by relname, relation, mode, locktype;

	-- Helper function
	CREATE or REPLACE FUNCTION wait_until_waiting_for_required_lock (rel_name text, lmode text, segment_id integer) /in func/
	RETURNS bool AS
	$$
	declare
	retries int; /* in func */
	begin /* in func */
	retries := 1200; /* in func */
	loop /* in func */
	if (select not granted from pg_locks l where granted='f' and l.relation::regclass = rel_name::regclass and l.mode=lmode and l.gp_segment_id=segment_id) then /* in func */
	return true; /* in func */
	end if; /* in func */
	if retries <= 0 then /* in func */
	return false; /* in func */
	end if; /* in func */
	perform pg_sleep(0.1); /* in func */
	retries := retries - 1; /* in func */
	end loop; /* in func */
	end; /* in func */
	$$ language plpgsql;


	--
	-- pg_ctl:
	-- datadir: data directory of process to target with `pg_ctl`
	-- command: commands valid for `pg_ctl`
	-- command_mode: modes valid for `pg_ctl -m`
	--
	create or replace function pg_ctl(datadir text, command text, command_mode text default 'immediate')
	returns text as $$
	class PgCtlError(Exception):
	def __init__(self, errmsg):
	self.errmsg = errmsg
	def __str__(self):
	return repr(self.errmsg)

	import subprocess
	if command == 'promote':
	cmd = 'pg_ctl promote -D %s' % datadir
	elif command in ('stop', 'restart'):
	cmd = 'pg_ctl -l postmaster.log -D %s ' % datadir
	cmd = cmd + '-w -t 600 -m %s %s' % (command_mode, command)
	else:
	return 'Invalid command input'

	proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
	shell=True)
	stdout, stderr = proc.communicate()

	if proc.returncode == 0:
	return 'OK'
	else:
	raise PgCtlError(stdout.decode()+'\|'+stderr.decode())
	$$ language plpython3u;

	--
	-- pg_ctl_start:
	--
	-- Start a specific greenplum segment
	--
	-- intentionally separate from pg_ctl() because it needs more information
	--
	-- datadir: data directory of process to target with `pg_ctl`
	-- port: which port the server should start on
	--
	create or replace function pg_ctl_start(datadir text, port bigint)
	returns text as $$
	import subprocess
	cmd = 'pg_ctl -l postmaster.log -D %s ' % datadir
	opts = '-p %d' % (port)
	opts = opts + ' -c gp_role=execute'
	cmd = cmd + '-o "%s" start' % opts
	return subprocess.check_output(cmd, stderr=subprocess.STDOUT, shell=True).decode().replace('.', '')
	$$ language plpython3u;

	--
	-- restart_primary_segments_containing_data_for(table_name text):
	-- table_name: the table containing data whose segment that needs a restart
	--
	-- Note: this does an immediate restart, which forces recovery
	--
	create or replace function restart_primary_segments_containing_data_for(table_name text) returns setof integer as $$
	declare
	segment_id integer; /* in func */
	begin
	for segment_id in select * from primary_segments_containing_data_for(table_name)
	loop
	perform pg_ctl(
	(select get_data_directory_for(segment_id)),
	'restart',
	'immediate'
	); /* in func */
	end loop; /* in func */
	end; /* in func */
	$$ language plpgsql;

	--
	-- clean_restart_primary_segments_containing_data_for(table_name text):
	-- table_name: the table containing data whose segment that needs a restart
	--
	-- Note: this does a fast restart, which does not require recovery
	--
	create or replace function clean_restart_primary_segments_containing_data_for(table_name text) returns setof integer as $$
	declare
	segment_id integer; /* in func */
	begin
	for segment_id in select * from primary_segments_containing_data_for(table_name)
	loop
	perform pg_ctl(
	(select get_data_directory_for(segment_id)),
	'restart',
	'fast'
	); /* in func */
	end loop; /* in func */
	end; /* in func */
	$$ language plpgsql;

	create or replace function primary_segments_containing_data_for(table_name text) returns setof integer as $$
	begin
	return query execute 'select distinct gp_segment_id from ' \|\| table_name; /* in func */
	end; /* in func */
	$$ language plpgsql;


	create or replace function get_data_directory_for(segment_number int, segment_role text default 'p') returns text as $$
	BEGIN
	return (
	select datadir
	from gp_segment_configuration
	where role=segment_role and
	content=segment_number
	); /* in func */
	END; /* in func */
	$$ language plpgsql;

	create or replace function master() returns setof gp_segment_configuration as $$
	select * from gp_segment_configuration where role='p' and content=-1; /* in func */
	$$ language sql;

	create or replace function wait_until_segment_synchronized(segment_number int) returns text as $$
	begin
	for i in 1..6000 loop
	if (select count(*) = 0 from gp_segment_configuration where content = segment_number and mode != 's') then
	return 'OK'; /* in func */
	end if; /* in func */
	perform pg_sleep(0.1); /* in func */
	perform gp_request_fts_probe_scan(); /* in func */
	end loop; /* in func */
	return 'Fail'; /* in func */
	end; /* in func */
	$$ language plpgsql;

	create or replace function wait_until_all_segments_synchronized() returns text as $$
	begin
	for i in 1..6000 loop
	if (select count(*) = 0 from gp_segment_configuration where content != -1 and mode != 's') then
	return 'OK'; /* in func */
	end if; /* in func */
	perform pg_sleep(0.1); /* in func */
	perform gp_request_fts_probe_scan(); /* in func */
	end loop; /* in func */
	return 'Fail'; /* in func */
	end; /* in func */
	$$ language plpgsql;

	CREATE OR REPLACE FUNCTION is_query_waiting_for_syncrep(iterations int, check_query text) RETURNS bool AS $$
	for i in range(iterations):
	results = plpy.execute("SELECT gp_execution_segment() AS content, query, wait_event\
	FROM gp_dist_random('pg_stat_activity')\
	WHERE gp_execution_segment() = 1 AND\
	query = '%s' AND\
	wait_event = 'SyncRep'" % check_query )
	if results:
	return True
	return False
	$$ LANGUAGE plpython3u VOLATILE;

	create or replace function wait_for_replication_replay (segid int, retries int) returns bool as
	$$
	declare
	i int; /* in func */
	result bool; /* in func */
	begin
	i := 0; /* in func */
	-- Wait until the mirror/standby has replayed up to flush location
	loop
	SELECT flush_lsn = replay_lsn INTO result from gp_stat_replication where gp_segment_id = segid; /* in func */
	if result then
	return true; /* in func */
	end if; /* in func */

	if i >= retries then
	return false; /* in func */
	end if; /* in func */
	perform pg_sleep(0.1); /* in func */
	perform pg_stat_clear_snapshot(); /* in func */
	i := i + 1; /* in func */
	end loop; /* in func */
	end; /* in func */
	$$ language plpgsql;

	create or replace function wait_until_standby_in_state(targetstate text)
	returns text as $$
	declare
	replstate text; /* in func */
	i int; /* in func */
	begin
	i := 0; /* in func */
	while i < 1200 loop
	select state into replstate from pg_stat_replication; /* in func */
	if replstate = targetstate then
	return replstate; /* in func */
	end if; /* in func */
	perform pg_sleep(0.1); /* in func */
	perform pg_stat_clear_snapshot(); /* in func */
	i := i + 1; /* in func */
	end loop; /* in func */
	return replstate; /* in func */
	end; /* in func */
	$$ language plpgsql;

	--
	-- pg_basebackup:
	-- host: host of the gpdb segment to back up
	-- port: port of the gpdb segment to back up
	-- slotname: desired slot name to create and associate with backup
	-- datadir: destination data directory of the backup
	-- forceoverwrite: overwrite the destination directory if it exists already
	-- xlog_method: (stream/fetch) how to obtain XLOG segment files from source
	--
	-- usage: `select pg_basebackup('somehost', 12345, 'some_slot_name', '/some/destination/data/directory')`
	--
	create or replace function pg_basebackup(host text, dbid int, port bigint, create_slot boolean, slotname text, datadir text, force_overwrite boolean, xlog_method text) returns text as $$
	import subprocess
	import os
	cmd = 'pg_basebackup --no-sync --checkpoint=fast -h %s -p %d -R -D %s --target-gp-dbid %d' % (host, port, datadir, dbid)

	if create_slot:
	cmd += ' --create-slot'

	if slotname is not None:
	cmd += ' --slot %s' % (slotname)

	if force_overwrite:
	cmd += ' --force-overwrite'

	if xlog_method == 'stream':
	cmd += ' --wal-method stream'
	elif xlog_method == 'fetch':
	cmd += ' --wal-method fetch'
	else:
	plpy.error('invalid xlog method')

	cmd += ' --no-verify-checksums'

	try:
	# Unset PGAPPNAME so that the pg_stat_replication.application_name is not affected
	if os.getenv('PGAPPNAME') is not None:
	os.environ.pop('PGAPPNAME')
	results = subprocess.check_output(cmd, stderr=subprocess.STDOUT, shell=True).replace(b'.', b'').decode()
	except subprocess.CalledProcessError as e:
	results = str(e) + "\ncommand output: " + (e.output.decode())

	return results
	$$ language plpython3u;

	create or replace function count_of_items_in_directory(user_path text) returns text as $$
	import subprocess
	cmd = 'ls {user_path}'.format(user_path=user_path)
	results = subprocess.check_output(cmd, stderr=subprocess.STDOUT, shell=True).replace(b'.', b'').decode()
	return len([result for result in results.splitlines() if result != ''])
	$$ language plpython3u;

	create or replace function count_of_items_in_database_directory(user_path text, database_oid oid) returns int as $$
	import subprocess
	import os
	directory = os.path.join(user_path, str(database_oid))
	cmd = 'ls ' + directory
	results = subprocess.check_output(cmd, stderr=subprocess.STDOUT, shell=True).replace(b'.', b'').decode()
	return len([result for result in results.splitlines() if result != ''])
	$$ language plpython3u;

	create or replace function validate_tablespace_symlink(datadir text, tablespacedir text, dbid int, tablespace_oid oid) returns boolean as $$
	import os
	return os.readlink('%s/pg_tblspc/%d' % (datadir, tablespace_oid)) == ('%s/%d' % (tablespacedir, dbid))
	$$ language plpython3u;

	-- This function is used to loop until master shutsdown, to make sure
	-- next command executed is only after restart and doesn't go through
	-- while PANIC is still being processed by master, as master continues
	-- to accept connections for a while despite undergoing PANIC.
	CREATE OR REPLACE FUNCTION wait_till_master_shutsdown()
	RETURNS void AS
	$$
	DECLARE
	i int; /* in func */
	BEGIN
	i := 0; /* in func */
	while i < 120 loop
	i := i + 1; /* in func */
	PERFORM pg_sleep(.5); /* in func */
	end loop; /* in func */
	END; /* in func */
	$$ LANGUAGE plpgsql;

	-- start_ignore
	set statement_timeout='180s';
	-- end_ignore
	-- Helper function that ensures stats collector receives stat from the latest operation.
	create or replace function wait_until_dead_tup_change_to(relid oid, stat_val_expected bigint)
	returns text as $$
	declare
	stat_val int; /* in func */
	i int; /* in func */
	begin
	i := 0; /* in func */
	while i < 1200 loop
	select pg_stat_get_dead_tuples(relid) into stat_val; /* in func */
	if stat_val = stat_val_expected then /* in func */
	return 'OK'; /* in func */
	end if; /* in func */
	perform pg_sleep(0.1); /* in func */
	perform pg_stat_clear_snapshot(); /* in func */
	i := i + 1; /* in func */
	end loop; /* in func */
	return 'Fail'; /* in func */
	end; /* in func */
	$$ language plpgsql;

	-- Helper function that ensures mirror of the specified contentid is down.
	create or replace function wait_for_mirror_down(contentid smallint, timeout_sec integer) returns bool as
	$$
	declare i int; /* in func */
	begin /* in func */
	i := 0; /* in func */
	loop /* in func */
	perform gp_request_fts_probe_scan(); /* in func */
	if (select count(1) from gp_segment_configuration where role='m' and content=$1 and status='d') = 1 then /* in func */
	return true; /* in func */
	end if; /* in func */
	if i >= 2 * $2 then /* in func */
	return false; /* in func */
	end if; /* in func */
	perform pg_sleep(0.5); /* in func */
	i = i + 1; /* in func */
	end loop; /* in func */
	end; /* in func */
	$$ language plpgsql;

	-- Helper function that ensures stats collector receives stat from the latest operation.
	create or replace function wait_until_vacuum_count_change_to(relid oid, stat_val_expected bigint)
	returns text as $$
	declare
	stat_val int; /* in func */
	i int; /* in func */
	begin
	i := 0; /* in func */
	while i < 1200 loop
	select pg_stat_get_vacuum_count(relid) into stat_val; /* in func */
	if stat_val = stat_val_expected then /* in func */
	return 'OK'; /* in func */
	end if; /* in func */
	perform pg_sleep(0.1); /* in func */
	perform pg_stat_clear_snapshot(); /* in func */
	i := i + 1; /* in func */
	end loop; /* in func */
	return 'Fail'; /* in func */
	end; /* in func */
	$$ language plpgsql;

	-- Helper function to get the number of blocks in a relation.
	CREATE OR REPLACE FUNCTION nblocks(rel regclass) RETURNS int AS $$ /* in func */
	BEGIN /* in func */
	RETURN pg_relation_size(rel) / current_setting('block_size')::int; /* in func */
	END; $$ /* in func */
	LANGUAGE PLPGSQL;

	-- Helper function to populate logical heap pages in a certain block sequence.
	-- Can be used for both heap and AO/CO tables. The target block sequence into
	-- which we insert the pages depends on the session which is inserting the data.
	-- This is currently meant to be used with a single column integer table.
	--
	-- Sample usage: SELECT populate_pages('foo', 1, tid '(33554435,0)')
	-- This will insert tuples with value=1 into a single QE such that logical
	-- heap blocks [33554432, 33554434] will be full and 33554435 will have only
	-- 1 tuple.
	--
	-- Note: while using this with AO/CO tables, please account for how the block
	-- sequences start/end based on the concurrency level (see AOSegmentGet_startHeapBlock())
	CREATE OR REPLACE FUNCTION populate_pages(relname text, value int, upto tid) RETURNS VOID AS $$ /* in func */
	DECLARE curtid tid; /* in func */
	BEGIN /* in func */
	LOOP /* in func */
	EXECUTE format('INSERT INTO %I VALUES($1) RETURNING ctid', relname) INTO curtid USING value; /* in func */
	EXIT WHEN curtid > upto; /* in func */
	END LOOP; /* in func */
	END; $$ /* in func */
	LANGUAGE PLPGSQL;