From bf66baf3c5bd43184a85648212aa1b0cfc3c265d Mon Sep 17 00:00:00 2001
From: DevMiner <devminer@devminer.xyz>
Date: Tue, 23 Jul 2024 17:55:58 +0200
Subject: [PATCH] chore(pgx): make it possible to get the raw hooked driver

---
 pgx_tracing.go | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

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 {