chore: add Execute function for SpannerLib

Adds an Execute function for SpannerLib that can be used to execute any type
of SQL statement. The return type is always a Rows object. The Rows object
is empty for DDL statements, it only contains ResultSetStats for DML statements
without a THEN RETURN clause, and it contains actual row data for queries and
DML statements with a THEN RETURN clause.

The Execute function can also be used to execute client-side SQL statements,
like BEGIN, COMMIT, SET, SHOW, etc.

Author: olavloite

spannerlib/api/connection.go

@@ -17,8 +17,14 @@ package api
 import (
 	"context"
 	"database/sql"
+	"fmt"
+	"strings"
 	"sync"
 	"sync/atomic"
+
+	"cloud.google.com/go/spanner"
+	"cloud.google.com/go/spanner/apiv1/spannerpb"
+	spannerdriver "github.com/googleapis/go-sql-spanner"
 )
 
 // CloseConnection looks up the connection with the given poolId and connId and closes it.
@@ -36,6 +42,14 @@ func CloseConnection(ctx context.Context, poolId, connId int64) error {
 	return conn.close(ctx)
 }
 
+func Execute(ctx context.Context, poolId, connId int64, executeSqlRequest *spannerpb.ExecuteSqlRequest) (int64, error) {
+	conn, err := findConnection(poolId, connId)
+	if err != nil {
+		return 0, err
+	}
+	return conn.Execute(ctx, executeSqlRequest)
+}
+
 type Connection struct {
 	// results contains the open query results for this connection.
 	results    *sync.Map
@@ -45,6 +59,11 @@ type Connection struct {
 	backend *sql.Conn
 }
 
+type queryExecutor interface {
+	ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error)
+	QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error)
+}
+
 func (conn *Connection) close(ctx context.Context) error {
 	conn.closeResults(ctx)
 	err := conn.backend.Close()
@@ -60,3 +79,73 @@ func (conn *Connection) closeResults(ctx context.Context) {
 		return true
 	})
 }
+
+func (conn *Connection) Execute(ctx context.Context, statement *spannerpb.ExecuteSqlRequest) (int64, error) {
+	return execute(ctx, conn, conn.backend, statement)
+}
+
+func execute(ctx context.Context, conn *Connection, executor queryExecutor, statement *spannerpb.ExecuteSqlRequest) (int64, error) {
+	params := extractParams(statement)
+	it, err := executor.QueryContext(ctx, statement.Sql, params...)
+	if err != nil {
+		return 0, err
+	}
+	// The first result set should contain the metadata.
+	if !it.Next() {
+		return 0, fmt.Errorf("query returned no metadata")
+	}
+	metadata := &spannerpb.ResultSetMetadata{}
+	if err := it.Scan(&metadata); err != nil {
+		return 0, err
+	}
+	// Move to the next result set, which contains the normal data.
+	if !it.NextResultSet() {
+		return 0, fmt.Errorf("no results found after metadata")
+	}
+	id := conn.resultsIdx.Add(1)
+	res := &rows{
+		backend:  it,
+		metadata: metadata,
+	}
+	if len(metadata.RowType.Fields) == 0 {
+		// No rows returned. Read the stats now.
+		res.readStats(ctx)
+	}
+	conn.results.Store(id, res)
+	return id, nil
+}
+
+func extractParams(statement *spannerpb.ExecuteSqlRequest) []any {
+	paramsLen := 1
+	if statement.Params != nil {
+		paramsLen = 1 + len(statement.Params.Fields)
+	}
+	params := make([]any, paramsLen)
+	params = append(params, spannerdriver.ExecOptions{
+		DecodeOption: spannerdriver.DecodeOptionProto,
+		// TODO: Implement support for passing in stale query options
+		// TimestampBound:          extractTimestampBound(statement),
+		ReturnResultSetMetadata: true,
+		ReturnResultSetStats:    true,
+		DirectExecuteQuery:      true,
+	})
+	if statement.Params != nil {
+		if statement.ParamTypes == nil {
+			statement.ParamTypes = make(map[string]*spannerpb.Type)
+		}
+		for param, value := range statement.Params.Fields {
+			genericValue := spanner.GenericColumnValue{
+				Value: value,
+				Type:  statement.ParamTypes[param],
+			}
+			if strings.HasPrefix(param, "_") {
+				// Prefix the parameter name with a 'p' to work around the fact that database/sql does not allow
+				// named arguments to start with anything else than a letter.
+				params = append(params, sql.Named("p"+param, spannerdriver.SpannerNamedArg{NameInQuery: param, Value: genericValue}))
+			} else {
+				params = append(params, sql.Named(param, genericValue))
+			}
+		}
+	}
+	return params
+}

