package store

import (
	"context"
	"time"

	"github.com/jackc/pgx/v5"

	"github.com/ponzischeme89/memby/server/internal/timing"
)

// queryTracer attributes Postgres time to the request that asked for it.
//
// pgx's tracer interface is the only place every query in the gateway passes through —
// there are several hundred call sites across internal/store and instrumenting them by
// hand would guarantee the one that matters is the one nobody wrapped. The SQL itself
// is deliberately not recorded: a breakdown is read beside a request line in the admin
// console, and a query text there would put table and column names in front of anybody
// who can open the log.
type queryTracer struct{}

type queryStartKey struct{}

func (queryTracer) TraceQueryStart(
	ctx context.Context, _ *pgx.Conn, _ pgx.TraceQueryStartData,
) context.Context {
	if timing.From(ctx) == nil {
		return ctx
	}
	return context.WithValue(ctx, queryStartKey{}, time.Now())
}

func (queryTracer) TraceQueryEnd(ctx context.Context, _ *pgx.Conn, _ pgx.TraceQueryEndData) {
	began, ok := ctx.Value(queryStartKey{}).(time.Time)
	if !ok {
		return
	}
	timing.Record(ctx, timing.StageDB, time.Since(began))
}