--
-- stat_statements
--
-- tests citus_stat_statements functionality
SHOW server_version \gset
SELECT substring(:'server_version', '\d+')::int >= 14 AS server_version_ge_14
\gset
\if :server_version_ge_14
SET compute_query_id = 'on';
\endif
-- check if pg_stat_statements is available
SELECT name FROM pg_available_extensions WHERE name = 'pg_stat_statements';
        name
---------------------------------------------------------------------
 pg_stat_statements
(1 row)

SELECT regexp_split_to_array(setting, ',') @> ARRAY['pg_stat_statements'] AS stats_loaded
FROM pg_settings WHERE name = 'shared_preload_libraries';
 stats_loaded
---------------------------------------------------------------------
 t
(1 row)

DROP EXTENSION IF EXISTS pg_stat_statements;
NOTICE:  extension "pg_stat_statements" does not exist, skipping
-- verify it is not loaded yet
SELECT extname FROM pg_extension WHERE extname = 'pg_stat_statements';
 extname
---------------------------------------------------------------------
(0 rows)

-- this should error out since extension is not created yet
SELECT * FROM citus_stat_statements;
ERROR:  pg_stat_statements is not installed
HINT:  install pg_stat_statements extension and try again
CONTEXT:  PL/pgSQL function citus_stat_statements() line XX at RAISE
-- create extension if available
SELECT CASE WHEN COUNT(*) > 0 THEN
    'CREATE EXTENSION  pg_stat_statements'
ELSE 'SELECT false as pg_stat_statements_available' END
AS create_cmd FROM pg_available_extensions()
WHERE name = 'pg_stat_statements'
\gset
:create_cmd;
CREATE FUNCTION normalize_query_string(query_string text)
	RETURNS TEXT
    LANGUAGE plpgsql
    AS $function$
BEGIN
	RETURN rtrim(regexp_replace(query_string, '\$\d+', '?', 'g'), ';');
END;
$function$;
-- verify citus stat statements reset
SELECT citus_stat_statements_reset();
 citus_stat_statements_reset
---------------------------------------------------------------------

(1 row)

CREATE TABLE test(a int);
SELECT create_distributed_table('test','a');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

insert into test values(1);
select query, calls from citus_stat_statements();
            query            | calls
---------------------------------------------------------------------
 insert into test values($1) |     1
(1 row)

\if :server_version_ge_14
SET compute_query_id = 'off';
\else
set citus.stat_statements_track = 'none';
\endif
-- for pg >= 14, since compute_query_id is off, this insert
-- shouldn't be tracked
-- for pg < 14, we disable it explicitly so that we don't need
-- to add an alternative output file.
insert into test values(1);
select query, calls from citus_stat_statements();
            query            | calls
---------------------------------------------------------------------
 insert into test values($1) |     1
(1 row)

\if :server_version_ge_14
SET compute_query_id = 'on';
\else
RESET citus.stat_statements_track;
\endif
SELECT citus_stat_statements_reset();
 citus_stat_statements_reset
---------------------------------------------------------------------

(1 row)

-- it should now succeed, but with empty result
SELECT normalize_query_string(query) FROM citus_stat_statements;
 normalize_query_string
---------------------------------------------------------------------
(0 rows)

-- run some queries
SELECT count(*) FROM lineitem_hash_part;
 count
---------------------------------------------------------------------
     0
(1 row)

SELECT count(*) FROM lineitem_hash_part;
 count
---------------------------------------------------------------------
     0
(1 row)

SELECT l_orderkey FROM lineitem_hash_part;
 l_orderkey
---------------------------------------------------------------------
(0 rows)

SELECT l_orderkey FROM lineitem_hash_part WHERE l_orderkey > 100;
 l_orderkey
---------------------------------------------------------------------
(0 rows)

SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = 4;
 count
---------------------------------------------------------------------
     0
(1 row)

SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = 4;
 count
---------------------------------------------------------------------
     0
(1 row)

SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = 1;
 count
---------------------------------------------------------------------
     0
(1 row)

SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = 1200;
 count
---------------------------------------------------------------------
     0
(1 row)

SELECT normalize_query_string(query), executor, partition_key, calls
FROM citus_stat_statements
ORDER BY 1, 2, 3, 4;
                     normalize_query_string                     | executor | partition_key | calls
