feat(interceptor): add PrepareContext support (#15)

Author: tmzane

README.md

@@ -94,17 +94,14 @@ db, _ := sql.Open("interceptor", "dsn")
 
 db.ExecContext(ctx, "INSERT INTO users VALUES (1, 'John Doe')")
 // stderr: INFO ExecContext query="INSERT INTO users VALUES (1, 'John Doe')"
+
 db.QueryContext(ctx, "SELECT id, name FROM users")
 // stderr: INFO QueryContext query="SELECT id, name FROM users"
 ```
 
-> [!note]
-> To keep the implementation simple, only `ExecContext` and `QueryContext` callbacks are supported.
-> If you need to intercept other database operations, such as `sql.DB.BeginTx`, consider using [ngrok/sqlmw][3] instead.
-
 Integration tests cover the following databases and drivers:
-- PostgreSQL with [jackx/pgx][4]
-- MySQL with [go-sql-driver/mysql][5]
+- PostgreSQL with [jackx/pgx][3]
+- MySQL with [go-sql-driver/mysql][4]
 
 ## 🚧 TODOs
 
@@ -114,6 +111,5 @@ Integration tests cover the following databases and drivers:
 
 [1]: https://github.com/golang/go/issues/61637
 [2]: https://grpc.io/docs/guides/interceptors
-[3]: https://github.com/ngrok/sqlmw
-[4]: https://github.com/jackc/pgx
-[5]: https://github.com/go-sql-driver/mysql
+[3]: https://github.com/jackc/pgx
+[4]: https://github.com/go-sql-driver/mysql

interceptor.go

@@ -10,6 +10,8 @@ var (
 	_ driver.DriverContext = Interceptor{}
 )
 
+// TODO: document that database/sql falls back to Prepare if the driver returns ErrSkip for Exec/Query.
+
 // Interceptor is a [driver.Driver] wrapper that allows to register callbacks for database queries.
 // It must first be registered with [sql.Register] with the same name that is then passed to [sql.Open]:
 //
@@ -31,6 +33,11 @@ type Interceptor struct {
 	// The implementation must call queryer.QueryContext(ctx, query, args) and return the result.
 	// Optional.
 	QueryContext func(ctx context.Context, query string, args []driver.NamedValue, queryer driver.QueryerContext) (driver.Rows, error)
+
+	// PrepareContext is a callback for [sql.DB.PrepareContext].
+	// The implementation must call preparer.ConnPrepareContext(ctx, query) and return the result.
+	// Optional.
+	PrepareContext func(ctx context.Context, query string, preparer driver.ConnPrepareContext) (driver.Stmt, error)
 }
 
 // Open implements [driver.Driver].
@@ -54,9 +61,10 @@ func (i Interceptor) OpenConnector(name string) (driver.Connector, error) {
 }
 
 var (
-	_ driver.Conn           = wrappedConn{}
-	_ driver.ExecerContext  = wrappedConn{}
-	_ driver.QueryerContext = wrappedConn{}
+	_ driver.Conn               = wrappedConn{}
+	_ driver.ExecerContext      = wrappedConn{}
+	_ driver.QueryerContext     = wrappedConn{}
+	_ driver.ConnPrepareContext = wrappedConn{}
 )
 
 type wrappedConn struct {
@@ -90,6 +98,18 @@ func (c wrappedConn) QueryContext(ctx context.Context, query string, args []driv
 
 var _ driver.Connector = wrappedConnector{}
 
+// PrepareContext implements [driver.ConnPrepareContext].
+func (c wrappedConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
+	preparer, ok := c.Conn.(driver.ConnPrepareContext)
+	if !ok {
+		panic("queries: driver does not implement driver.ConnPrepareContext")
+	}
+	if c.interceptor.PrepareContext != nil {
+		return c.interceptor.PrepareContext(ctx, query, preparer)
+	}
+	return preparer.PrepareContext(ctx, query)
+}
+
 type wrappedConnector struct {
 	driver.Connector
 	interceptor Interceptor

interceptor_test.go

@@ -15,19 +15,24 @@ import (
 func TestInterceptor(t *testing.T) {
 	ctx := context.Background()
 
-	var execIntercepted bool
-	var queryIntercepted bool
+	var execCalled bool
+	var queryCalled bool
+	var prepareCalled bool
 
 	interceptor := queries.Interceptor{
 		Driver: mockDriver{conn: spyConn{}},
 		ExecContext: func(ctx context.Context, query string, args []driver.NamedValue, execer driver.ExecerContext) (driver.Result, error) {
-			execIntercepted = true
+			execCalled = true
 			return execer.ExecContext(ctx, query, args)
 		},
 		QueryContext: func(ctx context.Context, query string, args []driver.NamedValue, queryer driver.QueryerContext) (driver.Rows, error) {
-			queryIntercepted = true
+			queryCalled = true
 			return queryer.QueryContext(ctx, query, args)
 		},
+		PrepareContext: func(ctx context.Context, query string, preparer driver.ConnPrepareContext) (driver.Stmt, error) {
+			prepareCalled = true
+			return preparer.PrepareContext(ctx, query)
+		},
 	}
 
 	driverName := t.Name() + "_interceptor"
@@ -39,11 +44,15 @@ func TestInterceptor(t *testing.T) {
 
 	_, err = db.ExecContext(ctx, "")
 	assert.IsErr[E](t, err, errCalled)
-	assert.Equal[E](t, execIntercepted, true)
+	assert.Equal[E](t, execCalled, true)
 
 	_, err = db.QueryContext(ctx, "") //nolint:gocritic // sqlQuery: unused result is fine here.
 	assert.IsErr[E](t, err, errCalled)
-	assert.Equal[E](t, queryIntercepted, true)
+	assert.Equal[E](t, queryCalled, true)
+
+	_, err = db.PrepareContext(ctx, "")
+	assert.IsErr[E](t, err, errCalled)
+	assert.Equal[E](t, prepareCalled, true)
 }
 
 func TestInterceptor_passthrough(t *testing.T) {
@@ -65,6 +74,9 @@ func TestInterceptor_passthrough(t *testing.T) {
 
 	_, err = db.QueryContext(ctx, "") //nolint:gocritic // sqlQuery: unused result is fine here.
 	assert.IsErr[E](t, err, errCalled)
+
+	_, err = db.PrepareContext(ctx, "")
+	assert.IsErr[E](t, err, errCalled)
 }
 
 func TestInterceptor_unimplemented(t *testing.T) {
@@ -86,6 +98,9 @@ func TestInterceptor_unimplemented(t *testing.T) {
 
 	queryFn := func() { _, _ = db.QueryContext(ctx, "") } //nolint:gocritic // sqlQuery: unused result is fine here.
 	assert.Panics[E](t, queryFn, "queries: driver does not implement driver.QueryerContext")
+
+	prepareFn := func() { _, _ = db.PrepareContext(ctx, "") }
+	assert.Panics[E](t, prepareFn, "queries: driver does not implement driver.ConnPrepareContext")
 }
 
 func TestInterceptor_driver(t *testing.T) {
@@ -123,3 +138,7 @@ func (spyConn) ExecContext(context.Context, string, []driver.NamedValue) (driver
 func (spyConn) QueryContext(context.Context, string, []driver.NamedValue) (driver.Rows, error) {
 	return nil, errCalled
 }
+
+func (spyConn) PrepareContext(context.Context, string) (driver.Stmt, error) {
+	return nil, errCalled
+}

tests/integration_test.go

@@ -40,6 +40,7 @@ func TestIntegration(t *testing.T) {
 	for name, database := range DBs {
 		var execCalls int
 		var queryCalls int
+		var prepareCalls int
 
 		interceptor := queries.Interceptor{
 			Driver: database.driver,
@@ -53,6 +54,11 @@ func TestIntegration(t *testing.T) {
 				t.Logf("[%s] QueryContext: %s %v", name, query, namedToAny(args))
 				return queryer.QueryContext(ctx, query, args)
 			},
+			PrepareContext: func(ctx context.Context, query string, preparer driver.ConnPrepareContext) (driver.Stmt, error) {
+				prepareCalls++
+				t.Logf("[%s] PrepareContext: %s", name, query)
+				return preparer.PrepareContext(ctx, query)
+			},
 		}
 
 		driverName := name + "_interceptor"
@@ -111,6 +117,12 @@ func TestIntegration(t *testing.T) {
 		assert.NoErr[F](t, tx.Commit())
 		assert.Equal[E](t, execCalls, 2)
 		assert.Equal[E](t, queryCalls, 5*2)
+		if name == "mysql" {
+			// github.com/go-sql-driver/mysql falls back to PrepareContext for queries with arguments.
+			assert.Equal[E](t, prepareCalls, 1)
+		} else {
+			assert.Equal[E](t, prepareCalls, 0)
+		}
 	}
 }