-- Create test function to sample
CREATE OR REPLACE FUNCTION test_function_result(a int, b text) RETURNS void AS
$BODY$
BEGIN
    INSERT INTO pg_tracing_test(a, b) VALUES (a, b);
END;
$BODY$
LANGUAGE plpgsql;
CREATE OR REPLACE FUNCTION test_function_project_set(a int) RETURNS SETOF oid AS
$BODY$
BEGIN
	RETURN QUERY SELECT oid from pg_class where oid = a;
END;
$BODY$
LANGUAGE plpgsql;
CREATE OR REPLACE FUNCTION test_2_nested_levels(a int)  RETURNS SETOF oid AS
$BODY$
BEGIN
    RETURN QUERY SELECT * FROM test_function_project_set(1);
END;
$BODY$
LANGUAGE plpgsql;
-- Trace a statement with a function call
/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000051-0000000000000051-01'*/ select test_function_project_set(1);
 test_function_project_set 
---------------------------
(0 rows)

-- The test function call will generate the following spans (non exhaustive list):
-- +---------------------------------------------------------------------------------+
-- | A: Select test_function_project_set(1);                                         |
-- +-+----------++----------------++-------------------------------------------------+
--   |B: Planner||C: ExecutorStart||D: ExecutorRun                                |
--   +----------++----------------+-+---------------------------------------------+
--                                  |E: ProjectSet                              |
--                                  ++---------+------------------------------+-+
--                                   |F: Result| G: Select a from b where...  |
--                                   +---------+--------------+---------------+
--                                                            |H: ExecutorRun|
--                                                            +--------------+
-- Gather span_id, span start and span end of function call statement
SELECT span_id AS span_a_id,
        get_epoch(span_start) as span_a_start,
        get_epoch(span_end) as span_a_end
		from pg_tracing_peek_spans where parent_id='0000000000000051' AND span_type='Select query' \gset
SELECT span_id AS span_d_id,
        get_epoch(span_start) as span_d_start,
        get_epoch(span_end) as span_d_end
		from pg_tracing_peek_spans where parent_id=:'span_a_id' and span_operation='ExecutorRun' \gset
SELECT span_id AS span_e_id,
        get_epoch(span_start) as span_e_start,
        get_epoch(span_end) as span_e_end
		from pg_tracing_peek_spans where parent_id=:'span_d_id' and span_type='ProjectSet' \gset
SELECT span_id AS span_g_id,
        get_epoch(span_start) as span_g_start,
        get_epoch(span_end) as span_g_end
		from pg_tracing_peek_spans where parent_id=:'span_e_id' and span_type='Result' \gset
SELECT span_id AS span_g_id,
        get_epoch(span_start) as span_g_start,
        get_epoch(span_end) as span_g_end
		from pg_tracing_peek_spans where parent_id=:'span_e_id' and span_type='Select query' \gset
SELECT span_id AS span_h_id,
        get_epoch(span_start) as span_h_start,
        get_epoch(span_end) as span_h_end
		from pg_tracing_peek_spans where parent_id=:'span_g_id' and span_operation='ExecutorRun' \gset
-- Check that spans' start and end are within expection
SELECT :span_a_start <= :span_d_start AS top_query_before_run,
		:span_a_end >= :span_d_end AS top_ends_after_run_end,
		:span_d_start <= :span_e_start AS top_run_starts_before_project,
		:span_d_end >= :span_e_end AS top_run_ends_after_project_end,
		:span_d_end >= :span_h_end AS top_run_ends_before_select_end,
		:span_d_end >= :span_g_end AS top_run_ends_after_nested_run_end;
 top_query_before_run | top_ends_after_run_end | top_run_starts_before_project | top_run_ends_after_project_end | top_run_ends_before_select_end | top_run_ends_after_nested_run_end 
----------------------+------------------------+-------------------------------+--------------------------------+--------------------------------+-----------------------------------
 t                    | t                      | t                             | t                              | t                              | t
(1 row)

SELECT
		:span_g_end >= :span_h_start AS nested_result_ends_before_parse,
		:span_h_end <= :span_g_end AS nested_parse_ends_before_select,
		:span_h_start >= :span_g_start AS run_starts_after_parent_select,
		:span_h_end <= :span_g_end AS run_ends_after_select_end;
 nested_result_ends_before_parse | nested_parse_ends_before_select | run_starts_after_parent_select | run_ends_after_select_end 