---------------------------------------------------------------------
 SELECT count(*) FROM lineitem_hash_part                        | adaptive |               |     2
 SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = ?   | adaptive | 1             |     1
 SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = ?   | adaptive | 1200          |     1
 SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = ?   | adaptive | 4             |     2
 SELECT l_orderkey FROM lineitem_hash_part                      | adaptive |               |     1
 SELECT l_orderkey FROM lineitem_hash_part WHERE l_orderkey > ? | adaptive |               |     1
(6 rows)

-- test GUC citus.stat_statements_track
SET citus.stat_statements_track TO 'none';
-- this shouldn't increment the call count since tracking is disabled
SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = 4;
 count
---------------------------------------------------------------------
     0
(1 row)

-- this should give the same output as above, since the last query is not counted
SELECT normalize_query_string(query), executor, partition_key, calls
FROM citus_stat_statements
ORDER BY 1, 2, 3, 4;
                     normalize_query_string                     | executor | partition_key | calls
---------------------------------------------------------------------
 SELECT count(*) FROM lineitem_hash_part                        | adaptive |               |     2
 SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = ?   | adaptive | 1             |     1
 SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = ?   | adaptive | 1200          |     1
 SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = ?   | adaptive | 4             |     2
 SELECT l_orderkey FROM lineitem_hash_part                      | adaptive |               |     1
 SELECT l_orderkey FROM lineitem_hash_part WHERE l_orderkey > ? | adaptive |               |     1
(6 rows)

-- reset the GUC to track stats
SET citus.stat_statements_track TO 'all';
-- reset pg_stat_statements and verify it also cleans citus_stat_statements output
-- verify that entries are actually removed from citus_stat_statements
SELECT pg_stat_statements_reset();
 pg_stat_statements_reset
---------------------------------------------------------------------

(1 row)

SELECT * FROM citus_stat_statements;
 queryid | userid | dbid | query | executor | partition_key | calls
---------------------------------------------------------------------
(0 rows)

-- run some queries
SELECT count(*) FROM lineitem_hash_part;
 count
---------------------------------------------------------------------
     0
(1 row)

SELECT count(*) FROM lineitem_hash_part;
 count
---------------------------------------------------------------------
     0
(1 row)

SELECT l_orderkey FROM lineitem_hash_part;
 l_orderkey
---------------------------------------------------------------------
(0 rows)

SELECT l_orderkey FROM lineitem_hash_part WHERE l_orderkey > 100;
 l_orderkey
---------------------------------------------------------------------
(0 rows)

SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = 4;
 count
---------------------------------------------------------------------
     0
(1 row)

SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = 4;
 count
---------------------------------------------------------------------
     0
(1 row)

SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = 1;
 count
---------------------------------------------------------------------
     0
(1 row)

SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = 1200;
 count
---------------------------------------------------------------------
     0
(1 row)

-- show current list, and reset pg_stat_statements
SELECT normalize_query_string(query), executor, partition_key, calls
FROM citus_stat_statements
ORDER BY 1, 2, 3, 4;
                     normalize_query_string                     | executor | partition_key | calls
---------------------------------------------------------------------
 SELECT count(*) FROM lineitem_hash_part                        | adaptive |               |     2
 SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = ?   | adaptive | 1             |     1
 SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = ?   | adaptive | 1200          |     1
 SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = ?   | adaptive | 4             |     2
 SELECT l_orderkey FROM lineitem_hash_part                      | adaptive |               |     1
 SELECT l_orderkey FROM lineitem_hash_part WHERE l_orderkey > ? | adaptive |               |     1
(6 rows)

SELECT pg_stat_statements_reset();
 pg_stat_statements_reset
---------------------------------------------------------------------

(1 row)

SELECT count(*) FROM lineitem_hash_part;
 count
---------------------------------------------------------------------
     0
(1 row)

SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = 4;
 count
---------------------------------------------------------------------
     0
(1 row)

-- verify although pg_stat_statements was reset, some call counts are not
-- if a query is re-issued between pg_stat_statements_reset() and citus_stat_statements()
-- its call count is preserved.
SELECT normalize_query_string(query), executor, partition_key, calls
FROM citus_stat_statements
ORDER BY 1, 2, 3, 4;
                    normalize_query_string                    | executor | partition_key | calls
---------------------------------------------------------------------
 SELECT count(*) FROM lineitem_hash_part                      | adaptive |               |     3
 SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = ? | adaptive | 1             |     1
 SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = ? | adaptive | 1200          |     1
 SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = ? | adaptive | 4             |     3
(4 rows)

-- citus_stat_statements_reset() must be called to reset call counts
SELECT citus_stat_statements_reset();
 citus_stat_statements_reset
