Start a tx block on stmt sampled within an explicit block

expected/extended.out

@@ -228,21 +228,22 @@ BEGIN;
 
 COMMIT;
 SELECT span_type, span_operation, parameters, lvl FROM peek_ordered_spans WHERE trace_id='00000000000000000000000000000002';
-  span_type   |  span_operation   | parameters | lvl 
---------------+-------------------+------------+-----
- Select query | SELECT $1         | {1}        |   1
- Planner      | Planner           |            |   2
- ExecutorRun  | ExecutorRun       |            |   2
- Result       | Result            |            |   3
- Select query | SELECT $1, $2     | {1,2}      |   1
- Planner      | Planner           |            |   2
- ExecutorRun  | ExecutorRun       |            |   2
- Result       | Result            |            |   3
- Select query | SELECT $1, $2, $3 | {1,2,3}    |   1
- Planner      | Planner           |            |   2
- ExecutorRun  | ExecutorRun       |            |   2
- Result       | Result            |            |   3
-(12 rows)
+    span_type     |  span_operation   | parameters | lvl 
+------------------+-------------------+------------+-----
+ TransactionBlock | TransactionBlock  |            |   1
+ Select query     | SELECT $1         | {1}        |   2
+ Planner          | Planner           |            |   3
+ ExecutorRun      | ExecutorRun       |            |   3
+ Result           | Result            |            |   4
+ Select query     | SELECT $1, $2     | {1,2}      |   2
+ Planner          | Planner           |            |   3
+ ExecutorRun      | ExecutorRun       |            |   3
+ Result           | Result            |            |   4
+ Select query     | SELECT $1, $2, $3 | {1,2,3}    |   2
+ Planner          | Planner           |            |   3
+ ExecutorRun      | ExecutorRun       |            |   3
+ Result           | Result            |            |   4
+(13 rows)
 
 -- Test tracing only subset of individual statements with extended protocol
 BEGIN;
@@ -272,17 +273,18 @@ SELECT $1 \bind 1 \g
 
 COMMIT;
 SELECT span_type, span_operation, parameters, lvl FROM peek_ordered_spans WHERE trace_id='00000000000000000000000000000003';
-  span_type   |  span_operation   | parameters | lvl 
---------------+-------------------+------------+-----
- Select query | SELECT $1         | {1}        |   1
- Planner      | Planner           |            |   2
- ExecutorRun  | ExecutorRun       |            |   2
- Result       | Result            |            |   3
- Select query | SELECT $1, $2, $3 | {1,2,3}    |   1
- Planner      | Planner           |            |   2
- ExecutorRun  | ExecutorRun       |            |   2
- Result       | Result            |            |   3
-(8 rows)
+    span_type     |  span_operation   | parameters | lvl 
+------------------+-------------------+------------+-----
+ TransactionBlock | TransactionBlock  |            |   1
+ Select query     | SELECT $1         | {1}        |   2
+ Planner          | Planner           |            |   3
+ ExecutorRun      | ExecutorRun       |            |   3
+ Result           | Result            |            |   4
+ Select query     | SELECT $1, $2, $3 | {1,2,3}    |   2
+ Planner          | Planner           |            |   3
+ ExecutorRun      | ExecutorRun       |            |   3
+ Result           | Result            |            |   4
+(9 rows)
 
 -- Test tracing the whole transaction with extended protocol with BEGIN sent through extended protocol
 /*dddbs='postgres.db',traceparent='00-00000000000000000000000000000004-0000000000000004-01'*/ BEGIN \bind \g

expected/parallel.out

@@ -32,8 +32,10 @@ set local pg_tracing.trace_parallel_workers = false;
 (1 row)
 
 commit;
--- Get root top span id
-SELECT span_id as root_span_id from pg_tracing_peek_spans where span_type='Select query' and trace_id='00000000000000000000000000000001' and parent_id='0000000000000001' \gset
+-- get tx block
+select span_id as tx_block_id from pg_tracing_peek_spans where span_type='TransactionBlock' and trace_id='00000000000000000000000000000001' and parent_id='0000000000000001' \gset
+-- get root top span id
+select span_id as root_span_id from pg_tracing_peek_spans where span_type='Select query' and trace_id='00000000000000000000000000000001' and parent_id=:'tx_block_id' \gset
 -- Get executor top span id
 SELECT span_id as executor_span_id from pg_tracing_peek_spans where span_operation='ExecutorRun' and trace_id='00000000000000000000000000000001' and parent_id=:'root_span_id' \gset
 -- Check the select spans that are attached to the root top span

expected/transaction.out

@@ -86,17 +86,18 @@ SELECT count(distinct(trace_id)) = 1 FROM pg_tracing_peek_spans;
 (1 row)
 
 SELECT span_type, span_operation, lvl FROM peek_ordered_spans;
