citus/src/test/regress/sql/cursors.sql

113 lines
3.2 KiB
PL/PgSQL

CREATE SCHEMA cursors;
SET search_path TO cursors;
CREATE TABLE distributed_table (key int, value text);
SELECT create_distributed_table('distributed_table', 'key');
-- load some data, but not very small amounts because RETURN QUERY in plpgsql
-- hard codes the cursor fetch to 50 rows on PG 12, though they might increase
-- it sometime in the future, so be mindful
INSERT INTO distributed_table SELECT i % 10, i::text FROM generate_series(0, 1000) i;
CREATE OR REPLACE FUNCTION simple_cursor_on_dist_table(cursor_name refcursor) RETURNS refcursor AS '
BEGIN
OPEN $1 FOR SELECT DISTINCT key FROM distributed_table ORDER BY 1;
RETURN $1;
END;
' LANGUAGE plpgsql;
CREATE OR REPLACE FUNCTION cursor_with_intermediate_result_on_dist_table(cursor_name refcursor) RETURNS refcursor AS '
BEGIN
OPEN $1 FOR
WITH cte_1 AS (SELECT * FROM distributed_table OFFSET 0)
SELECT DISTINCT key FROM distributed_table WHERE value in (SELECT value FROM cte_1) ORDER BY 1;
RETURN $1;
END;
' LANGUAGE plpgsql;
CREATE OR REPLACE FUNCTION cursor_with_intermediate_result_on_dist_table_with_param(cursor_name refcursor, filter text) RETURNS refcursor AS '
BEGIN
OPEN $1 FOR
WITH cte_1 AS (SELECT * FROM distributed_table WHERE value < $2 OFFSET 0)
SELECT DISTINCT key FROM distributed_table WHERE value in (SELECT value FROM cte_1) ORDER BY 1;
RETURN $1;
END;
' LANGUAGE plpgsql;
-- pretty basic query with cursors
-- Citus should plan/execute once and pull
-- the results to coordinator, then serve it
-- from the coordinator
BEGIN;
SELECT simple_cursor_on_dist_table('cursor_1');
SET LOCAL citus.log_intermediate_results TO ON;
SET LOCAL client_min_messages TO DEBUG1;
FETCH 5 IN cursor_1;
FETCH 50 IN cursor_1;
FETCH ALL IN cursor_1;
COMMIT;
BEGIN;
SELECT cursor_with_intermediate_result_on_dist_table('cursor_1');
-- multiple FETCH commands should not trigger re-running the subplans
SET LOCAL citus.log_intermediate_results TO ON;
SET LOCAL client_min_messages TO DEBUG1;
FETCH 5 IN cursor_1;
FETCH 1 IN cursor_1;
FETCH ALL IN cursor_1;
FETCH 5 IN cursor_1;
COMMIT;
BEGIN;
SELECT cursor_with_intermediate_result_on_dist_table_with_param('cursor_1', '600');
-- multiple FETCH commands should not trigger re-running the subplans
-- also test with parameters
SET LOCAL citus.log_intermediate_results TO ON;
SET LOCAL client_min_messages TO DEBUG1;
FETCH 1 IN cursor_1;
FETCH 1 IN cursor_1;
FETCH 1 IN cursor_1;
FETCH 1 IN cursor_1;
FETCH 1 IN cursor_1;
FETCH 1 IN cursor_1;
FETCH ALL IN cursor_1;
COMMIT;
CREATE OR REPLACE FUNCTION value_counter() RETURNS TABLE(counter text) LANGUAGE PLPGSQL AS $function$
BEGIN
return query
WITH cte AS
(SELECT dt.value
FROM distributed_table dt
WHERE dt.value in
(SELECT value
FROM distributed_table p
GROUP BY p.value
HAVING count(*) > 0))
SELECT * FROM cte;
END;
$function$ ;
SET citus.log_intermediate_results TO ON;
SET client_min_messages TO DEBUG1;
\set VERBOSITY terse
SELECT count(*) from (SELECT value_counter()) as foo;
BEGIN;
SELECT count(*) from (SELECT value_counter()) as foo;
COMMIT;
-- suppress NOTICEs
SET client_min_messages TO ERROR;
DROP SCHEMA cursors CASCADE;