| 138 | } |
| 139 | |
| 140 | func TestTraceQuery(t *testing.T) { |
| 141 | t.Parallel() |
| 142 | |
| 143 | tracer := &testTracer{} |
| 144 | |
| 145 | ctr := defaultConnTestRunner |
| 146 | ctr.CreateConfig = func(ctx context.Context, t testing.TB) *pgx.ConnConfig { |
| 147 | config := defaultConnTestRunner.CreateConfig(ctx, t) |
| 148 | config.Tracer = tracer |
| 149 | return config |
| 150 | } |
| 151 | |
| 152 | ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second) |
| 153 | defer cancel() |
| 154 | |
| 155 | pgxtest.RunWithQueryExecModes(ctx, t, ctr, nil, func(ctx context.Context, t testing.TB, conn *pgx.Conn) { |
| 156 | traceQueryStartCalled := false |
| 157 | tracer.traceQueryStart = func(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryStartData) context.Context { |
| 158 | traceQueryStartCalled = true |
| 159 | require.Equal(t, `select $1::text`, data.SQL) |
| 160 | require.Len(t, data.Args, 1) |
| 161 | require.Equal(t, `testing`, data.Args[0]) |
| 162 | return context.WithValue(ctx, ctxKey("fromTraceQueryStart"), "foo") |
| 163 | } |
| 164 | |
| 165 | traceQueryEndCalled := false |
| 166 | tracer.traceQueryEnd = func(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryEndData) { |
| 167 | traceQueryEndCalled = true |
| 168 | require.Equal(t, "foo", ctx.Value(ctxKey("fromTraceQueryStart"))) |
| 169 | require.Equal(t, `SELECT 1`, data.CommandTag.String()) |
| 170 | require.NoError(t, data.Err) |
| 171 | } |
| 172 | |
| 173 | var s string |
| 174 | err := conn.QueryRow(ctx, `select $1::text`, "testing").Scan(&s) |
| 175 | require.NoError(t, err) |
| 176 | require.Equal(t, "testing", s) |
| 177 | require.True(t, traceQueryStartCalled) |
| 178 | require.True(t, traceQueryEndCalled) |
| 179 | }) |
| 180 | } |
| 181 | |
| 182 | func TestTraceBatchNormal(t *testing.T) { |
| 183 | t.Parallel() |