---------------------------------------------------------------------

(1 row)

SELECT count(*) FROM lineitem_hash_part;
 count
---------------------------------------------------------------------
     0
(1 row)

SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = 4;
 count
---------------------------------------------------------------------
     0
(1 row)

-- verify citus stats has only 2 rows
SELECT normalize_query_string(query), executor, partition_key, calls
FROM citus_stat_statements
ORDER BY 1, 2, 3, 4;
                    normalize_query_string                    | executor | partition_key | calls
---------------------------------------------------------------------
 SELECT count(*) FROM lineitem_hash_part                      | adaptive |               |     1
 SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = ? | adaptive | 4             |     1
(2 rows)

-- create test tables to run update/delete scenarios
CREATE TABLE stat_test_text(user_id text, value int);
CREATE TABLE stat_test_bigint(user_id bigint, value int);
SELECT citus_stat_statements_reset();
 citus_stat_statements_reset
---------------------------------------------------------------------

(1 row)

-- verify regular tables are not included in citus_stat_statements
SELECT * FROM stat_test_text;
 user_id | value
---------------------------------------------------------------------
(0 rows)

SELECT * FROM stat_test_bigint WHERE user_id = 1::bigint;
 user_id | value
---------------------------------------------------------------------
(0 rows)

SELECT normalize_query_string(query), executor, partition_key, calls
FROM citus_stat_statements
ORDER BY 1, 2, 3, 4;
 normalize_query_string | executor | partition_key | calls
---------------------------------------------------------------------
(0 rows)

SELECT create_distributed_table('stat_test_text', 'user_id');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT create_distributed_table('stat_test_bigint', 'user_id');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT normalize_query_string(query), executor, partition_key, calls
FROM citus_stat_statements
ORDER BY 1, 2, 3, 4;
 normalize_query_string | executor | partition_key | calls
---------------------------------------------------------------------
(0 rows)

SELECT * FROM stat_test_text;
 user_id | value
---------------------------------------------------------------------
(0 rows)

SELECT * FROM stat_test_text WHERE user_id = 'me';
 user_id | value
---------------------------------------------------------------------
(0 rows)

SELECT * FROM stat_test_bigint;
 user_id | value
---------------------------------------------------------------------
(0 rows)

SELECT * FROM stat_test_bigint WHERE user_id = 2::bigint;
 user_id | value
---------------------------------------------------------------------
(0 rows)

SELECT normalize_query_string(query), executor, partition_key, calls
FROM citus_stat_statements
ORDER BY 1, 2, 3, 4;
                  normalize_query_string                  | executor | partition_key | calls
---------------------------------------------------------------------
 SELECT * FROM stat_test_bigint                           | adaptive |               |     1
 SELECT * FROM stat_test_bigint WHERE user_id = ?::bigint | adaptive | 2             |     1
 SELECT * FROM stat_test_text                             | adaptive |               |     1
 SELECT * FROM stat_test_text WHERE user_id = ?           | adaptive | me            |     1
(4 rows)

-- insert some rows and check stats
INSERT INTO stat_test_bigint VALUES (1, 1);
INSERT INTO stat_test_bigint VALUES (7, 1);
INSERT INTO stat_test_bigint VALUES (7, 1), (2,3);
INSERT INTO stat_test_bigint VALUES (8, 1), (8,3);
SELECT normalize_query_string(query), executor, partition_key, calls
FROM citus_stat_statements
ORDER BY 1, 2, 3, 4;
                  normalize_query_string                  | executor | partition_key | calls
---------------------------------------------------------------------
 INSERT INTO stat_test_bigint VALUES (?, ?)               | adaptive | 1             |     1
 INSERT INTO stat_test_bigint VALUES (?, ?)               | adaptive | 7             |     1
 INSERT INTO stat_test_bigint VALUES (?, ?), (?,?)        | adaptive | 8             |     1
 INSERT INTO stat_test_bigint VALUES (?, ?), (?,?)        | adaptive |               |     1
 SELECT * FROM stat_test_bigint                           | adaptive |               |     1
 SELECT * FROM stat_test_bigint WHERE user_id = ?::bigint | adaptive | 2             |     1
 SELECT * FROM stat_test_text                             | adaptive |               |     1
 SELECT * FROM stat_test_text WHERE user_id = ?           | adaptive | me            |     1
(8 rows)

-- delete some rows and check stats
SELECT citus_stat_statements_reset();
 citus_stat_statements_reset
---------------------------------------------------------------------