-  span_type   | span_operation | lvl 
---------------+----------------+-----
- 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
-(8 rows)
+    span_type     |  span_operation  | lvl 
+------------------+------------------+-----
+ TransactionBlock | TransactionBlock |   1
+ Select query     | SELECT $1;       |   2
+ Planner          | Planner          |   3
+ ExecutorRun      | ExecutorRun      |   3
+ Result           | Result           |   4
+ Select query     | SELECT $1;       |   2
+ Planner          | Planner          |   3
+ ExecutorRun      | ExecutorRun      |   3
+ Result           | Result           |   4
+(9 rows)
 
 CALL clean_spans();
 -- Test with a 0 parent_id in the middle of a transaction
@@ -134,17 +135,18 @@ SELECT count(distinct(trace_id)) = 1, count(distinct(parent_id)) = 1 FROM peek_o
 (1 row)
 
 SELECT span_type, span_operation, lvl FROM peek_ordered_spans;
-  span_type   | span_operation | lvl 
---------------+----------------+-----
- 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
-(8 rows)
+    span_type     |  span_operation  | lvl 
+------------------+------------------+-----
+ TransactionBlock | TransactionBlock |   1
+ Select query     | SELECT $1;       |   2
+ Planner          | Planner          |   3
+ ExecutorRun      | ExecutorRun      |   3
+ Result           | Result           |   4
+ Select query     | SELECT $1;       |   2
+ Planner          | Planner          |   3
+ ExecutorRun      | ExecutorRun      |   3
+ Result           | Result           |   4
+(9 rows)
 
 CALL clean_spans();
 -- Test modification within transaction block

sql/parallel.sql

@@ -15,8 +15,10 @@ set local pg_tracing.trace_parallel_workers = false;
 /*dddbs='postgres.db',traceparent='00-00000000000000000000000000000004-0000000000000004-01'*/ select 4 from pg_class limit 1;
 commit;
 
--- Get root top span id
-SELECT span_id as root_span_id from pg_tracing_peek_spans where span_type='Select query' and trace_id='00000000000000000000000000000001' and parent_id='0000000000000001' \gset
+-- get tx block
+select span_id as tx_block_id from pg_tracing_peek_spans where span_type='TransactionBlock' and trace_id='00000000000000000000000000000001' and parent_id='0000000000000001' \gset
+-- get root top span id
+select span_id as root_span_id from pg_tracing_peek_spans where span_type='Select query' and trace_id='00000000000000000000000000000001' and parent_id=:'tx_block_id' \gset
 -- Get executor top span id
 SELECT span_id as executor_span_id from pg_tracing_peek_spans where span_operation='ExecutorRun' and trace_id='00000000000000000000000000000001' and parent_id=:'root_span_id' \gset
 

src/pg_tracing.c

@@ -1442,35 +1442,45 @@ initialise_span_context(SpanContext * span_context,
 		span_context->query_id = query->queryId;
 }
 
+static bool
+should_start_tx_block(const Node *utilityStmt)
+{
+	TransactionStmt *stmt;
+    if (GetCurrentTransactionStartTimestamp() != GetCurrentStatementStartTimestamp())
+        /* There's an ongoing tx block, we can create the matching span */
+        return true;
+    if (utilityStmt != NULL && nodeTag(utilityStmt) == T_TransactionStmt) {
+        stmt = (TransactionStmt *) utilityStmt;
+        /* If we have an explicit BEGIN statement, start a tx block */
+        return (stmt->kind == TRANS_STMT_BEGIN);
+    }
+    return false;
+}
+
 /*
  * Start a tx_block span if we have an explicit BEGIN utility statement
  */
 static void
 start_tx_block_span(const Node *utilityStmt, SpanContext * span_context)
 {
-	TransactionStmt *stmt;
-
 	if (nested_level > 0)
 		return;
-	if (utilityStmt == NULL || nodeTag(utilityStmt) != T_TransactionStmt)
-		return;
 
-	stmt = (TransactionStmt *) utilityStmt;
-	if (stmt->kind != TRANS_STMT_BEGIN)
-		/* Not an explicit begin, bail out */
+	if (tx_block_span.span_id > 0)
+		/* We already have an ongoing tx_block span */
 		return;
 
+    if (!should_start_tx_block(utilityStmt))
+        /* Not a candidate to start tx block */
+        return;
+
 	/*
 	 * ProcessUtility may start a new transaction with a BEGIN. We want to
 	 * prevent tracing to be cleared if next parse doesn't have the same lxid
 	 * so we udpate the latest_lxid observed.
 	 */
 	update_latest_lxid();
 
-	if (tx_block_span.span_id > 0)
-		/* We already have an ongoing tx_block span */
-		return;
-
 	/* We have an explicit BEGIN, start a transaction block span */
 	begin_span(span_context->traceparent->trace_id, &tx_block_span,
 			   SPAN_TRANSACTION_BLOCK, NULL, span_context->traceparent->parent_id,