CREATE SCHEMA adaptive_executor;
SET search_path TO adaptive_executor;
CREATE TABLE test (x int, y int);
SET citus.shard_count TO 4;
SET citus.shard_replication_factor TO 1;
SET citus.next_shard_id TO 801009000;
SELECT create_distributed_table('test','x');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

-- Add 1 row to each shard
SELECT get_shard_id_for_distribution_column('test', 1);
 get_shard_id_for_distribution_column
---------------------------------------------------------------------
                            801009000
(1 row)

INSERT INTO test VALUES (1,2);
SELECT get_shard_id_for_distribution_column('test', 3);
 get_shard_id_for_distribution_column
---------------------------------------------------------------------
                            801009001
(1 row)

INSERT INTO test VALUES (3,2);
SELECT get_shard_id_for_distribution_column('test', 6);
 get_shard_id_for_distribution_column
---------------------------------------------------------------------
                            801009002
(1 row)

INSERT INTO test VALUES (8,2);
SELECT get_shard_id_for_distribution_column('test', 11);
 get_shard_id_for_distribution_column
---------------------------------------------------------------------
                            801009003
(1 row)

INSERT INTO test VALUES (11,2);
-- Set a very high slow start to avoid opening parallel connections
SET citus.executor_slow_start_interval TO '60s';
SET citus.max_adaptive_executor_pool_size TO 2;
BEGIN;
SELECT count(*) FROM test a JOIN (SELECT x, pg_sleep(0.1) FROM test) b USING (x);
 count
---------------------------------------------------------------------
     4
(1 row)

SELECT sum(result::bigint) FROM run_command_on_workers($$
  SELECT count(*) FROM pg_stat_activity
  WHERE pid <> pg_backend_pid() AND query LIKE '%8010090%'
$$);
 sum
---------------------------------------------------------------------
   2
(1 row)

END;
-- SELECT takes longer than slow start interval, should open multiple connections
SET citus.executor_slow_start_interval TO '10ms';
BEGIN;
SELECT count(*) FROM test a JOIN (SELECT x, pg_sleep(0.2) FROM test) b USING (x);
 count
---------------------------------------------------------------------
     4
(1 row)

SELECT sum(result::bigint) FROM run_command_on_workers($$
  SELECT count(*) FROM pg_stat_activity
  WHERE pid <> pg_backend_pid() AND query LIKE '%8010090%'
$$);
 sum
---------------------------------------------------------------------
   4
(1 row)

END;
CREATE OR REPLACE FUNCTION select_for_update()
RETURNS void
AS $$
DECLARE
    my int;
BEGIN
    SELECT y INTO my FROM test WHERE x = 1 FOR UPDATE;
END;
$$ LANGUAGE plpgsql;
-- so that we can prove that we open a transaction block by logging below:
-- "NOTICE:  issuing BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;..."
SET citus.log_remote_commands TO on;
SELECT select_for_update();
NOTICE:  issuing BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;SELECT assign_distributed_transaction_id(xx, xx, 'xxxxxxx');
DETAIL:  on server postgres@localhost:xxxxx connectionId: xxxxxxx
CONTEXT:  SQL statement "SELECT y         FROM test WHERE x = 1 FOR UPDATE"
PL/pgSQL function select_for_update() line XX at SQL statement
NOTICE:  issuing SELECT y FROM adaptive_executor.test_801009000 test WHERE (x OPERATOR(pg_catalog.=) 1) FOR UPDATE OF test
DETAIL:  on server postgres@localhost:xxxxx connectionId: xxxxxxx
CONTEXT:  SQL statement "SELECT y         FROM test WHERE x = 1 FOR UPDATE"
PL/pgSQL function select_for_update() line XX at SQL statement
NOTICE:  issuing COMMIT
DETAIL:  on server postgres@localhost:xxxxx connectionId: xxxxxxx
 select_for_update
---------------------------------------------------------------------

(1 row)

SET citus.log_remote_commands TO off;
DROP SCHEMA adaptive_executor CASCADE;
NOTICE:  drop cascades to 2 other objects
DETAIL:  drop cascades to table test
drop cascades to function select_for_update()