Giters

jackc / pgx

PostgreSQL driver and toolkit for Go

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Missing `QueryTracer` `TraceQueryEnd()` call in conn `Query()` func

renanferr opened this issue · comments

commented

Describe the bug
The Query Tracer is not working properly with the Query() connection func. The current unit tests are using QueryRow() and actually there are missing calls for the TraceQueryEnd() func for several exec modes.

To Reproduce
Just run conn.Query() with an active query tracer and check if TraceQueryEnd() was called.

I wrote a unit test for this case in #1821 to better demonstrate.

Runnable example to reproduce:

package main

import (
	"context"
	"log"
	"os"

	"github.com/jackc/pgx/v5"
)

type testTracer struct {
	traceQueryStart    func(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryStartData) context.Context
	traceQueryEnd      func(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryEndData)
	traceBatchStart    func(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchStartData) context.Context
	traceBatchQuery    func(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchQueryData)
	traceBatchEnd      func(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchEndData)
	traceCopyFromStart func(ctx context.Context, conn *pgx.Conn, data pgx.TraceCopyFromStartData) context.Context
	traceCopyFromEnd   func(ctx context.Context, conn *pgx.Conn, data pgx.TraceCopyFromEndData)
	tracePrepareStart  func(ctx context.Context, conn *pgx.Conn, data pgx.TracePrepareStartData) context.Context
	tracePrepareEnd    func(ctx context.Context, conn *pgx.Conn, data pgx.TracePrepareEndData)
	traceConnectStart  func(ctx context.Context, data pgx.TraceConnectStartData) context.Context
	traceConnectEnd    func(ctx context.Context, data pgx.TraceConnectEndData)
}

func (tt *testTracer) TraceQueryStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryStartData) context.Context {
	if tt.traceQueryStart != nil {
		return tt.traceQueryStart(ctx, conn, data)
	}
	return ctx
}

func (tt *testTracer) TraceQueryEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryEndData) {
	if tt.traceQueryEnd != nil {
		tt.traceQueryEnd(ctx, conn, data)
	}
}

func (tt *testTracer) TraceBatchStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchStartData) context.Context {
	if tt.traceBatchStart != nil {
		return tt.traceBatchStart(ctx, conn, data)
	}
	return ctx
}

func (tt *testTracer) TraceBatchQuery(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchQueryData) {
	if tt.traceBatchQuery != nil {
		tt.traceBatchQuery(ctx, conn, data)
	}
}

func (tt *testTracer) TraceBatchEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchEndData) {
	if tt.traceBatchEnd != nil {
		tt.traceBatchEnd(ctx, conn, data)
	}
}

func (tt *testTracer) TraceCopyFromStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceCopyFromStartData) context.Context {
	if tt.traceCopyFromStart != nil {
		return tt.traceCopyFromStart(ctx, conn, data)
	}
	return ctx
}

func (tt *testTracer) TraceCopyFromEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceCopyFromEndData) {
	if tt.traceCopyFromEnd != nil {
		tt.traceCopyFromEnd(ctx, conn, data)
	}
}

func (tt *testTracer) TracePrepareStart(ctx context.Context, conn *pgx.Conn, data pgx.TracePrepareStartData) context.Context {
	if tt.tracePrepareStart != nil {
		return tt.tracePrepareStart(ctx, conn, data)
	}
	return ctx
}

func (tt *testTracer) TracePrepareEnd(ctx context.Context, conn *pgx.Conn, data pgx.TracePrepareEndData) {
	if tt.tracePrepareEnd != nil {
		tt.tracePrepareEnd(ctx, conn, data)
	}
}

func (tt *testTracer) TraceConnectStart(ctx context.Context, data pgx.TraceConnectStartData) context.Context {
	if tt.traceConnectStart != nil {
		return tt.traceConnectStart(ctx, data)
	}
	return ctx
}

func (tt *testTracer) TraceConnectEnd(ctx context.Context, data pgx.TraceConnectEndData) {
	if tt.traceConnectEnd != nil {
		tt.traceConnectEnd(ctx, data)
	}
}

func main() {
	var calledStart bool
	var calledEnd bool
	cfg, err := pgx.ParseConfig(os.Getenv("DATABASE_URL"))
	cfg.Tracer = &testTracer{
		traceQueryStart: func(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryStartData) context.Context {
			calledStart = true
			return ctx
		},
		traceQueryEnd: func(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryEndData) {
			calledEnd = true
		},
	}
	conn, err := pgx.ConnectConfig(context.Background(), cfg)
	if err != nil {
		log.Fatal(err)
	}
	defer conn.Close(context.Background())

	rows, err := conn.Query(context.Background(), "select 1")
	if err != nil {
		log.Fatal(err)
	}

	if rows.Err() != nil {
		log.Fatal(rows.Err())
	}

	if !rows.Next() {
		log.Fatal("no rows")
	}

	var d int
	if err := rows.Scan(&d); err != nil {
		log.Fatal(err)
	}

	log.Printf("result %d", d)
	log.Printf("calledStart %v", calledStart)
	log.Printf("calledEnd %v", calledEnd)
}

Expected behavior
The connection Query() func should call TracerQueryStart() and TracerQueryEnd() just like QueryRow(), Exec(), etc.

Actual behavior
Query() func calls TracerQueryStart() but only calls TracerQueryEnd() if this err check doesn't pass. Otherwise, no other calls are made neither in the conn Query() func or the PgConn ExecParams() or ExecPrepared() funcs.

