feat: add QueryTracer interface for SQL statement tracing (#1716)

connection.go

@@ -555,11 +555,14 @@
 		return nil, err
 	}
 
+	ctx, traceEnd := mc.traceQuery(ctx, query, args)
 	rows, err := mc.query(query, dargs)
 	if err != nil {
+		traceEnd(err)
 		mc.finish()
 		return nil, err
 	}
+	traceEnd(nil)
 	rows.finish = mc.finish
 	return rows, err
 }
@@ -575,7 +578,10 @@
 	}
 	defer mc.finish()
 
-	return mc.Exec(query, dargs)
+	_, traceEnd := mc.traceQuery(ctx, query, args)
+	result, err := mc.Exec(query, dargs)
+	traceEnd(err)
+	return result, err
 }
 
 func (mc *mysqlConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
@@ -595,6 +601,11 @@
 		stmt.Close()
 		return nil, ctx.Err()
 	}
+
+	// Store query string for tracing prepared statement execution.
+	if s, ok := stmt.(*mysqlStmt); ok {
+	s.queryStr = query
+	}
 	return stmt, nil
 }
 
@@ -608,11 +619,14 @@
 		return nil, err
 	}
 
+	ctx, traceEnd := stmt.mc.traceQuery(ctx, stmt.queryStr, args)
 	rows, err := stmt.query(dargs)
 	if err != nil {
+		traceEnd(err)
 		stmt.mc.finish()
 		return nil, err
 	}
+	traceEnd(nil)
 	rows.finish = stmt.mc.finish
 	return rows, err
 }
@@ -628,7 +642,10 @@
 	}
 	defer stmt.mc.finish()
 
-	return stmt.Exec(dargs)
+	_, traceEnd := stmt.mc.traceQuery(ctx, stmt.queryStr, args)
+	result, err := stmt.Exec(dargs)
+	traceEnd(err)
+	return result, err
 }
 
 func (mc *mysqlConn) watchCancel(ctx context.Context) error {

dsn.go

@@ -81,6 +81,7 @@ type Config struct {
 	pubKey        *rsa.PublicKey                       // Server public key
 	timeTruncate  time.Duration                        // Truncate time.Time values to the specified duration
 	charsets      []string                             // Connection charset. When set, this will be set in SET NAMES <charset> query
+	tracer        QueryTracer                          // Tracer for SQL query tracing
 }
 
 // Functional Options Pattern
@@ -135,6 +136,16 @@ func EnableCompression(yes bool) Option {
 	}
 }
 
+// WithTracer sets the query tracer for tracing SQL query execution.
+// The tracer is called before and after each query with the query string,
+// arguments, error, and execution duration.
+func WithTracer(tracer QueryTracer) Option {
+	return func(cfg *Config) error {
+		cfg.tracer = tracer
+		return nil
+	}
+}
+
 // Charset sets the connection charset and collation.
 //
 // charset is the connection charset.

statement.go