(1 row)

DELETE FROM stat_test_bigint WHERE value > 1000;
DELETE FROM stat_test_bigint WHERE value > 1200;
DELETE FROM stat_test_bigint WHERE user_id > 1000;
DELETE FROM stat_test_bigint WHERE user_id = 1000;
DELETE FROM stat_test_bigint WHERE user_id = 1000;
SELECT normalize_query_string(query), executor, partition_key, calls
FROM citus_stat_statements
ORDER BY 1, 2, 3, 4;
             normalize_query_string             | executor | partition_key | calls
---------------------------------------------------------------------
 DELETE FROM stat_test_bigint WHERE user_id = ? | adaptive | 1000          |     2
 DELETE FROM stat_test_bigint WHERE user_id > ? | adaptive |               |     1
 DELETE FROM stat_test_bigint WHERE value > ?   | adaptive |               |     2
(3 rows)

-- update some rows and check stats
SELECT citus_stat_statements_reset();
 citus_stat_statements_reset
---------------------------------------------------------------------

(1 row)

UPDATE stat_test_bigint SET value = 300 WHERE value = 3000;
UPDATE stat_test_bigint SET value = 320 WHERE value = 3200;
UPDATE stat_test_bigint SET value = 300 WHERE user_id = 3;
UPDATE stat_test_bigint SET value = 300 WHERE user_id = 3;
UPDATE stat_test_bigint SET value = 3000 WHERE user_id > 500;
SELECT normalize_query_string(query), executor, partition_key, calls
FROM citus_stat_statements
ORDER BY 1, 2, 3, 4;
                 normalize_query_string                  | executor | partition_key | calls
---------------------------------------------------------------------
 UPDATE stat_test_bigint SET value = ? WHERE user_id = ? | adaptive | 3             |     2
 UPDATE stat_test_bigint SET value = ? WHERE user_id > ? | adaptive |               |     1
 UPDATE stat_test_bigint SET value = ? WHERE value = ?   | adaptive |               |     2
(3 rows)

-- test joins
CREATE TABLE stat_test_bigint_other(LIKE stat_test_bigint);
SELECT create_distributed_table('stat_test_bigint_other', 'user_id');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT citus_stat_statements_reset();
 citus_stat_statements_reset
---------------------------------------------------------------------

(1 row)

INSERT INTO stat_test_bigint_other SELECT * FROM stat_test_bigint;
INSERT INTO stat_test_bigint_other SELECT * FROM stat_test_bigint WHERE user_id = 3;
INSERT INTO stat_test_bigint_other SELECT * FROM stat_test_bigint WHERE user_id = 3;
SELECT count(*) FROM stat_test_bigint b JOIN stat_test_bigint_other o USING (user_id);
 count
---------------------------------------------------------------------
    10
(1 row)

SELECT count(*) FROM stat_test_bigint b JOIN stat_test_bigint_other o USING (user_id);
 count
---------------------------------------------------------------------
    10
(1 row)

SELECT count(*) FROM stat_test_bigint b JOIN stat_test_bigint_other o USING (user_id)
WHERE b.user_id = 3;
 count
---------------------------------------------------------------------
     0
(1 row)

SELECT count(*) FROM stat_test_bigint b JOIN stat_test_bigint_other o USING (user_id)
WHERE b.user_id = 3;
 count
---------------------------------------------------------------------
     0
(1 row)

SELECT count(*) FROM stat_test_bigint b JOIN stat_test_bigint_other o USING (user_id)
WHERE o.user_id = 3;
 count
---------------------------------------------------------------------
     0
(1 row)

SELECT normalize_query_string(query), executor, partition_key, calls FROM citus_stat_statements
ORDER BY 1, 2, 3, 4;
                                normalize_query_string                                 | executor | partition_key | calls
---------------------------------------------------------------------
 INSERT INTO stat_test_bigint_other SELECT * FROM stat_test_bigint                     | adaptive |               |     1
 INSERT INTO stat_test_bigint_other SELECT * FROM stat_test_bigint WHERE user_id = ?   | adaptive |               |     2
 SELECT count(*) FROM stat_test_bigint b JOIN stat_test_bigint_other o USING (user_id) | adaptive |               |     2
 SELECT count(*) FROM stat_test_bigint b JOIN stat_test_bigint_other o USING (user_id)+| adaptive | 3             |     2
 WHERE b.user_id = ?                                                                   |          |               |
 SELECT count(*) FROM stat_test_bigint b JOIN stat_test_bigint_other o USING (user_id)+| adaptive | 3             |     1
 WHERE o.user_id = ?                                                                   |          |               |
