From e9aad0fb0b55c42f1ff2e2106c65aab8b3da724c Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Sat, 26 Apr 2025 09:55:31 -0500
Subject: [PATCH] Add test for tracer within transaction

https://github.com/jackc/pgx/issues/2304
---
 tracer_test.go | 40 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/tracer_test.go b/tracer_test.go
index a0fea71e..8920313f 100644
--- a/tracer_test.go
+++ b/tracer_test.go
@@ -566,3 +566,43 @@ func TestTraceConnect(t *testing.T) {
 	require.True(t, traceConnectStartCalled)
 	require.True(t, traceConnectEndCalled)
 }
+
+// Ensure tracer runs within a transaction.
+//
+// https://github.com/jackc/pgx/issues/2304
+func TestTraceWithinTx(t *testing.T) {
+	t.Parallel()
+
+	tracer := &testTracer{}
+
+	ctr := defaultConnTestRunner
+	ctr.CreateConfig = func(ctx context.Context, t testing.TB) *pgx.ConnConfig {
+		config := defaultConnTestRunner.CreateConfig(ctx, t)
+		config.Tracer = tracer
+		return config
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
+	defer cancel()
+
+	pgxtest.RunWithQueryExecModes(ctx, t, ctr, nil, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
+		var queries []string
+		tracer.traceQueryStart = func(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryStartData) context.Context {
+			queries = append(queries, data.SQL)
+			return ctx
+		}
+
+		tx, err := conn.Begin(ctx)
+		require.NoError(t, err)
+		defer tx.Rollback(ctx)
+		_, err = tx.Exec(ctx, `select $1::text`, "testing")
+		require.NoError(t, err)
+		err = tx.Commit(ctx)
+		require.NoError(t, err)
+
+		require.Len(t, queries, 3)
+		require.Equal(t, `begin`, queries[0])
+		require.Equal(t, `select $1::text`, queries[1])
+		require.Equal(t, `commit`, queries[2])
+	})
+}