| 98 | } |
| 99 | |
| 100 | func TestTraceExec(t *testing.T) { |
| 101 | t.Parallel() |
| 102 | |
| 103 | tracer := &testTracer{} |
| 104 | |
| 105 | ctr := defaultConnTestRunner |
| 106 | ctr.CreateConfig = func(ctx context.Context, t testing.TB) *pgx.ConnConfig { |
| 107 | config := defaultConnTestRunner.CreateConfig(ctx, t) |
| 108 | config.Tracer = tracer |
| 109 | return config |
| 110 | } |
| 111 | |
| 112 | ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second) |
| 113 | defer cancel() |
| 114 | |
| 115 | pgxtest.RunWithQueryExecModes(ctx, t, ctr, nil, func(ctx context.Context, t testing.TB, conn *pgx.Conn) { |
| 116 | traceQueryStartCalled := false |
| 117 | tracer.traceQueryStart = func(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryStartData) context.Context { |
| 118 | traceQueryStartCalled = true |
| 119 | require.Equal(t, `select $1::text`, data.SQL) |
| 120 | require.Len(t, data.Args, 1) |
| 121 | require.Equal(t, `testing`, data.Args[0]) |
| 122 | return context.WithValue(ctx, ctxKey("fromTraceQueryStart"), "foo") |
| 123 | } |
| 124 | |
| 125 | traceQueryEndCalled := false |
| 126 | tracer.traceQueryEnd = func(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryEndData) { |
| 127 | traceQueryEndCalled = true |
| 128 | require.Equal(t, "foo", ctx.Value(ctxKey("fromTraceQueryStart"))) |
| 129 | require.Equal(t, `SELECT 1`, data.CommandTag.String()) |
| 130 | require.NoError(t, data.Err) |
| 131 | } |
| 132 | |
| 133 | _, err := conn.Exec(ctx, `select $1::text`, "testing") |
| 134 | require.NoError(t, err) |
| 135 | require.True(t, traceQueryStartCalled) |
| 136 | require.True(t, traceQueryEndCalled) |
| 137 | }) |
| 138 | } |
| 139 | |
| 140 | func TestTraceQuery(t *testing.T) { |
| 141 | t.Parallel() |