---------------------------------+---------------------------------+--------------------------------+---------------------------
 t                               | t                               | t                              | t
(1 row)

-- Check that the root span is the longest one
WITH max_end AS (select max(span_end) from pg_tracing_peek_spans WHERE span_type != 'Commit')
SELECT span_end = max_end.max from pg_tracing_peek_spans, max_end
    where span_id = :'span_a_id';
 ?column? 
----------
 t
(1 row)

-- Check that ExecutorRun is attached to the nested top span
SELECT span_operation, deparse_info from pg_tracing_peek_spans where parent_id=:'span_h_id' order by span_operation;
                   span_operation                   |         deparse_info         
----------------------------------------------------+------------------------------
 IndexOnlyScan using pg_class_oid_index on pg_class | Index Cond: (oid = '1'::oid)
(1 row)

-- Check tracking with top tracking
SET pg_tracing.track = 'top';
/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000052-0000000000000052-01'*/ select test_function_project_set(1);
 test_function_project_set 
---------------------------
(0 rows)

SELECT count(*) from pg_tracing_consume_spans where trace_id='00000000000000000000000000000052';
 count 
-------
     5
(1 row)

-- Check tracking with no tracking
SET pg_tracing.track = 'none';
/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000053-0000000000000053-01'*/ select test_function_project_set(1);
 test_function_project_set 
---------------------------
(0 rows)

SELECT count(*) from pg_tracing_consume_spans where trace_id='00000000000000000000000000000053';
 count 
-------
     0
(1 row)

-- Reset tracking setting
SET pg_tracing.track TO DEFAULT;
-- Create test procedure
CREATE OR REPLACE PROCEDURE sum_one() AS $$
DECLARE
  r int;
BEGIN
  SELECT (10 + 10)::int INTO r;
END; $$ LANGUAGE plpgsql;
-- Test tracking of procedure with utility tracking enabled
SET pg_tracing.track_utility=on;
/*traceparent='00-00000000000000000000000000000054-0000000000000054-01'*/ CALL sum_one();
select span_operation, lvl FROM peek_ordered_spans where trace_id='00000000000000000000000000000054';
    span_operation     | lvl 
-----------------------+-----
 CALL sum_one();       |   0
 ProcessUtility        |   1
 SELECT ($1 + $2)::int |   2
 Planner               |   3
 ExecutorRun           |   3
 Result                |   4
(6 rows)

-- Test again with utility tracking disabled
SET pg_tracing.track_utility=off;
/*traceparent='00-00000000000000000000000000000055-0000000000000055-01'*/ CALL sum_one();
select span_operation, lvl FROM peek_ordered_spans where trace_id='00000000000000000000000000000055';
 span_operation | lvl 
----------------+-----
(0 rows)

-- Create immutable function
CREATE OR REPLACE FUNCTION test_immutable_function(a int) RETURNS oid
AS 'SELECT oid from pg_class where oid = a;'
LANGUAGE sql IMMUTABLE;
/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000056-0000000000000056-01'*/ select test_immutable_function(1);
 test_immutable_function 
-------------------------
                        
(1 row)

select span_operation, lvl FROM peek_ordered_spans where trace_id='00000000000000000000000000000056';
                   span_operation                   | lvl 
----------------------------------------------------+-----
 select test_immutable_function($1)                 |   0
 Planner                                            |   1
 SELECT oid from pg_class where oid = a;            |   2
 Planner                                            |   3
 ExecutorRun                                        |   3
 IndexOnlyScan using pg_class_oid_index on pg_class |   4
 ExecutorRun                                        |   1
 Result                                             |   2
(8 rows)

-- Create function with generate series
CREATE OR REPLACE FUNCTION test_generate_series(IN anyarray, OUT x anyelement)
    RETURNS SETOF anyelement
    LANGUAGE sql
    AS 'select * from pg_catalog.generate_series(array_lower($1, 1), array_upper($1, 1), 1)';
/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000057-0000000000000057-01'*/ select test_generate_series('{1,2,3,4}'::int[]) FROM (VALUES (1,2)) as t;
 test_generate_series 
----------------------
                    1
                    2
                    3
                    4
(4 rows)

SELECT span_id AS span_project_set_id,
        get_epoch(span_start) as span_project_set_start,
        get_epoch(span_end) as span_project_set_end
		from pg_tracing_peek_spans where span_type='ProjectSet' \gset