spannerlib/api/pool.go

@@ -17,6 +17,7 @@ package api
 import (
 	"context"
 	"database/sql"
+	"fmt"
 	"sync"
 	"sync/atomic"
 
@@ -131,3 +132,16 @@ func findConnection(poolId, connId int64) (*Connection, error) {
 	conn := c.(*Connection)
 	return conn, nil
 }
+
+func findRows(poolId, connId, rowsId int64) (*rows, error) {
+	conn, err := findConnection(poolId, connId)
+	if err != nil {
+		return nil, err
+	}
+	r, ok := conn.results.Load(rowsId)
+	if !ok {
+		return nil, fmt.Errorf("rows %v not found", rowsId)
+	}
+	res := r.(*rows)
+	return res, nil
+}

spannerlib/api/rows.go

@@ -0,0 +1,203 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package api
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+
+	"cloud.google.com/go/spanner"
+	"cloud.google.com/go/spanner/apiv1/spannerpb"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/structpb"
+)
+
+type EncodeRowOption int32
+
+const (
+	EncodeRowOptionProto EncodeRowOption = iota
+)
+
+// Metadata returns the ResultSetMetadata of the given rows.
+// This function can be called for any type of statement (queries, DML, DDL).
+func Metadata(_ context.Context, poolId, connId, rowsId int64) (*spannerpb.ResultSetMetadata, error) {
+	res, err := findRows(poolId, connId, rowsId)
+	if err != nil {
+		return nil, err
+	}
+	return res.Metadata()
+}
+
+// ResultSetStats returns the result statistics of the given rows.
+// This function can only be called once all data in the rows have been fetched.
+// The stats are empty for queries and DDL statements.
+func ResultSetStats(ctx context.Context, poolId, connId, rowsId int64) (*spannerpb.ResultSetStats, error) {
+	res, err := findRows(poolId, connId, rowsId)
+	if err != nil {
+		return nil, err
+	}
+	return res.ResultSetStats(ctx)
+}
+
+// NextEncoded returns the next row data in encoded form.
+// Using NextEncoded instead of Next can be more efficient for large result sets,
+// as it allows the library to re-use the encoding buffer.
+// TODO: Add an encoder function as input argument, instead of hardcoding protobuf encoding here.
+func NextEncoded(ctx context.Context, poolId, connId, rowsId int64) ([]byte, error) {
+	_, bytes, err := next(ctx, poolId, connId, rowsId, true)
+	if err != nil {
+		return nil, err
+	}
+	return bytes, nil
+}
+
+// Next returns the next row as a protobuf ListValue.
+func Next(ctx context.Context, poolId, connId, rowsId int64) (*structpb.ListValue, error) {
+	values, _, err := next(ctx, poolId, connId, rowsId, false)
+	if err != nil {
+		return nil, err
+	}
+	return values, nil
+}
+
+// next returns the next row of data.
+// The row is returned as a protobuf ListValue if marshalResult==false.
+// The row is returned as a byte slice if marshalResult==true.
+// TODO: Add generics to the function and add input arguments for encoding instead of hardcoding it.
+func next(ctx context.Context, poolId, connId, rowsId int64, marshalResult bool) (*structpb.ListValue, []byte, error) {
+	rows, err := findRows(poolId, connId, rowsId)
+	if err != nil {
+		return nil, nil, err
+	}
+	values, err := rows.Next(ctx)
+	if err != nil {
+		return nil, nil, err
+	}
+	if !marshalResult || values == nil {
+		return values, nil, nil
+	}
+
+	rows.marshalBuffer, err = proto.MarshalOptions{}.MarshalAppend(rows.marshalBuffer[:0], rows.values)
+	if err != nil {
+		return nil, nil, err
+	}
+	return values, rows.marshalBuffer, nil
+}
+
+// CloseRows closes the given rows. Callers must always call this to clean up any resources
+// that are held by the underlying cursor.
+func CloseRows(ctx context.Context, poolId, connId, rowsId int64) error {
+	conn, err := findConnection(poolId, connId)
+	if err != nil {
+		return err
+	}
+	r, ok := conn.results.LoadAndDelete(rowsId)
+	if !ok {
+		return nil
+	}
+	res := r.(*rows)
+	return res.Close(ctx)
+}
+
+type rows struct {
+	backend  *sql.Rows
+	metadata *spannerpb.ResultSetMetadata
+	stats    *spannerpb.ResultSetStats
+	done     bool
+
+	buffer        []any
+	values        *structpb.ListValue
+	marshalBuffer []byte
+}
+
+func (rows *rows) Close(ctx context.Context) error {
+	err := rows.backend.Close()
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (rows *rows) Metadata() (*spannerpb.ResultSetMetadata, error) {
+	return rows.metadata, nil
+}
+
+func (rows *rows) ResultSetStats(ctx context.Context) (*spannerpb.ResultSetStats, error) {
+	if rows.stats == nil {
+		rows.readStats(ctx)
+	}
+	return rows.stats, nil
+}
+
+type genericValue struct {
+	v *structpb.Value
+}
+
+func (gv *genericValue) Scan(src any) error {
+	if v, ok := src.(spanner.GenericColumnValue); ok {
+		gv.v = v.Value
+		return nil
+	}
+	return errors.New("cannot convert value to generic column value")
+}
+
+func (rows *rows) Next(ctx context.Context) (*structpb.ListValue, error) {
+	// No columns means no rows, so just return nil to indicate that there are no (more) rows.
+	if len(rows.metadata.RowType.Fields) == 0 || rows.done {
+		return nil, nil
+	}
+	if rows.stats != nil {
+		return nil, spanner.ToSpannerError(status.Error(codes.FailedPrecondition, "cannot read more data after returning stats"))
+	}
+	ok := rows.backend.Next()
+	if !ok {
+		rows.done = true
+		// No more rows. Read stats and return nil.
+		rows.readStats(ctx)
+		// nil indicates no more rows.
+		return nil, nil
+	}
+
+	if rows.buffer == nil {
+		rows.buffer = make([]any, len(rows.metadata.RowType.Fields))
+		for i := range rows.buffer {
+			rows.buffer[i] = &genericValue{}
+		}
+		rows.values = &structpb.ListValue{
+			Values: make([]*structpb.Value, len(rows.buffer)),
+		}
+		rows.marshalBuffer = make([]byte, 0)
+	}
+	if err := rows.backend.Scan(rows.buffer...); err != nil {
+		return nil, err
+	}
+	for i := range rows.buffer {
+		rows.values.Values[i] = rows.buffer[i].(*genericValue).v
+	}
+	return rows.values, nil
+}
+
+func (rows *rows) readStats(ctx context.Context) {
+	rows.stats = &spannerpb.ResultSetStats{}
+	if !rows.backend.NextResultSet() {
+		return
+	}
+	if rows.backend.Next() {
+		_ = rows.backend.Scan(&rows.stats)
+	}
+}