Expose SQL query capture via govisual.WrapDriver

internal/profiling/profiler.go

@@ -384,7 +384,13 @@ func (p *Profiler) RecordSQLQuery(ctx context.Context, query string, duration ti
 	if !p.enabled.Load() {
 		return
 	}
+	RecordSQL(ctx, query, duration, rows, err)
+}
 
+// RecordSQL attaches a query to the request profile carried by ctx. Without
+// an active profile it is a no-op, so it is safe to call unconditionally —
+// this is what the WrapDriver instrumentation uses.
+func RecordSQL(ctx context.Context, query string, duration time.Duration, rows int, err error) {
 	metrics, ok := ctx.Value(profileContextKey{}).(*Metrics)
 	if !ok || metrics == nil {
 		return

internal/profiling/sqlhook.go

@@ -2,52 +2,41 @@ package profiling
 
 import (
 	"context"
-	"database/sql"
 	"database/sql/driver"
 	"time"
 )
 
-// SQLHook provides hooks for SQL operations
-type SQLHook struct {
-	profiler *Profiler
-}
-
-// NewSQLHook creates a new SQL hook
-func NewSQLHook(profiler *Profiler) *SQLHook {
-	return &SQLHook{profiler: profiler}
-}
-
-// WrapDriver wraps a SQL driver with profiling hooks
-func (h *SQLHook) WrapDriver(d driver.Driver) driver.Driver {
-	return &hookedDriver{driver: d, hook: h}
+// WrapDriver wraps a database/sql driver so queries executed with a request
+// context are recorded on that request's profile. Queries executed with
+// contexts that carry no profile are passed through untouched.
+func WrapDriver(d driver.Driver) driver.Driver {
+	return &hookedDriver{driver: d}
 }
 
 // hookedDriver wraps a driver.Driver with hooks
 type hookedDriver struct {
 	driver driver.Driver
-	hook   *SQLHook
 }
 
 func (d *hookedDriver) Open(name string) (driver.Conn, error) {
 	conn, err := d.driver.Open(name)
 	if err != nil {
 		return nil, err
 	}
-	return &hookedConn{conn: conn, hook: d.hook}, nil
+	return &hookedConn{conn: conn}, nil
 }
 
 // hookedConn wraps a driver.Conn with hooks
 type hookedConn struct {
 	conn driver.Conn
-	hook *SQLHook
 }
 
 func (c *hookedConn) Prepare(query string) (driver.Stmt, error) {
 	stmt, err := c.conn.Prepare(query)
 	if err != nil {
 		return nil, err
 	}
-	return &hookedStmt{stmt: stmt, query: query, hook: c.hook}, nil
+	return &hookedStmt{stmt: stmt, query: query}, nil
 }
 
 func (c *hookedConn) Close() error {
@@ -59,7 +48,7 @@ func (c *hookedConn) Begin() (driver.Tx, error) {
 	if err != nil {
 		return nil, err
 	}
-	return &hookedTx{tx: tx, hook: c.hook}, nil
+	return &hookedTx{tx: tx}, nil
 }
 
 // Implement other required methods
@@ -76,12 +65,12 @@ func (c *hookedConn) PrepareContext(ctx context.Context, query string) (driver.S
 	}
 
 	duration := time.Since(start)
-	c.hook.profiler.RecordSQLQuery(ctx, query, duration, 0, err)
+	RecordSQL(ctx, query, duration, 0, err)
 
 	if err != nil {
 		return nil, err
 	}
-	return &hookedStmt{stmt: stmt, query: query, hook: c.hook}, nil
+	return &hookedStmt{stmt: stmt, query: query}, nil
 }
 
 func (c *hookedConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
@@ -111,7 +100,7 @@ func (c *hookedConn) ExecContext(ctx context.Context, query string, args []drive
 		}
 	}
 
-	c.hook.profiler.RecordSQLQuery(ctx, query, duration, int(rows), err)
+	RecordSQL(ctx, query, duration, int(rows), err)
 
 	return result, err
 }
@@ -135,7 +124,7 @@ func (c *hookedConn) QueryContext(ctx context.Context, query string, args []driv
 	}
 
 	duration := time.Since(start)
-	c.hook.profiler.RecordSQLQuery(ctx, query, duration, 0, err)
+	RecordSQL(ctx, query, duration, 0, err)
 
 	return rows, err
 }
@@ -153,14 +142,13 @@ func (c *hookedConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver
 	if err != nil {
 		return nil, err
 	}
-	return &hookedTx{tx: tx, hook: c.hook}, nil
+	return &hookedTx{tx: tx}, nil
 }
 
 // hookedStmt wraps a driver.Stmt with hooks
 type hookedStmt struct {
 	stmt  driver.Stmt
 	query string
-	hook  *SQLHook
 }
 
 func (s *hookedStmt) Close() error {
@@ -184,7 +172,7 @@ func (s *hookedStmt) Exec(args []driver.Value) (driver.Result, error) {
 	}
 
 	// Use background context as we don't have access to request context here
-	s.hook.profiler.RecordSQLQuery(context.Background(), s.query, duration, int(rows), err)
+	RecordSQL(context.Background(), s.query, duration, int(rows), err)
 
 	return result, err
 }
@@ -194,7 +182,7 @@ func (s *hookedStmt) Query(args []driver.Value) (driver.Rows, error) {
 	rows, err := s.stmt.Query(args)
 	duration := time.Since(start)
 
-	s.hook.profiler.RecordSQLQuery(context.Background(), s.query, duration, 0, err)
+	RecordSQL(context.Background(), s.query, duration, 0, err)
 
 	return rows, err
 }
@@ -224,7 +212,7 @@ func (s *hookedStmt) ExecContext(ctx context.Context, args []driver.NamedValue)
 		}
 	}
 
