-- Test functions for copying intermediate results
CREATE SCHEMA intermediate_results;
SET search_path TO 'intermediate_results';
-- in the same transaction we can read a result
BEGIN;
SELECT create_intermediate_result('squares', 'SELECT s, s*s FROM generate_series(1,5) s');
 create_intermediate_result 
----------------------------
                          5
(1 row)

SELECT * FROM read_intermediate_result('squares', 'binary') AS res (x int, x2 int);
 x | x2 
---+----
 1 |  1
 2 |  4
 3 |  9
 4 | 16
 5 | 25
(5 rows)

COMMIT;
-- in separate transactions, the result is no longer available
SELECT create_intermediate_result('squares', 'SELECT s, s*s FROM generate_series(1,5) s');
 create_intermediate_result 
----------------------------
                          5
(1 row)

SELECT * FROM read_intermediate_result('squares', 'binary') AS res (x int, x2 int);
ERROR:  result "squares" does not exist
BEGIN;
CREATE TABLE interesting_squares (user_id text, interested_in text);
SELECT create_distributed_table('interesting_squares', 'user_id');
 create_distributed_table 
--------------------------
 
(1 row)

INSERT INTO interesting_squares VALUES ('jon', '2'), ('jon', '5'), ('jack', '3');
-- put an intermediate result on all workers
SELECT broadcast_intermediate_result('squares', 'SELECT s, s*s FROM generate_series(1,5) s');
 broadcast_intermediate_result 
-------------------------------
                             5
(1 row)

-- query the intermediate result in a router query
SELECT x, x2
FROM interesting_squares JOIN (SELECT * FROM read_intermediate_result('squares', 'binary') AS res (x int, x2 int)) squares ON (x::text = interested_in)
WHERE user_id = 'jon'
ORDER BY x;
 x | x2 
---+----
 2 |  4
 5 | 25
(2 rows)

END;
BEGIN;
-- put an intermediate result on all workers
SELECT broadcast_intermediate_result('squares', 'SELECT s, s*s FROM generate_series(1,5) s');
 broadcast_intermediate_result 
-------------------------------
                             5
(1 row)

-- query the intermediate result in a distributed query
SELECT x, x2
FROM interesting_squares
JOIN (SELECT * FROM read_intermediate_result('squares', 'binary') AS res (x int, x2 int)) squares ON (x::text = interested_in)
ORDER BY x;
 x | x2 
---+----
 2 |  4
 3 |  9
 5 | 25
(3 rows)

END;
CREATE FUNCTION raise_failed_execution_int_result(query text) RETURNS void AS $$
BEGIN
        EXECUTE query;
        EXCEPTION WHEN OTHERS THEN
        IF SQLERRM LIKE '%does not exist%' THEN
                RAISE 'Task failed to execute';
        ELSIF SQLERRM LIKE '%could not receive query results%' THEN
          RAISE 'Task failed to execute';
        END IF;
END;
$$LANGUAGE plpgsql;
-- don't print the worker port
\set VERBOSITY terse
SET client_min_messages TO ERROR;
-- files should now be cleaned up
SELECT raise_failed_execution_int_result($$
	SELECT x, x2
	FROM interesting_squares JOIN (SELECT * FROM read_intermediate_result('squares', 'binary') AS res (x text, x2 int)) squares ON (x = interested_in)
	WHERE user_id = 'jon'
	ORDER BY x;
$$);
ERROR:  Task failed to execute
\set VERBOSITY DEFAULT
SET client_min_messages TO DEFAULT;
-- try to read the file as text, will fail because of binary encoding
BEGIN;
SELECT create_intermediate_result('squares', 'SELECT s, s*s FROM generate_series(1,5) s');
 create_intermediate_result 
----------------------------
                          5
(1 row)

SELECT * FROM read_intermediate_result('squares', 'binary') AS res (x text, x2 int);
ERROR:  invalid byte sequence for encoding "UTF8": 0x00
END;
-- try to read the file with wrong encoding
BEGIN;
SELECT create_intermediate_result('squares', 'SELECT s, s*s FROM generate_series(1,5) s');
 create_intermediate_result 
----------------------------
                          5
(1 row)

SELECT * FROM read_intermediate_result('squares', 'csv') AS res (x int, x2 int);
ERROR:  invalid input syntax for type integer: "PGCOPY"
END;
-- try a composite type
CREATE TYPE intermediate_results.square_type AS (x text, x2 int);
CREATE TABLE stored_squares (user_id text, square intermediate_results.square_type, metadata jsonb);
INSERT INTO stored_squares VALUES ('jon', '(2,4)'::intermediate_results.square_type, '{"value":2}');
INSERT INTO stored_squares VALUES ('jon', '(3,9)'::intermediate_results.square_type, '{"value":3}');
INSERT INTO stored_squares VALUES ('jon', '(4,16)'::intermediate_results.square_type, '{"value":4}');
INSERT INTO stored_squares VALUES ('jon', '(5,25)'::intermediate_results.square_type, '{"value":5}');
-- composite types change the format to text
BEGIN;
SELECT create_intermediate_result('stored_squares', 'SELECT square FROM stored_squares');
 create_intermediate_result 
