Add ConnWrapper

connector.go → db.go

@@ -1,6 +1,7 @@
 package qg
 
 import (
+	"database/sql"
 	"database/sql/driver"
 	"fmt"
 	"net/url"
@@ -35,3 +36,26 @@ func GetConnectorFromConnStr(connStr string) (driver.Connector, error) {
 
 	return connector, nil
 }
+
+// driver.Conn Wrapper
+type ConnWrapper interface {
+	WrappedConn() driver.Conn
+}
+
+// Get *pgx.Conn from *sql.Conn and pass it to function.
+func rawConn(conn *sql.Conn, f func(*pgx.Conn) error) error {
+	err := conn.Raw(func(driverConn any) error {
+		var stdlibConn *stdlib.Conn
+
+		if tracedConn, ok := driverConn.(ConnWrapper); ok {
+			stdlibConn = tracedConn.WrappedConn().(*stdlib.Conn)
+		} else {
+			stdlibConn = driverConn.(*stdlib.Conn)
+		}
+
+		pgxConn := stdlibConn.Conn()
+		return f(pgxConn)
+	})
+
+	return err
+}

util.go

@@ -1,13 +1,5 @@
 package qg
 
-import (
-	"database/sql"
-
-	"github.com/jackc/pgx/v4"
-	"github.com/jackc/pgx/v4/stdlib"
-	sqltracer "gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql"
-)
-
 // intPow returns x**y, the base-x exponential of y.
 func intPow(x, y int) (r int) {
 	if x == r || y < r {
@@ -32,21 +24,3 @@ func intPow(x, y int) (r int) {
 	}
 	return
 }
-
-// Get *pgx.Conn from *sql.Conn and pass it to function.
-func rawConn(conn *sql.Conn, f func(*pgx.Conn) error) error {
-	err := conn.Raw(func(driverConn any) error {
-		var stdlibConn *stdlib.Conn
-
-		if tracedConn, ok := driverConn.(*sqltracer.TracedConn); ok {
-			stdlibConn = tracedConn.WrappedConn().(*stdlib.Conn)
-		} else {
-			stdlibConn = driverConn.(*stdlib.Conn)
-		}
-
-		pgxConn := stdlibConn.Conn()
-		return f(pgxConn)
-	})
-
-	return err
-}