feat: move metadata and stats to separate result sets

Author: olavloite

checksum_row_iterator.go

@@ -252,6 +252,9 @@ func (it *checksumRowIterator) Metadata() (*sppb.ResultSetMetadata, error) {
 	return it.metadata, nil
 }
 
-func (it *checksumRowIterator) RowCount() int64 {
-	return it.RowIterator.RowCount
+func (it *checksumRowIterator) ResultSetStats() *sppb.ResultSetStats {
+	return &sppb.ResultSetStats{
+		RowCount:  &sppb.ResultSetStats_RowCountExact{RowCountExact: it.RowIterator.RowCount},
+		QueryPlan: it.RowIterator.QueryPlan,
+	}
 }

client_side_statement.go

@@ -445,6 +445,6 @@ func (t *clientSideIterator) Metadata() (*spannerpb.ResultSetMetadata, error) {
 	return t.metadata, nil
 }
 
-func (t *clientSideIterator) RowCount() int64 {
-	return 0
+func (t *clientSideIterator) ResultSetStats() *spannerpb.ResultSetStats {
+	return &spannerpb.ResultSetStats{}
 }

conn.go

@@ -815,8 +815,15 @@ func (c *conn) queryContext(ctx context.Context, query string, execOptions ExecO
 			return nil, err
 		}
 	}
-	res := &rows{it: iter, decodeOption: execOptions.DecodeOption, decodeToNativeArrays: execOptions.DecodeToNativeArrays}
-	if execOptions.DecodeOption == DecodeOptionProto {
+	res := &rows{
+		it:                      iter,
+		decodeOption:            execOptions.DecodeOption,
+		decodeToNativeArrays:    execOptions.DecodeToNativeArrays,
+		returnResultSetMetadata: execOptions.ReturnResultSetMetadata,
+		returnResultSetStats:    execOptions.ReturnResultSetStats,
+	}
+	if execOptions.DirectExecute {
+		// This forces the execution of the statement.
 		res.getColumns()
 		if res.dirtyErr != nil && !errors.Is(res.dirtyErr, iterator.Done) {
 			_ = res.Close()

driver.go

@@ -163,6 +163,10 @@ type ExecOptions struct {
 	// AutoCommitDMLMode determines the type of transaction that DML statements
 	// that are executed outside explicit transactions use.
 	AutocommitDMLMode AutocommitDMLMode
+
+	ReturnResultSetMetadata bool
+	DirectExecute           bool
+	ReturnResultSetStats    bool
 }
 
 type DecodeOption int

merged_row_iterator.go

@@ -264,6 +264,6 @@ func (m *mergedRowIterator) Metadata() (*sppb.ResultSetMetadata, error) {
 	return m.metadata, nil
 }
 
-func (m *mergedRowIterator) RowCount() int64 {
-	return 0
+func (m *mergedRowIterator) ResultSetStats() *sppb.ResultSetStats {
+	return &sppb.ResultSetStats{}
 }

rows.go

@@ -16,7 +16,7 @@ package spannerdriver
 
 import (
 	"database/sql/driver"
-	"encoding/base64"
+	"errors"
 	"fmt"
 	"io"
 	"sync"
@@ -27,11 +27,20 @@ import (
 	sppb "cloud.google.com/go/spanner/apiv1/spannerpb"
 	"github.com/google/uuid"
 	"google.golang.org/api/iterator"
-	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/structpb"
 )
 
 var _ driver.RowsColumnTypeDatabaseTypeName = (*rows)(nil)
+var _ driver.RowsNextResultSet = (*rows)(nil)
+
+type resultSetType int
+
+const (
+	resultSetTypeMetadata resultSetType = iota
+	resultSetTypeResults
+	resultSetTypeStats
+	resultSetTypeNoMoreResults
+)
 
 type rows struct {
 	it    rowIterator
@@ -46,6 +55,31 @@ type rows struct {
 	decodeToNativeArrays bool
 
 	dirtyRow *spanner.Row
+
+	currentResultSetType    resultSetType
+	returnResultSetMetadata bool
+	returnResultSetStats    bool
+
+	hasReturnedResultSetMetadata bool
+	hasReturnedResultSetStats    bool
+}
+
+func (r *rows) HasNextResultSet() bool {
+	if r.currentResultSetType == resultSetTypeMetadata && r.returnResultSetMetadata {
+		return true
+	}
+	if r.currentResultSetType == resultSetTypeResults && r.returnResultSetStats {
+		return true
+	}
+	return false
+}
+
+func (r *rows) NextResultSet() error {
+	if !r.HasNextResultSet() {
+		return io.EOF
+	}
+	r.currentResultSetType++
+	return nil
 }
 
 // Columns returns the names of the columns. The number of
@@ -54,12 +88,32 @@ type rows struct {
 // string should be returned for that entry.
 func (r *rows) Columns() []string {
 	r.getColumns()
-	return r.cols
+	switch r.currentResultSetType {
+	case resultSetTypeMetadata:
+		return []string{"metadata"}
+	case resultSetTypeResults:
+		return r.cols
+	case resultSetTypeStats:
+		return []string{"stats"}
+	case resultSetTypeNoMoreResults:
+		return nil
+	}
+	return nil
 }
 
 func (r *rows) ColumnTypeDatabaseTypeName(index int) string {
 	r.getColumns()
-	return r.colTypeNames[index]
+	switch r.currentResultSetType {
+	case resultSetTypeMetadata:
+		return "ResultSetMetadata"
+	case resultSetTypeResults:
+		return r.colTypeNames[index]
+	case resultSetTypeStats:
+		return "ResultSetStats"
+	case resultSetTypeNoMoreResults:
+		return ""
+	}
+	return ""
 }
 
 // Close closes the rows iterator.
@@ -75,6 +129,9 @@ func (r *rows) Close() error {
 
 func (r *rows) getColumns() {
 	r.colsOnce.Do(func() {
+		if r.currentResultSetType == resultSetTypeMetadata && !r.returnResultSetMetadata {
+			r.currentResultSetType = resultSetTypeResults
+		}
 		row, err := r.it.Next()
 		if err == nil {
 			r.dirtyRow = row
@@ -92,22 +149,10 @@ func (r *rows) getColumns() {
 		rowType := metadata.RowType
 		r.cols = make([]string, len(rowType.Fields))
 		r.colTypeNames = make([]string, len(rowType.Fields))
-		if r.decodeOption == DecodeOptionProto {
-			if len(rowType.Fields) == 0 {
-				r.cols = make([]string, 1)
-				r.colTypeNames = make([]string, 1)
-			}
-			metadataBytes, err := proto.Marshal(metadata)
-			if err == nil {
-				r.colTypeNames[0] = base64.StdEncoding.EncodeToString(metadataBytes)
-			}
-			r.cols[0] = fmt.Sprintf("%v", r.it.RowCount())
-		} else {
-			for i, c := range rowType.Fields {
-				r.cols[i] = c.Name
-				if r.decodeOption != DecodeOptionProto {
-					r.colTypeNames[i] = c.Type.Code.String()
-				}
+		for i, c := range rowType.Fields {
+			r.cols[i] = c.Name
+			if r.decodeOption != DecodeOptionProto {
+				r.colTypeNames[i] = c.Type.Code.String()
 			}
 		}
 	})
@@ -124,6 +169,30 @@ func (r *rows) getColumns() {
 // a buffer held in dest.
 func (r *rows) Next(dest []driver.Value) error {
 	r.getColumns()
+	if r.currentResultSetType == resultSetTypeMetadata {
+		if r.dirtyErr != nil && !errors.Is(r.dirtyErr, iterator.Done) {
+			return r.dirtyErr
+		}
+		if r.hasReturnedResultSetMetadata {
+			return io.EOF
+		}
+		r.hasReturnedResultSetMetadata = true
+		metadata, err := r.it.Metadata()
+		if err != nil {
+			return err
+		}
+		dest[0] = metadata
+		return nil
+	}
+	if r.currentResultSetType == resultSetTypeStats {
+		if r.hasReturnedResultSetStats {
+			return io.EOF
+		}
+		r.hasReturnedResultSetStats = true
+		dest[0] = r.it.ResultSetStats()
+		return nil
+	}
+
 	var row *spanner.Row
 	if r.dirtyErr != nil {
 		err := r.dirtyErr
@@ -132,7 +201,8 @@ func (r *rows) Next(dest []driver.Value) error {
 			return io.EOF
 		}
 		return err
-	} else if r.dirtyRow != nil {
+	}
+	if r.dirtyRow != nil {
 		row = r.dirtyRow
 		r.dirtyRow = nil
 	} else {

rows_test.go

@@ -49,8 +49,8 @@ func (t *testIterator) Metadata() (*sppb.ResultSetMetadata, error) {
 	return t.metadata, nil
 }
 
-func (t *testIterator) RowCount() int64 {
-	return 0
+func (t *testIterator) ResultSetStats() *sppb.ResultSetStats {
+	return &sppb.ResultSetStats{}
 }
 
 func newRow(t *testing.T, cols []string, vals []interface{}) *spanner.Row {

spannerlib/exported/connection.go

@@ -4,6 +4,7 @@ import "C"
 import (
 	"context"
 	"database/sql"
+	"fmt"
 	"sync"
 	"sync/atomic"
 
@@ -99,11 +100,11 @@ func (conn *Connection) BeginTransaction(txOpts *spannerpb.TransactionOptions) *
 	if txOpts.GetReadOnly() != nil {
 		tx, err = spannerdriver.BeginReadOnlyTransactionOnConn(
 			context.Background(), conn.backend.Conn, convertToReadOnlyOpts(txOpts))
-	} else if txOpts.GetReadWrite() != nil {
+	} else if txOpts.GetPartitionedDml() != nil {
+		err = spanner.ToSpannerError(status.Error(codes.InvalidArgument, "transaction type not supported"))
+	} else {
 		tx, err = spannerdriver.BeginReadWriteTransactionOnConn(
 			context.Background(), conn.backend.Conn, convertToReadWriteTransactionOptions(txOpts))
-	} else {
-		err = spanner.ToSpannerError(status.Error(codes.InvalidArgument, "transaction type not supported"))
 	}
 	if err != nil {
 		return errMessage(err)
@@ -140,10 +141,14 @@ func convertTimestampBound(txOpts *spannerpb.TransactionOptions) spanner.Timesta
 }
 
 func convertToReadWriteTransactionOptions(txOpts *spannerpb.TransactionOptions) spannerdriver.ReadWriteTransactionOptions {
+	readLockMode := spannerpb.TransactionOptions_ReadWrite_READ_LOCK_MODE_UNSPECIFIED
+	if txOpts.GetReadWrite() != nil {
+		readLockMode = txOpts.GetReadWrite().GetReadLockMode()
+	}
 	return spannerdriver.ReadWriteTransactionOptions{
 		TransactionOptions: spanner.TransactionOptions{
 			IsolationLevel: txOpts.GetIsolationLevel(),
-			ReadLockMode:   txOpts.GetReadWrite().GetReadLockMode(),
+			ReadLockMode:   readLockMode,
 		},
 	}
 }
@@ -172,9 +177,26 @@ func execute(conn *Connection, executor queryExecutor, statement *spannerpb.Exec
 	if err != nil {
 		return errMessage(err)
 	}
+	// The first result set should contain the metadata.
+	if !it.Next() {
+		return errMessage(fmt.Errorf("query returned no metadata"))
+	}
+	metadata := &spannerpb.ResultSetMetadata{}
+	if err := it.Scan(&metadata); err != nil {
+		return errMessage(err)
+	}
+	// Move to the next result set, which contains the normal data.
+	if !it.NextResultSet() {
+		return errMessage(fmt.Errorf("no results found after metadata"))
+	}
 	id := conn.resultsIdx.Add(1)
 	res := &rows{
-		backend: it,
+		backend:  it,
+		metadata: metadata,
+	}
+	if len(metadata.RowType.Fields) == 0 {
+		// No rows returned. Read the stats now.
+		res.readStats()
 	}
 	conn.results.Store(id, res)
 	return idMessage(id)
@@ -224,7 +246,12 @@ func extractParams(statement *spannerpb.ExecuteBatchDmlRequest_Statement) []any
 		paramsLen = 1 + len(statement.Params.Fields)
 	}
 	params := make([]any, paramsLen)
-	params = append(params, spannerdriver.ExecOptions{DecodeOption: spannerdriver.DecodeOptionProto})
+	params = append(params, spannerdriver.ExecOptions{
+		DecodeOption:            spannerdriver.DecodeOptionProto,
+		ReturnResultSetMetadata: true,
+		ReturnResultSetStats:    true,
+		DirectExecute:           true,
+	})
 	if statement.Params != nil {
 		if statement.ParamTypes == nil {
 			statement.ParamTypes = make(map[string]*spannerpb.Type)

spannerlib/exported/pool_test.go

@@ -68,7 +68,7 @@ func TestExecuteDml(t *testing.T) {
 	if len(metadataValue.RowType.Fields) > 0 {
 		fmt.Printf("Row type: %v\n", metadataValue.RowType)
 	} else {
-		rowCount := UpdateCount(pool.ObjectId, conn.ObjectId, rows.ObjectId)
+		rowCount := ResultSetStats(pool.ObjectId, conn.ObjectId, rows.ObjectId)
 		fmt.Printf("Update count: %v\n", string(rowCount.Res))
 	}
 	for {