----------------------------
                          4
(1 row)

SELECT * FROM read_intermediate_result('stored_squares', 'binary') AS res (s intermediate_results.square_type);
ERROR:  COPY file signature not recognized
COMMIT;
BEGIN;
SELECT create_intermediate_result('stored_squares', 'SELECT square FROM stored_squares');
 create_intermediate_result 
----------------------------
                          4
(1 row)

SELECT * FROM read_intermediate_result('stored_squares', 'text') AS res (s intermediate_results.square_type);
   s    
--------
 (2,4)
 (3,9)
 (4,16)
 (5,25)
(4 rows)

COMMIT;
BEGIN;
-- put an intermediate result in text format on all workers
SELECT broadcast_intermediate_result('stored_squares', 'SELECT square, metadata FROM stored_squares');
 broadcast_intermediate_result 
-------------------------------
                             4
(1 row)

-- query the intermediate result in a router query using text format
SELECT * FROM interesting_squares JOIN (
  SELECT * FROM
    read_intermediate_result('stored_squares', 'text') AS res (s intermediate_results.square_type, m jsonb)
) squares
ON ((s).x = interested_in) WHERE user_id = 'jon' ORDER BY 1,2;
 user_id | interested_in |   s    |      m       
---------+---------------+--------+--------------
 jon     | 2             | (2,4)  | {"value": 2}
 jon     | 5             | (5,25) | {"value": 5}
(2 rows)

-- query the intermediate result in a real-time query using text format
SELECT * FROM interesting_squares JOIN (
  SELECT * FROM
    read_intermediate_result('stored_squares', 'text') AS res (s intermediate_results.square_type, m jsonb)
) squares
ON ((s).x = interested_in) ORDER BY 1,2;
 user_id | interested_in |   s    |      m       
---------+---------------+--------+--------------
 jack    | 3             | (3,9)  | {"value": 3}
 jon     | 2             | (2,4)  | {"value": 2}
 jon     | 5             | (5,25) | {"value": 5}
(3 rows)

END;
BEGIN;
-- accurate row count estimates for primitive types
SELECT create_intermediate_result('squares', 'SELECT s, s*s FROM generate_series(1,632) s');
 create_intermediate_result 
----------------------------
                        632
(1 row)

EXPLAIN (COSTS OFF) SELECT * FROM read_intermediate_result('squares', 'binary') AS res (x int, x2 int);
                  QUERY PLAN                   
-----------------------------------------------
 Function Scan on read_intermediate_result res
(1 row)

-- less accurate results for variable types
SELECT create_intermediate_result('hellos', $$SELECT s, 'hello-'||s FROM generate_series(1,63) s$$);
 create_intermediate_result 
----------------------------
                         63
(1 row)

EXPLAIN (COSTS OFF) SELECT * FROM read_intermediate_result('hellos', 'binary') AS res (x int, y text);
                  QUERY PLAN                   
-----------------------------------------------
 Function Scan on read_intermediate_result res
(1 row)

-- not very accurate results for text encoding
SELECT create_intermediate_result('stored_squares', 'SELECT square FROM stored_squares');
 create_intermediate_result 
----------------------------
                          4
(1 row)

EXPLAIN (COSTS OFF) SELECT * FROM read_intermediate_result('stored_squares', 'text') AS res (s intermediate_results.square_type);
                  QUERY PLAN                   
-----------------------------------------------
 Function Scan on read_intermediate_result res
(1 row)

END;
-- pipe query output into a result file and create a table to check the result
COPY (SELECT s, s*s FROM generate_series(1,5) s)
TO PROGRAM
  $$psql -h localhost -p 57636 -U postgres -d regression -c "BEGIN; COPY squares FROM STDIN WITH (format result); CREATE TABLE intermediate_results.squares AS SELECT * FROM read_intermediate_result('squares', 'text') AS res(x int, x2 int); END;"$$
WITH (FORMAT text);
SELECT * FROM squares ORDER BY x;
 x | x2 
---+----
 1 |  1
 2 |  4
 3 |  9
 4 | 16
 5 | 25
(5 rows)