@@ -21,6 +21,7 @@ type mysqlStmt struct {
 	id         uint32
 	paramCount int
 	columns    []mysqlField
+	queryStr   string // original query string, stored for tracing
 }
 
 func (stmt *mysqlStmt) Close() error {

tracer.go

@@ -0,0 +1,39 @@
+package mysql
+
+import (
+	"context"
+	"database/sql/driver"
+	"time"
+)
+
+// QueryTracer is an interface for tracing SQL query execution.
+// It can be used for logging, metrics collection, or distributed tracing.
+//
+// TraceQueryStart is called before a query is executed. It receives the context,
+// the SQL query string, and the named arguments. It returns a new context that
+// will be passed to TraceQueryEnd — this allows attaching trace-specific metadata
+// (e.g. span IDs) to the context.
+//
+// TraceQueryEnd is called after the query completes (or fails). It receives the
+// context returned by TraceQueryStart, the error (nil on success), and the
+// wall-clock duration of the query execution.
+type QueryTracer interface {
+	TraceQueryStart(ctx context.Context, query string, args []driver.NamedValue) context.Context
+	TraceQueryEnd(ctx context.Context, err error, duration time.Duration)
+}
+
+// traceQuery starts tracing a query if a tracer is configured.
+// It returns the (possibly updated) context and a finish function.
+// The finish function must be called with the resulting error when the query completes.
+// If no tracer is configured, the returned context is unchanged and the finish function is a no-op.
+func (mc *mysqlConn) traceQuery(ctx context.Context, query string, args []driver.NamedValue) (context.Context, func(error)) {
+	t := mc.cfg.tracer
+	if t == nil {
+		return ctx, func(error) {}
+	}
+	start := time.Now()
+	ctx = t.TraceQueryStart(ctx, query, args)
+	return ctx, func(err error) {
+		t.TraceQueryEnd(ctx, err, time.Since(start))
+	}
+}

tracer_test.go

@@ -0,0 +1,165 @@
+package mysql
+
+import (
+	"context"
+	"database/sql/driver"
+	"testing"
+	"time"
+)
+
+// testTracer records trace calls for verification.
+type testTracer struct {
+	startCalled bool
+	endCalled   bool
+	query       string
+	args        []driver.NamedValue
+	err         error
+	duration    time.Duration
+	ctxKey      any
+	ctxVal      any
+}
+
+type tracerCtxKey struct{}
+
+func (t *testTracer) TraceQueryStart(ctx context.Context, query string, args []driver.NamedValue) context.Context {
+	t.startCalled = true
+	t.query = query
+	t.args = args
+	// Attach a value to context to verify it flows to TraceQueryEnd.
+	return context.WithValue(ctx, tracerCtxKey{}, "traced")
+}
+
+func (t *testTracer) TraceQueryEnd(ctx context.Context, err error, duration time.Duration) {
+	t.endCalled = true
+	t.err = err
+	t.duration = duration
+	t.ctxVal = ctx.Value(tracerCtxKey{})
+}
+
+func (t *testTracer) reset() {
+	t.startCalled = false
+	t.endCalled = false
+	t.query = ""
+	t.args = nil
+	t.err = nil
+	t.duration = 0
+	t.ctxVal = nil
+}
+
+func TestTraceQuery_WithTracer(t *testing.T) {
+	tr := &testTracer{}
+	mc := &mysqlConn{
+		cfg: &Config{
+			tracer: tr,
+		},
+	}
+
+	args := []driver.NamedValue{
+		{Ordinal: 1, Value: int64(42)},
+		{Ordinal: 2, Value: "hello"},
+	}
+
+	ctx, finish := mc.traceQuery(context.Background(), "SELECT * FROM users WHERE id = ?", args)
+	_ = ctx
+
+	if !tr.startCalled {
+		t.Fatal("TraceQueryStart was not called")
+	}
+	if tr.query != "SELECT * FROM users WHERE id = ?" {
+		t.Fatalf("unexpected query: %q", tr.query)
+	}
+	if len(tr.args) != 2 {
+		t.Fatalf("expected 2 args, got %d", len(tr.args))
+	}
+	if tr.args[0].Value != int64(42) {
+		t.Fatalf("unexpected arg[0]: %v", tr.args[0].Value)
+	}
+
+	// Simulate some work
+	time.Sleep(time.Millisecond)
+	finish(nil)
+
+	if !tr.endCalled {
+		t.Fatal("TraceQueryEnd was not called")
+	}
+	if tr.err != nil {
+		t.Fatalf("unexpected error: %v", tr.err)
+	}
+	if tr.duration < time.Millisecond {
+		t.Fatalf("duration too short: %v", tr.duration)
+	}
+}
+
+func TestTraceQuery_ContextFlows(t *testing.T) {
+	tr := &testTracer{}
+	mc := &mysqlConn{
+		cfg: &Config{
+			tracer: tr,
+		},
+	}
+
+	_, finish := mc.traceQuery(context.Background(), "INSERT INTO t VALUES (?)", nil)
+	finish(nil)
+
+	// The context value set in TraceQueryStart should be visible in TraceQueryEnd.
+	if tr.ctxVal != "traced" {
+		t.Fatalf("context value not propagated: got %v, want %q", tr.ctxVal, "traced")
+	}
+}
+
+func TestTraceQuery_WithError(t *testing.T) {
+	tr := &testTracer{}
+	mc := &mysqlConn{
+		cfg: &Config{
+			tracer: tr,
+		},
+	}
+
+	_, finish := mc.traceQuery(context.Background(), "BAD SQL", nil)
+	finish(ErrInvalidConn)
+
+	if !tr.endCalled {
+		t.Fatal("TraceQueryEnd was not called")
+	}
+	if tr.err != ErrInvalidConn {
+		t.Fatalf("unexpected error: %v, want %v", tr.err, ErrInvalidConn)
+	}
+}
+
+func TestTraceQuery_NilTracer(t *testing.T) {
+	mc := &mysqlConn{
+		cfg: &Config{
+			tracer: nil,
+		},
+	}
+
+	ctx := context.Background()
+	retCtx, finish := mc.traceQuery(ctx, "SELECT 1", nil)
+
+	// Context should be unchanged.
+	if retCtx != ctx {
+		t.Fatal("context should not be modified when tracer is nil")
+	}
+
+	// finish should be safe to call (no-op).
+	finish(nil)
+	finish(ErrInvalidConn)
+}
+
+func TestWithTracerOption(t *testing.T) {
+	tr := &testTracer{}
+	cfg := NewConfig()
+
+	if cfg.tracer != nil {
+		t.Fatal("tracer should be nil by default")
+	}
+
+	err := cfg.Apply(WithTracer(tr))
+	if err != nil {
+		t.Fatalf("Apply(WithTracer) failed: %v", err)
+	}
+
+	if cfg.tracer != tr {
+		t.Fatal("tracer was not set")
+	}
+}