-- Only trace queries with sample flag
SET pg_tracing.sample_rate = 0.0;
SET pg_tracing.caller_sample_rate = 1.0;
-- Set tracecontext at the start of the transaction
/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000001-01'*/ BEGIN; SELECT 1; COMMIT;
 ?column? 
----------
        1
(1 row)

SELECT span_type, span_operation, lvl FROM peek_ordered_spans where trace_id='00000000000000000000000000000001' AND span_type!='Commit';
    span_type     |  span_operation  | lvl 
------------------+------------------+-----
 TransactionBlock | TransactionBlock |   0
 Utility query    | BEGIN;           |   1
 ProcessUtility   | ProcessUtility   |   2
 Select query     | SELECT $1        |   1
 Planner          | Planner          |   2
 ExecutorRun      | ExecutorRun      |   2
 Result           | Result           |   3
 Utility query    | COMMIT;          |   1
 ProcessUtility   | ProcessUtility   |   2
(9 rows)

CALL clean_spans();
-- Test with override of the trace context in the middle of a transaction
/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000002-0000000000000002-01'*/ BEGIN;
/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000003-0000000000000003-01'*/ SELECT 1; COMMIT;
 ?column? 
----------
        1
(1 row)

SELECT span_type, span_operation, lvl FROM peek_ordered_spans where trace_id='00000000000000000000000000000002' AND span_type!='Commit';
    span_type     |  span_operation  | lvl 
------------------+------------------+-----
 TransactionBlock | TransactionBlock |   0
 Utility query    | BEGIN;           |   1
 ProcessUtility   | ProcessUtility   |   2
 Utility query    | COMMIT;          |   1
 ProcessUtility   | ProcessUtility   |   2
(5 rows)

SELECT span_type, span_operation, lvl FROM peek_ordered_spans where trace_id='00000000000000000000000000000003';
  span_type   | span_operation | lvl 
--------------+----------------+-----
 Select query | SELECT $1      |   1
 Planner      | Planner        |   2
 ExecutorRun  | ExecutorRun    |   2
 Result       | Result         |   3
(4 rows)

CALL clean_spans();
-- Test with a 0 traceid in the middle of a transaction
BEGIN;
SELECT 1;
 ?column? 
----------
        1
(1 row)

/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000000-0000000000000001-01'*/ SELECT 2;
 ?column? 
----------
        2
(1 row)

SELECT 3;
 ?column? 
----------
        3
(1 row)

/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000000-0000000000000001-01'*/ SELECT 4;
 ?column? 
----------
        4
(1 row)

END;
-- Only one trace id should have been generated
SELECT count(distinct(trace_id)) = 1 FROM pg_tracing_peek_spans;
 ?column? 
----------
 t
(1 row)

SELECT span_type, span_operation, lvl FROM peek_ordered_spans;
    span_type     |  span_operation  | lvl 
------------------+------------------+-----
 TransactionBlock | TransactionBlock |   0
 Select query     | SELECT $1        |   1
 Planner          | Planner          |   2
 ExecutorRun      | ExecutorRun      |   2
 Result           | Result           |   3
 Select query     | SELECT $1        |   1
 Planner          | Planner          |   2
 ExecutorRun      | ExecutorRun      |   2
 Result           | Result           |   3
(9 rows)

CALL clean_spans();
-- Test with a 0 parent_id in the middle of a transaction
BEGIN;
SELECT 1;
 ?column? 
----------
        1
(1 row)

/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000000-0000000000000000-01'*/ SELECT 2;
 ?column? 
----------
        2
(1 row)

SELECT 3;
 ?column? 
----------
        3
(1 row)

/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000000-0000000000000000-01'*/ SELECT 4;
 ?column? 
----------
        4
(1 row)

END;
-- Only one trace id and parent id should have been generated
SELECT count(distinct(trace_id)) = 1, count(distinct(parent_id)) = 1 FROM peek_ordered_spans WHERE lvl=1;
 ?column? | ?column? 
----------+----------
 t        | t
(1 row)

SELECT span_type, span_operation, lvl FROM peek_ordered_spans;
    span_type     |  span_operation  | lvl 
------------------+------------------+-----
 TransactionBlock | TransactionBlock |   0
 Select query     | SELECT $1        |   1
 Planner          | Planner          |   2
 ExecutorRun      | ExecutorRun      |   2
 Result           | Result           |   3
 Select query     | SELECT $1        |   1
 Planner          | Planner          |   2
 ExecutorRun      | ExecutorRun      |   2
 Result           | Result           |   3
(9 rows)

CALL clean_spans();
-- Test modification within transaction block
/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000000-01'*/ BEGIN;
INSERT INTO test_modifications(a, b) VALUES (1, 1);
END;
SELECT span_type, span_operation, lvl FROM peek_ordered_spans;
     span_type     |                    span_operation                    | lvl 
-------------------+------------------------------------------------------+-----
 TransactionBlock  | TransactionBlock                                     |   0
 Utility query     | BEGIN;                                               |   1
 ProcessUtility    | ProcessUtility                                       |   2
 Insert query      | INSERT INTO test_modifications(a, b) VALUES ($1, $2) |   1
 Planner           | Planner                                              |   2
 ExecutorRun       | ExecutorRun                                          |   2
 Insert            | Insert on test_modifications                         |   3
 Result            | Result                                               |   4
 Utility query     | END;                                                 |   1
 ProcessUtility    | ProcessUtility                                       |   2
 TransactionCommit | TransactionCommit                                    |   1
(11 rows)

SELECT span_id AS span_tx_block,
        get_epoch(span_start) AS span_tx_block_start,
        get_epoch(span_end) AS span_tx_block_end
		FROM pg_tracing_peek_spans
        WHERE trace_id='00000000000000000000000000000001' AND span_operation='TransactionBlock' \gset
SELECT span_id AS span_commit,
        get_epoch(span_start) as span_commit_start,
        get_epoch(span_end) as span_commit_end
		FROM pg_tracing_peek_spans
        WHERE trace_id='00000000000000000000000000000001'
            AND parent_id = :'span_tx_block'
            AND span_operation='TransactionCommit' \gset
-- Transaction block should end after TransactionCommit span
SELECT :span_tx_block_end >= :span_commit_end;
 ?column? 
----------
 t
(1 row)

CALL clean_spans();
-- Test with transaction block created with sample rate
SET pg_tracing.sample_rate = 1.0;
BEGIN;
/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000001-01'*/ SELECT 1;
 ?column? 
----------
        1
(1 row)

/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000002-01'*/ SELECT 2;
 ?column? 
----------
        2
(1 row)

COMMIT;
SELECT span_type, span_operation, lvl FROM peek_ordered_spans where trace_id='00000000000000000000000000000001';
    span_type     |  span_operation  | lvl 
------------------+------------------+-----
 TransactionBlock | TransactionBlock |   0
 Utility query    | BEGIN;           |   1
 ProcessUtility   | ProcessUtility   |   2
 Select query     | SELECT $1        |   1
 Planner          | Planner          |   2
 ExecutorRun      | ExecutorRun      |   2
 Result           | Result           |   3
 Select query     | SELECT $1        |   1
 Planner          | Planner          |   2
 ExecutorRun      | ExecutorRun      |   2
 Result           | Result           |   3
 Utility query    | COMMIT;          |   1
 ProcessUtility   | ProcessUtility   |   2
(13 rows)

CALL reset_settings();
CALL clean_spans();