Tracing and Testing - Githubissues

I am creating a new span before calling generated code and would like the pggen layer to be observable. How would you recommend doing this?

Observability is lackluster pgx/v4, but much improved in pgx/v5. Since pggen currently targets pgx/v4, we created our own pool, implementing pgxpool.Pool, to read the pggen query name from context and emit spans. The main difficulty is that the Pool methods returns pgx.Rows which delay work until scanned.

Here's one method from our tracing pool implementation.

// TracingPool is a wrapper around pgxpool.Pool that just records traces.
type TracingPool struct {
    *pgxpool.Pool
}

func NewTracingPool(pool *pgxpool.Pool) *TracingPool {
    return &TracingPool{Pool: pool}
}

func (t *TracingPool) Query(ctx context.Context, sql string, args ...interface{}) (pgx.Rows, error) {
    waiter := newAfterReleaseWaiter()
    ctx, span := trace.OpenSpan(waiter.embedCtx(ctx), "DB Pool Query - "+ExtractPggenQueryName(ctx))
    // Span closed in tracingRows.Close.
    rows, err := t.Pool.Query(ctx, annotateQuery(ctx, sql), args...)
    if err != nil {
        return rows, err
    }
    return newTracingRows(rows, ctx, span, waiter), nil
}

The painful part is tracking the span by implementing pgx.Rows.

If I were to implement this today, I wouldn't try to track the span through Rows.scan. I'd instrument each top-level Pool method with BeginSpan/defer EndSpan and maybe use the pool lifecycle hooks to emit events:

Here's how we did the pool lifecycle hooks.

// AddPoolConfigTracing wraps existing hooks in a pgxpool.Config to add trace
// events for pool lifecycle events.
func AddPoolConfigTracing(cfg *pgxpool.Config) {
    oldBeforeConnect := cfg.BeforeConnect
    oldAfterConnect := cfg.AfterConnect
    oldBeforeAcquire := cfg.BeforeAcquire
    oldAfterRelease := cfg.AfterRelease

    oldLogger := cfg.ConnConfig.Logger
    cfg.ConnConfig.Logger = connEventLogger{next: oldLogger}
    cfg.ConnConfig.LogLevel = pgx.LogLevelInfo

    // Use a pointer to conn to relate BeforeAcquire and AfterRelease since
    // AfterRelease doesn't take a context. Using sync.Map since each go routine
    // is only modifying a single map entry, which sync.Map is specialized for.
    connCtxs := &sync.Map{}

    cfg.BeforeConnect = func(ctx context.Context, config *pgx.ConnConfig) error {
        if oldBeforeConnect != nil {
            if err := oldBeforeConnect(ctx, config); err != nil {
                return err
            }
        }
        trace.SpanFromContext(ctx).AddEvent(trace.NewEvent("PgPool: BeforeConnect"))
        return nil
    }

    cfg.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
        if oldAfterConnect != nil {
            if err := oldAfterConnect(ctx, conn); err != nil {
                return err
            }
        }
        trace.SpanFromContext(ctx).AddEvent(trace.NewEvent("PgPool: AfterConnect"))
        return nil
    }

    cfg.BeforeAcquire = func(ctx context.Context, conn *pgx.Conn) bool {
        if oldBeforeAcquire != nil {
            if shouldContinue := oldBeforeAcquire(ctx, conn); !shouldContinue {
                return shouldContinue
            }
        }
        trace.SpanFromContext(ctx).AddEvent(trace.NewEvent("PgPool: BeforeAcquire"))
        connCtxs.Store(conn, ctx)
        return true
    }

    cfg.AfterRelease = func(conn *pgx.Conn) bool {
        if oldAfterRelease != nil {
            if shouldContinue := oldAfterRelease(conn); !shouldContinue {
                return shouldContinue
            }
        }
        ctxVal, ok := connCtxs.LoadAndDelete(conn)
        if !ok {
            return true
        }
        ctx := ctxVal.(context.Context)
        trace.SpanFromContext(ctx).AddEvent(trace.NewEvent("PgPool: AfterRelease"))
        if w := extractAfterReleaseWaiterContext(ctx); w != nil {
            w.close()
        }

        return true
    }
}

Here's what it looks like:

Additionally, can you please explain more about how you perform testing of the generated code?

We use a real database for all tests. I avoid mocking databases since database interactions are usually the most important part of a backend server and should be covered by tests.

Each test gets a fresh copy of the database. In essence, each test does the following:

initdb
Run all migrations
Run the test

We employ many tricks to keep testing fast:

Cache the data directory with Bazel so we only run migrations once for all tests.
Run Postgres with unsafe values, like fsync=off: https://stackoverflow.com/questions/9407442/optimise-postgresql-for-fast-testing
In each test file, run the common database setup code and snapshot the resulting data directory. Then, for each test, copy the snapshotted directory.
Use CloneFile when available for cheap directory copies: https://www.manpagez.com/man/2/clonefile/.

We started much simpler. You can get most of the speedups by running migrations once and caching the directory.

For more Bazel-specific info, see https://news.ycombinator.com/item?id=28082628 and https://news.ycombinator.com/item?id=28470034.

jschaf / pggen

Tracing and Testing #91