34 lines
754 B
Go
34 lines
754 B
Go
package postgres
|
|
|
|
import (
|
|
"context"
|
|
"time"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
|
|
"telesrv/internal/observability/dbtrace"
|
|
)
|
|
|
|
type queryStatsTracer struct{}
|
|
|
|
type queryStatsStartKey struct{}
|
|
|
|
func (queryStatsTracer) TraceQueryStart(ctx context.Context, _ *pgx.Conn, _ pgx.TraceQueryStartData) context.Context {
|
|
if _, ok := dbtrace.FromContext(ctx); !ok {
|
|
return ctx
|
|
}
|
|
return context.WithValue(ctx, queryStatsStartKey{}, time.Now())
|
|
}
|
|
|
|
func (queryStatsTracer) TraceQueryEnd(ctx context.Context, _ *pgx.Conn, data pgx.TraceQueryEndData) {
|
|
stats, ok := dbtrace.FromContext(ctx)
|
|
if !ok {
|
|
return
|
|
}
|
|
start, _ := ctx.Value(queryStatsStartKey{}).(time.Time)
|
|
if start.IsZero() {
|
|
stats.Add(0, data.Err)
|
|
return
|
|
}
|
|
stats.Add(time.Since(start), data.Err)
|
|
}
|