-	s.hook.profiler.RecordSQLQuery(ctx, s.query, duration, int(rows), err)
+	RecordSQL(ctx, s.query, duration, int(rows), err)
 
 	return result, err
 }
@@ -246,15 +234,14 @@ func (s *hookedStmt) QueryContext(ctx context.Context, args []driver.NamedValue)
 	}
 
 	duration := time.Since(start)
-	s.hook.profiler.RecordSQLQuery(ctx, s.query, duration, 0, err)
+	RecordSQL(ctx, s.query, duration, 0, err)
 
 	return rows, err
 }
 
 // hookedTx wraps a driver.Tx with hooks
 type hookedTx struct {
-	tx   driver.Tx
-	hook *SQLHook
+	tx driver.Tx
 }
 
 func (t *hookedTx) Commit() error {
@@ -265,8 +252,3 @@ func (t *hookedTx) Rollback() error {
 	return t.tx.Rollback()
 }
 
-// RegisterSQLDriver registers a wrapped SQL driver with profiling support
-func RegisterSQLDriver(name string, driver driver.Driver, profiler *Profiler) {
-	hook := NewSQLHook(profiler)
-	sql.Register(name+"-profiled", hook.WrapDriver(driver))
-}

internal/profiling/sqlhook_test.go

@@ -0,0 +1,69 @@
+package profiling
+
+import (
+	"context"
+	"database/sql"
+	"database/sql/driver"
+	"testing"
+	"time"
+)
+
+type fakeDriver struct{}
+
+func (fakeDriver) Open(name string) (driver.Conn, error) { return &fakeConn{}, nil }
+
+type fakeConn struct{}
+
+func (*fakeConn) Prepare(query string) (driver.Stmt, error) { return nil, driver.ErrSkip }
+func (*fakeConn) Close() error                              { return nil }
+func (*fakeConn) Begin() (driver.Tx, error)                 { return nil, driver.ErrSkip }
+
+func (*fakeConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
+	return driver.RowsAffected(3), nil
+}
+
+func TestWrapDriverRecordsQueriesOnProfile(t *testing.T) {
+	sql.Register("fake-viz", WrapDriver(fakeDriver{}))
+	db, err := sql.Open("fake-viz", "")
+	if err != nil {
+		t.Fatalf("open: %v", err)
+	}
+	defer db.Close()
+
+	profiler := NewProfiler(10)
+	profiler.SetProfileType(ProfileMemory)
+	profiler.SetThreshold(0)
+	ctx := profiler.StartProfiling(context.Background(), "req-sql")
+
+	if _, err := db.ExecContext(ctx, "UPDATE things SET x = 1"); err != nil {
+		t.Fatalf("exec: %v", err)
+	}
+
+	metrics := profiler.EndProfiling(ctx)
+	if metrics == nil {
+		t.Fatal("expected metrics")
+	}
+	if len(metrics.SQLQueries) != 1 {
+		t.Fatalf("expected 1 recorded query, got %d", len(metrics.SQLQueries))
+	}
+	q := metrics.SQLQueries[0]
+	if q.Query != "UPDATE things SET x = 1" || q.Rows != 3 {
+		t.Fatalf("recorded query = %+v", q)
+	}
+}
+
+func TestWrapDriverNoProfileIsNoOp(t *testing.T) {
+	sql.Register("fake-viz-2", WrapDriver(fakeDriver{}))
+	db, err := sql.Open("fake-viz-2", "")
+	if err != nil {
+		t.Fatalf("open: %v", err)
+	}
+	defer db.Close()
+
+	// A context without an active profile must pass through untouched.
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	if _, err := db.ExecContext(ctx, "SELECT 1"); err != nil {
+		t.Fatalf("exec: %v", err)
+	}
+}

sql.go

@@ -0,0 +1,21 @@
+package govisual
+
+import (
+	"database/sql/driver"
+
+	"github.com/doganarif/govisual/v2/internal/profiling"
+)
+
+// WrapDriver instruments a database/sql driver so queries executed with a
+// request's context show up on that request's profile in the dashboard.
+// Register the wrapped driver once and open the database through it:
+//
+//	sql.Register("postgres+viz", govisual.WrapDriver(&pq.Driver{}))
+//	db, err := sql.Open("postgres+viz", dsn)
+//
+// Queries are attributed through the context, so they must run through the
+// *Context variants (QueryContext, ExecContext) with the incoming request's
+// context, and profiling must be enabled via WithProfiling(true).
+func WrapDriver(d driver.Driver) driver.Driver {
+	return profiling.WrapDriver(d)
+}