Override existing tx_traceparent when relevant

expected/transaction.out

@@ -187,4 +187,39 @@ SELECT :span_tx_block_end >= :span_commit_end;
 (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 |   1
+ Utility query    | BEGIN;           |   2
+ ProcessUtility   | ProcessUtility   |   3
+ 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
+ Utility query    | COMMIT;          |   2
+ ProcessUtility   | ProcessUtility   |   3
+(13 rows)
+
 CALL reset_settings();
+CALL clean_spans();

sql/transaction.sql

@@ -62,7 +62,15 @@ SELECT span_id AS span_commit,
             AND span_operation='TransactionCommit' \gset
 -- Transaction block should end after TransactionCommit span
 SELECT :span_tx_block_end >= :span_commit_end;
+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;
+/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000002-01'*/ SELECT 2;
+COMMIT;
+SELECT span_type, span_operation, lvl FROM peek_ordered_spans where trace_id='00000000000000000000000000000001';
 
-CALL clean_spans();
 CALL reset_settings();
+CALL clean_spans();

src/pg_tracing.c

@@ -1389,6 +1389,28 @@ end_nested_level(const TimestampTz *input_span_end_time)
 	}
 }
 
+/*
+ * Adjust the tx_traceparent if the latest captured traceparent is more
+ * relevant, i.e. provided from upstream and not automatically generated.
+ */
+static void
+adjust_tx_traceparent(const Traceparent * traceparent)
+{
+	if (traceid_equal(traceparent->trace_id, tx_block_span.trace_id))
+		return;
+	if (!tx_traceparent.generated || traceparent->generated)
+		return;
+
+	/*
+	 * We have a generated tx_traceparent and a provided traceparent, give
+	 * priority to the provided traceparent and amend the existing spans
+	 */
+	tx_block_span.trace_id = traceparent->trace_id;
+	tx_traceparent.trace_id = traceparent->trace_id;
+	for (int i = 0; i < current_trace_spans->end; i++)
+		current_trace_spans->spans[i].trace_id = traceparent->trace_id;
+}
+
 static void
 initialise_span_context(SpanContext * span_context,
 						Traceparent * traceparent,
@@ -1402,6 +1424,7 @@ initialise_span_context(SpanContext * span_context,
 	span_context->traceparent = traceparent;
 	if (!traceid_zero(tx_block_span.trace_id) && nested_level == 0)
 	{
+		adjust_tx_traceparent(span_context->traceparent);
 		/* We have an ongoing transaction block. Use it as parent for level 0 */
 		span_context->traceparent->parent_id = tx_block_span.span_id;
 	}