feat: add DirectExecuteQuery option (#455)

Adds an option to directly execute the query when QueryContext is called. Without this
option, the query execution is delayed until the first call to Rows.Next. This again
also delays any query errors until the first call to Rows.Next, which can be confusing
in some cases.

Author: olavloite

conn.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"database/sql"
 	"database/sql/driver"
+	"errors"
 	"log/slog"
 	"slices"
 	"time"
@@ -797,7 +798,16 @@ func (c *conn) queryContext(ctx context.Context, query string, execOptions ExecO
 			return nil, err
 		}
 	}
-	return &rows{it: iter, decodeOption: execOptions.DecodeOption, decodeToNativeArrays: execOptions.DecodeToNativeArrays}, nil
+	res := &rows{it: iter, decodeOption: execOptions.DecodeOption, decodeToNativeArrays: execOptions.DecodeToNativeArrays}
+	if execOptions.DirectExecuteQuery {
+		// This call to res.getColumns() triggers the execution of the statement, as it needs to fetch the metadata.
+		res.getColumns()
+		if res.dirtyErr != nil && !errors.Is(res.dirtyErr, iterator.Done) {
+			_ = res.Close()
+			return nil, res.dirtyErr
+		}
+	}
+	return res, nil
 }
 
 func (c *conn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {

driver.go

@@ -175,6 +175,13 @@ type ExecOptions struct {
 	// AutoCommitDMLMode determines the type of transaction that DML statements
 	// that are executed outside explicit transactions use.
 	AutocommitDMLMode AutocommitDMLMode
+
+	// DirectExecute determines whether a query is executed directly when the
+	// [sql.DB.QueryContext] method is called, or whether the actual query execution
+	// is delayed until the first call to [sql.Rows.Next]. The default is to delay
+	// the execution. Set this flag to true to execute the query directly when
+	// [sql.DB.QueryContext] is called.
+	DirectExecuteQuery bool
 }
 
 type DecodeOption int

driver_with_mockserver_test.go

@@ -235,6 +235,53 @@ func TestSimpleQuery(t *testing.T) {
 	}
 }
 
+func TestDirectExecuteQuery(t *testing.T) {
+	t.Parallel()
+
+	db, server, teardown := setupTestDBConnection(t)
+	defer teardown()
+
+	// This does not use DirectExecuteQuery. The query is only sent to Spanner when
+	// rows.Next is called.
+	rows, err := db.QueryContext(context.Background(), testutil.SelectFooFromBar)
+	if err != nil {
+		t.Fatal(err)
+	}
+	// There should be no request on the server.
+	requests := drainRequestsFromServer(server.TestSpanner)
+	sqlRequests := requestsOfType(requests, reflect.TypeOf(&sppb.ExecuteSqlRequest{}))
+	if g, w := len(sqlRequests), 0; g != w {
+		t.Fatalf("sql requests count mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	if !rows.Next() {
+		t.Fatal("no rows")
+	}
+	// The request should now be present on the server.
+	requests = drainRequestsFromServer(server.TestSpanner)
+	sqlRequests = requestsOfType(requests, reflect.TypeOf(&sppb.ExecuteSqlRequest{}))
+	if g, w := len(sqlRequests), 1; g != w {
+		t.Fatalf("sql requests count mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	_ = rows.Close()
+
+	// Now repeat the same with the DirectExecuteQuery option.
+	rows, err = db.QueryContext(context.Background(), testutil.SelectFooFromBar, ExecOptions{DirectExecuteQuery: true})
+	if err != nil {
+		t.Fatal(err)
+	}
+	// The request should be present on the server.
+	requests = drainRequestsFromServer(server.TestSpanner)
+	sqlRequests = requestsOfType(requests, reflect.TypeOf(&sppb.ExecuteSqlRequest{}))
+	if g, w := len(sqlRequests), 1; g != w {
+		t.Fatalf("sql requests count mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	// Verify that we can get the row that we selected.
+	if !rows.Next() {
+		t.Fatal("no rows")
+	}
+	_ = rows.Close()
+}
+
 func TestConcurrentScanAndClose(t *testing.T) {
 	t.Parallel()