(5 rows)

-- test reference table
CREATE TABLE stat_test_reference(LIKE stat_test_bigint);
SELECT create_reference_table('stat_test_reference');
 create_reference_table
---------------------------------------------------------------------

(1 row)

INSERT INTO stat_test_reference SELECT user_id, count(*) FROM stat_test_bigint GROUP BY user_id;
SELECT citus_stat_statements_reset();
 citus_stat_statements_reset
---------------------------------------------------------------------

(1 row)

SELECT count(*) FROM stat_test_reference;
 count
---------------------------------------------------------------------
     4
(1 row)

SELECT count(*) FROM stat_test_reference WHERE user_id = 2;
 count
---------------------------------------------------------------------
     1
(1 row)

SELECT count(*) FROM stat_test_reference WHERE user_id = 2;
 count
---------------------------------------------------------------------
     1
(1 row)

SELECT count(*) FROM stat_test_bigint b JOIN stat_test_reference r USING (user_id);
 count
---------------------------------------------------------------------
     6
(1 row)

SELECT count(*) FROM stat_test_bigint b JOIN stat_test_reference r USING (user_id);
 count
---------------------------------------------------------------------
     6
(1 row)

SELECT count(*) FROM stat_test_bigint b JOIN stat_test_reference r USING (user_id)
WHERE b.user_id = 1;
 count
---------------------------------------------------------------------
     1
(1 row)

SELECT count(*) FROM stat_test_bigint b JOIN stat_test_reference r USING (user_id)
WHERE b.user_id = 1 and r.value > 0;
 count
---------------------------------------------------------------------
     1
(1 row)

SELECT count(*) FROM stat_test_bigint b JOIN stat_test_reference r USING (user_id)
WHERE r.user_id = 1;
 count
---------------------------------------------------------------------
     1
(1 row)

SELECT normalize_query_string(query), executor, partition_key, calls
FROM citus_stat_statements
ORDER BY 1, 2, 3, 4;
                               normalize_query_string                               | executor | partition_key | calls
---------------------------------------------------------------------
 SELECT count(*) FROM stat_test_bigint b JOIN stat_test_reference r USING (user_id) | adaptive |               |     2
 SELECT count(*) FROM stat_test_bigint b JOIN stat_test_reference r USING (user_id)+| adaptive | 1             |     1
 WHERE b.user_id = ?                                                                |          |               |
 SELECT count(*) FROM stat_test_bigint b JOIN stat_test_reference r USING (user_id)+| adaptive | 1             |     1
 WHERE b.user_id = ? and r.value > ?                                                |          |               |
 SELECT count(*) FROM stat_test_bigint b JOIN stat_test_reference r USING (user_id)+| adaptive | 1             |     1
 WHERE r.user_id = ?                                                                |          |               |
 SELECT count(*) FROM stat_test_reference                                           | adaptive |               |     1
 SELECT count(*) FROM stat_test_reference WHERE user_id = ?                         | adaptive |               |     2
(6 rows)

-- non-stats role should only see its own entries, even when calling citus_query_stats directly
CREATE USER nostats;
GRANT SELECT ON TABLE lineitem_hash_part TO nostats;
SET ROLE nostats;
SELECT count(*) FROM lineitem_hash_part WHERE l_orderkey = 2;
 count
---------------------------------------------------------------------
     0
(1 row)

SELECT partition_key FROM citus_query_stats();
 partition_key
---------------------------------------------------------------------
 2
(1 row)

RESET ROLE;
-- stats-role/superuser should be able to see entries belonging to other users
SELECT partition_key FROM citus_query_stats() WHERE partition_key = '2';
 partition_key
---------------------------------------------------------------------
 2
(1 row)

-- drop pg_stat_statements and verify citus_stat_statement does not work anymore
DROP extension pg_stat_statements;
SELECT normalize_query_string(query), executor, partition_key, calls
FROM citus_stat_statements
ORDER BY 1, 2, 3, 4;
ERROR:  pg_stat_statements is not installed
HINT:  install pg_stat_statements extension and try again
CONTEXT:  PL/pgSQL function citus_stat_statements() line XX at RAISE
-- drop created tables
DROP TABLE stat_test_text, stat_test_bigint, stat_test_bigint_other, stat_test_reference;
DROP FUNCTION normalize_query_string(text);
\if :server_version_ge_14
SET compute_query_id = 'off';
\endif