-- empty shard interval array should raise error
SELECT worker_hash_partition_table(42,1,'SELECT a FROM generate_series(1,100) AS a', 'a', 23, ARRAY[0]);
ERROR:  invalid distribution column value
-- cannot use DDL commands
select broadcast_intermediate_result('a', 'create table foo(int serial)');
ERROR:  cannot execute utility commands
select broadcast_intermediate_result('a', 'prepare foo as select 1');
ERROR:  cannot execute utility commands
select create_intermediate_result('a', 'create table foo(int serial)');
ERROR:  cannot execute utility commands
--
-- read_intermediate_results
--
BEGIN;
SELECT create_intermediate_result('squares_1', 'SELECT s, s*s FROM generate_series(1,3) s'),
       create_intermediate_result('squares_2', 'SELECT s, s*s FROM generate_series(4,6) s'),
       create_intermediate_result('squares_3', 'SELECT s, s*s FROM generate_series(7,10) s');
 create_intermediate_result | create_intermediate_result | create_intermediate_result 
----------------------------+----------------------------+----------------------------
                          3 |                          3 |                          4
(1 row)

SELECT count(*) FROM read_intermediate_results(ARRAY[]::text[], 'binary') AS res (x int, x2 int);
 count 
-------
     0
(1 row)

SELECT * FROM read_intermediate_results(ARRAY['squares_1']::text[], 'binary') AS res (x int, x2 int);
 x | x2 
---+----
 1 |  1
 2 |  4
 3 |  9
(3 rows)

SELECT * FROM read_intermediate_results(ARRAY['squares_1', 'squares_2', 'squares_3']::text[], 'binary') AS res (x int, x2 int);
 x  | x2  
----+-----
  1 |   1
  2 |   4
  3 |   9
  4 |  16
  5 |  25
  6 |  36
  7 |  49
  8 |  64
  9 |  81
 10 | 100
(10 rows)

COMMIT;
-- in separate transactions, the result is no longer available
SELECT create_intermediate_result('squares_1', 'SELECT s, s*s FROM generate_series(1,5) s');
 create_intermediate_result 
----------------------------
                          5
(1 row)

SELECT * FROM read_intermediate_results(ARRAY['squares_1']::text[], 'binary') AS res (x int, x2 int);
ERROR:  result "squares_1" does not exist
-- error behaviour, and also check that results are deleted on rollback
BEGIN;
SELECT create_intermediate_result('squares_1', 'SELECT s, s*s FROM generate_series(1,3) s');
 create_intermediate_result 
----------------------------
                          3
(1 row)

SAVEPOINT s1;
SELECT * FROM read_intermediate_results(ARRAY['notexistingfile', 'squares_1'], 'binary') AS res (x int, x2 int);
ERROR:  result "notexistingfile" does not exist
ROLLBACK TO SAVEPOINT s1;
SELECT * FROM read_intermediate_results(ARRAY['squares_1', 'notexistingfile'], 'binary') AS res (x int, x2 int);
ERROR:  result "notexistingfile" does not exist
ROLLBACK TO SAVEPOINT s1;
SELECT * FROM read_intermediate_results(ARRAY['squares_1', NULL], 'binary') AS res (x int, x2 int);
ERROR:  null array element not allowed in this context
ROLLBACK TO SAVEPOINT s1;
-- after rollbacks we should be able to run vail read_intermediate_results still.
SELECT count(*) FROM read_intermediate_results(ARRAY['squares_1']::text[], 'binary') AS res (x int, x2 int);
 count 
-------
     3
(1 row)

SELECT count(*) FROM read_intermediate_results(ARRAY[]::text[], 'binary') AS res (x int, x2 int);
 count 
-------
     0
(1 row)

END;
SELECT * FROM read_intermediate_results(ARRAY['squares_1']::text[], 'binary') AS res (x int, x2 int);
ERROR:  result "squares_1" does not exist
-- Test non-binary format: read_intermediate_results(..., 'text')
BEGIN;
-- ROW(...) types switch the output format to text
SELECT broadcast_intermediate_result('stored_squares_1',
                                     'SELECT s, s*s, ROW(1::text, 2) FROM generate_series(1,3) s'),
       broadcast_intermediate_result('stored_squares_2',
                                     'SELECT s, s*s, ROW(2::text, 3) FROM generate_series(4,6) s');
 broadcast_intermediate_result | broadcast_intermediate_result 
-------------------------------+-------------------------------
                             3 |                             3
(1 row)

-- query the intermediate result in a router query using text format
SELECT * FROM interesting_squares JOIN (
  SELECT * FROM
    read_intermediate_results(ARRAY['stored_squares_1', 'stored_squares_2'], 'binary') AS res (x int, x2 int, z intermediate_results.square_type)
) squares
ON (squares.x::text = interested_in) WHERE user_id = 'jon' ORDER BY 1,2;
 user_id | interested_in | x | x2 |   z   
---------+---------------+---+----+-------
 jon     | 2             | 2 |  4 | (1,2)
 jon     | 5             | 5 | 25 | (2,3)
(2 rows)

END;
DROP SCHEMA intermediate_results CASCADE;
NOTICE:  drop cascades to 5 other objects
DETAIL:  drop cascades to table interesting_squares
drop cascades to function raise_failed_execution_int_result(text)
drop cascades to type square_type
drop cascades to table stored_squares
drop cascades to table squares