diff --git a/pgx_tracing.go b/pgx_tracing.go index 06b8b82..353739f 100644 --- a/pgx_tracing.go +++ b/pgx_tracing.go @@ -3,6 +3,7 @@ package unitel import ( "context" "database/sql" + "database/sql/driver" "fmt" "strings" @@ -59,15 +60,19 @@ func (h *tracedPgxHooks) After(ctx context.Context, query string, args ...interf } func (t *Telemetry) RegisterTracedPGx(printQueries bool) { + sql.Register(TracedPGxDriverName, t.TracedPGx(printQueries)) + + log.Debug().Msgf("Registered %s driver", TracedPGxDriverName) +} + +func (t *Telemetry) TracedPGx(printQueries bool) driver.Driver { tracer := t.tracerProvider.Tracer(pgxClientID, trace.WithInstrumentationVersion(libVersion)) - sql.Register(TracedPGxDriverName, sqlhooks.Wrap(&pgx.Driver{}, &tracedPgxHooks{ + return sqlhooks.Wrap(&pgx.Driver{}, &tracedPgxHooks{ printQueries: printQueries, t: t, tracer: tracer, - })) - - log.Debug().Msgf("Registered %s driver", TracedPGxDriverName) + }) } func formatArgs(args []interface{}) []string {