SELECT span_id AS span_result_id,
        get_epoch(span_start) as span_result_start,
        get_epoch(span_end) as span_result_end
		from pg_tracing_peek_spans where parent_id=:'span_project_set_id' and span_type='Result' \gset
select span_operation, lvl FROM peek_ordered_spans where trace_id='00000000000000000000000000000057';
                                     span_operation                                     | lvl 
----------------------------------------------------------------------------------------+-----
 select test_generate_series($1::int[]) FROM (VALUES ($2,$3)) as t                      |   0
 Planner                                                                                |   1
 ExecutorRun                                                                            |   1
 ProjectSet                                                                             |   2
 Result                                                                                 |   3
 select * from pg_catalog.generate_series(array_lower($1, $2), array_upper($1, $3), $4) |   3
 Planner                                                                                |   4
(7 rows)

-- +-----------------------------------------------------------+
-- | A: Select test_function(1);                               |
-- +-+----------+---+--------------------------------------+---+
--   |B: Planner|   |C: ExecutorRun                        |
--   +----------+   ++-------------------------------------+
--                   |D: Result                           |
--                   +----+-----------------------------+-+
--                        |E: Insert INTO...            |
--                        +---+--------------+----------+
--                            |F: ExecutorRun|
--                            +--------------+
-- Check function with result node
/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000058-0000000000000058-01'*/ select test_function_result(1, 'test');
 test_function_result 
----------------------
 
(1 row)

-- Gather span_id, span start and span end of function call statement
SELECT span_id AS span_a_id,
        get_epoch(span_start) as span_a_start,
        get_epoch(span_end) as span_a_end
		from pg_tracing_peek_spans where parent_id='0000000000000058' AND span_type = 'Select query'  \gset
SELECT span_id AS span_c_id,
        get_epoch(span_start) as span_c_start,
        get_epoch(span_end) as span_c_end
		from pg_tracing_peek_spans where parent_id=:'span_a_id' and span_operation='ExecutorRun' \gset
SELECT span_id AS span_d_id,
        get_epoch(span_start) as span_d_start,
        get_epoch(span_end) as span_d_end
		from pg_tracing_peek_spans where parent_id=:'span_c_id' and span_type='Result' \gset
SELECT span_id AS span_e_id,
        get_epoch(span_start) as span_e_start,
        get_epoch(span_end) as span_e_end
		from pg_tracing_peek_spans where parent_id=:'span_d_id' and span_type='Insert query' \gset
select span_operation, lvl FROM peek_ordered_spans where trace_id='00000000000000000000000000000058';
                 span_operation                  | lvl 
-------------------------------------------------+-----
 select test_function_result($1, $2)             |   0
 Planner                                         |   1
 ExecutorRun                                     |   1
 Result                                          |   2
 INSERT INTO pg_tracing_test(a, b) VALUES (a, b) |   3
 Planner                                         |   4
 ExecutorRun                                     |   4
 Insert on pg_tracing_test                       |   5
 Result                                          |   6
 TransactionCommit                               |   0
(10 rows)

-- Make sure we have 2 query_id associated with the trace
SELECT count(distinct query_id)=2 from pg_tracing_consume_spans where trace_id='00000000000000000000000000000058';
 ?column? 
----------
 t
(1 row)

-- Trace a statement with multiple nested levels
/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000059-0000000000000059-01'*/ select * FROM test_2_nested_levels(1);
 test_2_nested_levels 
----------------------
(0 rows)

select span_operation, lvl FROM peek_ordered_spans where trace_id='00000000000000000000000000000059';
                   span_operation                   | lvl 
----------------------------------------------------+-----
 select * FROM test_2_nested_levels($1)             |   0
 Planner                                            |   1
 ExecutorRun                                        |   1
 FunctionScan on test_2_nested_levels               |   2
 SELECT * FROM test_function_project_set($1)        |   2
 Planner                                            |   3
 ExecutorRun                                        |   3
 FunctionScan on test_function_project_set          |   4
 SELECT oid from pg_class where oid = a             |   4
 Planner                                            |   5
 ExecutorRun                                        |   5
 IndexOnlyScan using pg_class_oid_index on pg_class |   6
(12 rows)

CALL clean_spans();
-- Cleanup
CALL clean_spans();