| 419 | } |
| 420 | |
| 421 | func TestTraceCopyFrom(t *testing.T) { |
| 422 | t.Parallel() |
| 423 | |
| 424 | tracer := &testTracer{} |
| 425 | |
| 426 | ctr := defaultConnTestRunner |
| 427 | ctr.CreateConfig = func(ctx context.Context, t testing.TB) *pgx.ConnConfig { |
| 428 | config := defaultConnTestRunner.CreateConfig(ctx, t) |
| 429 | config.Tracer = tracer |
| 430 | return config |
| 431 | } |
| 432 | |
| 433 | ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second) |
| 434 | defer cancel() |
| 435 | |
| 436 | pgxtest.RunWithQueryExecModes(ctx, t, ctr, nil, func(ctx context.Context, t testing.TB, conn *pgx.Conn) { |
| 437 | ctx, cancel := context.WithTimeout(ctx, 30*time.Second) |
| 438 | defer cancel() |
| 439 | |
| 440 | traceCopyFromStartCalled := false |
| 441 | tracer.traceCopyFromStart = func(ctx context.Context, conn *pgx.Conn, data pgx.TraceCopyFromStartData) context.Context { |
| 442 | traceCopyFromStartCalled = true |
| 443 | require.Equal(t, pgx.Identifier{"foo"}, data.TableName) |
| 444 | require.Equal(t, []string{"a"}, data.ColumnNames) |
| 445 | return context.WithValue(ctx, ctxKey("fromTraceCopyFromStart"), "foo") |
| 446 | } |
| 447 | |
| 448 | traceCopyFromEndCalled := false |
| 449 | tracer.traceCopyFromEnd = func(ctx context.Context, conn *pgx.Conn, data pgx.TraceCopyFromEndData) { |
| 450 | traceCopyFromEndCalled = true |
| 451 | require.Equal(t, "foo", ctx.Value(ctxKey("fromTraceCopyFromStart"))) |
| 452 | require.Equal(t, `COPY 2`, data.CommandTag.String()) |
| 453 | require.NoError(t, data.Err) |
| 454 | } |
| 455 | |
| 456 | _, err := conn.Exec(ctx, `create temporary table foo(a int4)`) |
| 457 | require.NoError(t, err) |
| 458 | |
| 459 | inputRows := [][]any{ |
| 460 | {int32(1)}, |
| 461 | {nil}, |
| 462 | } |
| 463 | |
| 464 | copyCount, err := conn.CopyFrom(ctx, pgx.Identifier{"foo"}, []string{"a"}, pgx.CopyFromRows(inputRows)) |
| 465 | require.NoError(t, err) |
| 466 | require.EqualValues(t, len(inputRows), copyCount) |
| 467 | require.True(t, traceCopyFromStartCalled) |
| 468 | require.True(t, traceCopyFromEndCalled) |
| 469 | }) |
| 470 | } |
| 471 | |
| 472 | func TestTracePrepare(t *testing.T) { |
| 473 | t.Parallel() |