Version

  • Go: go version go1.21.0 darwin/arm64
  • PostgreSQL: PostgreSQL 14.8 (Homebrew) on aarch64-apple-darwin22.4.0, compiled by Apple clang version 14.0.3 (clang-1403.0.22.14.1), 64-bit
  • pgx: v5.5.0

Additional context
I was working on a dd-trace contrib package to integrate datadog with pgx and noticed my tests were failing due to a pgx bug whenever I was working with the Query() func.
Since I already have a branch with unit tests on it I can just fix it as well. Just wanted to make sure this is an actual bug.

commented

The query is never finished in your example. Either Rows.Close() needs to be called explicitly or Rows.Next() needs to be called until it returns false.

Here is your example when it finishes the query.

package main

import (
	"context"
	"log"
	"os"

	"github.com/jackc/pgx/v5"
)

type testTracer struct {
	traceQueryStart    func(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryStartData) context.Context
	traceQueryEnd      func(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryEndData)
	traceBatchStart    func(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchStartData) context.Context
	traceBatchQuery    func(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchQueryData)
	traceBatchEnd      func(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchEndData)
	traceCopyFromStart func(ctx context.Context, conn *pgx.Conn, data pgx.TraceCopyFromStartData) context.Context
	traceCopyFromEnd   func(ctx context.Context, conn *pgx.Conn, data pgx.TraceCopyFromEndData)
	tracePrepareStart  func(ctx context.Context, conn *pgx.Conn, data pgx.TracePrepareStartData) context.Context
	tracePrepareEnd    func(ctx context.Context, conn *pgx.Conn, data pgx.TracePrepareEndData)
	traceConnectStart  func(ctx context.Context, data pgx.TraceConnectStartData) context.Context
	traceConnectEnd    func(ctx context.Context, data pgx.TraceConnectEndData)
}

func (tt *testTracer) TraceQueryStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryStartData) context.Context {
	if tt.traceQueryStart != nil {
		return tt.traceQueryStart(ctx, conn, data)
	}
	return ctx
}

func (tt *testTracer) TraceQueryEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryEndData) {
	if tt.traceQueryEnd != nil {
		tt.traceQueryEnd(ctx, conn, data)
	}
}

func (tt *testTracer) TraceBatchStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchStartData) context.Context {
	if tt.traceBatchStart != nil {
		return tt.traceBatchStart(ctx, conn, data)
	}
	return ctx
}

func (tt *testTracer) TraceBatchQuery(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchQueryData) {
	if tt.traceBatchQuery != nil {
		tt.traceBatchQuery(ctx, conn, data)
	}
}

func (tt *testTracer) TraceBatchEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchEndData) {
	if tt.traceBatchEnd != nil {
		tt.traceBatchEnd(ctx, conn, data)
	}
}

func (tt *testTracer) TraceCopyFromStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceCopyFromStartData) context.Context {
	if tt.traceCopyFromStart != nil {
		return tt.traceCopyFromStart(ctx, conn, data)
	}
	return ctx
}

func (tt *testTracer) TraceCopyFromEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceCopyFromEndData) {
	if tt.traceCopyFromEnd != nil {
		tt.traceCopyFromEnd(ctx, conn, data)
	}
}

func (tt *testTracer) TracePrepareStart(ctx context.Context, conn *pgx.Conn, data pgx.TracePrepareStartData) context.Context {
	if tt.tracePrepareStart != nil {
		return tt.tracePrepareStart(ctx, conn, data)
	}
	return ctx
}

func (tt *testTracer) TracePrepareEnd(ctx context.Context, conn *pgx.Conn, data pgx.TracePrepareEndData) {
	if tt.tracePrepareEnd != nil {
		tt.tracePrepareEnd(ctx, conn, data)
	}
}

func (tt *testTracer) TraceConnectStart(ctx context.Context, data pgx.TraceConnectStartData) context.Context {
	if tt.traceConnectStart != nil {
		return tt.traceConnectStart(ctx, data)
	}
	return ctx
}

func (tt *testTracer) TraceConnectEnd(ctx context.Context, data pgx.TraceConnectEndData) {
	if tt.traceConnectEnd != nil {
		tt.traceConnectEnd(ctx, data)
	}
}

func main() {
	var calledStart bool
	var calledEnd bool
	cfg, err := pgx.ParseConfig(os.Getenv("DATABASE_URL"))
	cfg.Tracer = &testTracer{
		traceQueryStart: func(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryStartData) context.Context {
			calledStart = true
			return ctx
		},
		traceQueryEnd: func(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryEndData) {
			calledEnd = true
		},
	}
	conn, err := pgx.ConnectConfig(context.Background(), cfg)
	if err != nil {
		log.Fatal(err)
	}
	defer conn.Close(context.Background())

	rows, err := conn.Query(context.Background(), "select 1")
	if err != nil {
		log.Fatal(err)
	}

	rowCount := 0
	for rows.Next() {
		var d int
		if err := rows.Scan(&d); err != nil {
			log.Fatal(err)
		}

		log.Printf("result %d", d)
		rowCount++
	}

	if rows.Err() != nil {
		log.Fatal(rows.Err())
	}

	if rowCount == 0 {
		log.Fatal("no rows")
	}

	log.Printf("calledStart %v", calledStart)
	log.Printf("calledEnd %v", calledEnd)
}

Results:

jack@glados ~/dev/pgx_issues/pgx-1820 ±master⚡ » go run .
2023/12/02 09:03:13 result 1
2023/12/02 09:03:13 calledStart true
2023/12/02 09:03:13 calledEnd true
commented

Ohh, I didn't figure you had to call Next() or Close() to actually trigger the query end. Thank you very much for your quick response!