refactor(db-tracing)!: allow passing in own drivers
This commit is contained in:
parent
89ed751e04
commit
62f70ab0c9
1 changed files with 27 additions and 20 deletions
|
@ -2,13 +2,11 @@ package unitel
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"database/sql"
|
|
||||||
"database/sql/driver"
|
"database/sql/driver"
|
||||||
"fmt"
|
"fmt"
|
||||||
"strings"
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
pgx "github.com/jackc/pgx/v5/stdlib"
|
|
||||||
"github.com/qustavo/sqlhooks/v2"
|
"github.com/qustavo/sqlhooks/v2"
|
||||||
"github.com/rs/zerolog/log"
|
"github.com/rs/zerolog/log"
|
||||||
"go.opentelemetry.io/otel/attribute"
|
"go.opentelemetry.io/otel/attribute"
|
||||||
|
@ -18,13 +16,19 @@ import (
|
||||||
)
|
)
|
||||||
|
|
||||||
const (
|
const (
|
||||||
pgxClientID = libBase + "#TracedPGx"
|
sqlClientID = libBase + "#TracedSQL"
|
||||||
TracedPGxDriverName = "pgx-traced"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
type dbCtxStart struct{}
|
var dbCtxKey = struct{}{}
|
||||||
|
|
||||||
|
type dbCtxVal struct {
|
||||||
|
start time.Time
|
||||||
|
attrs []attribute.KeyValue
|
||||||
|
}
|
||||||
|
|
||||||
type tracedPgxHooks struct {
|
type tracedPgxHooks struct {
|
||||||
|
dbType string
|
||||||
|
|
||||||
t *Telemetry
|
t *Telemetry
|
||||||
printQueries bool
|
printQueries bool
|
||||||
tracer trace.Tracer
|
tracer trace.Tracer
|
||||||
|
@ -41,8 +45,9 @@ func (h *tracedPgxHooks) Before(ctx context.Context, query string, args ...inter
|
||||||
operation := strings.Split(query, " ")[0]
|
operation := strings.Split(query, " ")[0]
|
||||||
|
|
||||||
attrs := []attribute.KeyValue{
|
attrs := []attribute.KeyValue{
|
||||||
semconv.DBSystemPostgreSQL,
|
semconv.DBSystemKey.String(h.dbType),
|
||||||
semconv.DBOperationName(operation),
|
semconv.DBOperationName(operation),
|
||||||
|
|
||||||
/* Sentry */
|
/* Sentry */
|
||||||
attribute.String("db.statement", cleanedQuery),
|
attribute.String("db.statement", cleanedQuery),
|
||||||
attribute.StringSlice("db.params", formatArgs(args)),
|
attribute.StringSlice("db.params", formatArgs(args)),
|
||||||
|
@ -67,7 +72,10 @@ func (h *tracedPgxHooks) Before(ctx context.Context, query string, args ...inter
|
||||||
l.Msg(cleanedQuery)
|
l.Msg(cleanedQuery)
|
||||||
}
|
}
|
||||||
|
|
||||||
return context.WithValue(s.Context(), dbCtxStart{}, time.Now()), nil
|
return context.WithValue(s.Context(), dbCtxKey, dbCtxVal{
|
||||||
|
start: time.Now(),
|
||||||
|
attrs: attrs,
|
||||||
|
}), nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func (h *tracedPgxHooks) After(ctx context.Context, query string, args ...interface{}) (context.Context, error) {
|
func (h *tracedPgxHooks) After(ctx context.Context, query string, args ...interface{}) (context.Context, error) {
|
||||||
|
@ -75,22 +83,18 @@ func (h *tracedPgxHooks) After(ctx context.Context, query string, args ...interf
|
||||||
s.End()
|
s.End()
|
||||||
}
|
}
|
||||||
|
|
||||||
if start, ok := ctx.Value(dbCtxStart{}).(time.Time); ok {
|
if val, ok := ctx.Value(dbCtxKey).(dbCtxVal); ok {
|
||||||
h.mDuration.Record(ctx, float64(time.Since(start).Milliseconds()))
|
attrs := val.attrs
|
||||||
|
|
||||||
|
h.mDuration.Record(ctx, float64(time.Since(val.start).Milliseconds()), metric.WithAttributes(attrs...))
|
||||||
}
|
}
|
||||||
|
|
||||||
return ctx, nil
|
return ctx, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func (t *Telemetry) RegisterTracedPGx(printQueries bool) {
|
func (t *Telemetry) TraceSQL(driver driver.Driver, dbType string, printQueries bool) driver.Driver {
|
||||||
sql.Register(TracedPGxDriverName, t.TracedPGx(printQueries))
|
tracer := t.tracerProvider.Tracer(sqlClientID, trace.WithInstrumentationVersion(libVersion))
|
||||||
|
meter := t.meterProvider.Meter(sqlClientID, metric.WithInstrumentationVersion(libVersion))
|
||||||
log.Debug().Msgf("Registered %s driver", TracedPGxDriverName)
|
|
||||||
}
|
|
||||||
|
|
||||||
func (t *Telemetry) TracedPGx(printQueries bool) driver.Driver {
|
|
||||||
tracer := t.tracerProvider.Tracer(pgxClientID, trace.WithInstrumentationVersion(libVersion))
|
|
||||||
meter := t.meterProvider.Meter(pgxClientID, metric.WithInstrumentationVersion(libVersion))
|
|
||||||
|
|
||||||
mDuration, err := meter.Float64Histogram(
|
mDuration, err := meter.Float64Histogram(
|
||||||
"db.client.operation.duration",
|
"db.client.operation.duration",
|
||||||
|
@ -102,11 +106,14 @@ func (t *Telemetry) TracedPGx(printQueries bool) driver.Driver {
|
||||||
log.Fatal().Err(err).Msg("failed to create request duration histogram")
|
log.Fatal().Err(err).Msg("failed to create request duration histogram")
|
||||||
}
|
}
|
||||||
|
|
||||||
return sqlhooks.Wrap(&pgx.Driver{}, &tracedPgxHooks{
|
return sqlhooks.Wrap(driver, &tracedPgxHooks{
|
||||||
|
dbType: dbType,
|
||||||
|
|
||||||
printQueries: printQueries,
|
printQueries: printQueries,
|
||||||
t: t,
|
t: t,
|
||||||
tracer: tracer,
|
tracer: tracer,
|
||||||
mDuration: mDuration,
|
|
||||||
|
mDuration: mDuration,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in a new issue