chore: add more features

Author: olavloite

checksum_row_iterator.go

@@ -41,6 +41,8 @@ func init() {
 	gob.Register(structpb.Value_StructValue{})
 }
 
+var _ rowIterator = (*checksumRowIterator)(nil)
+
 // checksumRowIterator implements rowIterator and keeps track of a running
 // checksum for all results that have been seen during the iteration of the
 // results. This checksum can be used to verify whether a retry returned the
@@ -249,3 +251,7 @@ func (it *checksumRowIterator) Stop() {
 func (it *checksumRowIterator) Metadata() (*sppb.ResultSetMetadata, error) {
 	return it.metadata, nil
 }
+
+func (it *checksumRowIterator) RowCount() int64 {
+	return it.RowIterator.RowCount
+}

client_side_statement.go

@@ -414,6 +414,8 @@ func createSingleValueIterator(column string, value interface{}, code spannerpb.
 	}, nil
 }
 
+var _ rowIterator = (*clientSideIterator)(nil)
+
 // clientSideIterator implements the rowIterator interface for client side
 // statements. All values are created and kept in memory, and this struct
 // should only be used for small result sets.
@@ -442,3 +444,7 @@ func (t *clientSideIterator) Stop() {
 func (t *clientSideIterator) Metadata() (*spannerpb.ResultSetMetadata, error) {
 	return t.metadata, nil
 }
+
+func (t *clientSideIterator) RowCount() int64 {
+	return 0
+}

conn.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"database/sql"
 	"database/sql/driver"
+	"errors"
 	"log/slog"
 	"slices"
 	"time"
@@ -792,7 +793,13 @@ 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}
+	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) {

merged_row_iterator.go

@@ -263,3 +263,7 @@ func (m *mergedRowIterator) Metadata() (*sppb.ResultSetMetadata, error) {
 	}
 	return m.metadata, nil
 }
+
+func (m *mergedRowIterator) RowCount() int64 {
+	return 0
+}

rows.go

@@ -16,6 +16,7 @@ package spannerdriver
 
 import (
 	"database/sql/driver"
+	"encoding/base64"
 	"fmt"
 	"io"
 	"sync"
@@ -26,16 +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)
+
 type rows struct {
 	it    rowIterator
 	close func() error
 
-	colsOnce sync.Once
-	dirtyErr error
-	cols     []string
+	colsOnce     sync.Once
+	dirtyErr     error
+	cols         []string
+	colTypeNames []string
 
 	decodeOption         DecodeOption
 	decodeToNativeArrays bool
@@ -52,6 +57,11 @@ func (r *rows) Columns() []string {
 	return r.cols
 }
 
+func (r *rows) ColumnTypeDatabaseTypeName(index int) string {
+	r.getColumns()
+	return r.colTypeNames[index]
+}
+
 // Close closes the rows iterator.
 func (r *rows) Close() error {
 	r.it.Stop()
@@ -81,8 +91,24 @@ func (r *rows) getColumns() {
 		}
 		rowType := metadata.RowType
 		r.cols = make([]string, len(rowType.Fields))
-		for i, c := range rowType.Fields {
-			r.cols[i] = c.Name
+		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()
+				}
+			}
 		}
 	})
 }

spannerlib/backend/build.sh renamed to spannerlib/build.sh

@@ -4,7 +4,6 @@ import "C"
 import (
 	"context"
 	"database/sql"
-	"fmt"
 	"sync"
 	"sync/atomic"
 
@@ -23,13 +22,23 @@ func CloseConnection(poolId, connId int64) *Message {
 	return conn.close()
 }
 
+func BeginTransaction(poolId, connId int64, txOptsBytes []byte) *Message {
+	txOpts := spannerpb.TransactionOptions{}
+	if err := proto.Unmarshal(txOptsBytes, &txOpts); err != nil {
+		return errMessage(err)
+	}
+	conn, err := findConnection(poolId, connId)
+	if err != nil {
+		return errMessage(err)
+	}
+	return conn.BeginTransaction(&txOpts)
+}
+
 func Execute(poolId, connId int64, statementBytes []byte) *Message {
 	statement := spannerpb.ExecuteBatchDmlRequest_Statement{}
 	if err := proto.Unmarshal(statementBytes, &statement); err != nil {
 		return errMessage(err)
 	}
-	fmt.Printf("Statement: %v\n", statement.Sql)
-	fmt.Printf("Params: %v\n", statement.Params)
 	conn, err := findConnection(poolId, connId)
 	if err != nil {
 		return errMessage(err)
@@ -41,6 +50,9 @@ type Connection struct {
 	results    *sync.Map
 	resultsIdx atomic.Int64
 
+	transactions    *sync.Map
+	transactionsIdx atomic.Int64
+
 	backend *backend.SpannerConnection
 }
 
@@ -57,6 +69,32 @@ func (conn *Connection) close() *Message {
 	return &Message{}
 }
 
+func (conn *Connection) BeginTransaction(txOpts *spannerpb.TransactionOptions) *Message {
+	tx, err := conn.backend.Conn.BeginTx(context.Background(), &sql.TxOptions{
+		Isolation: convertIsolationLevel(txOpts.IsolationLevel),
+		ReadOnly:  txOpts.GetReadOnly() != nil,
+	})
+	if err != nil {
+		return errMessage(err)
+	}
+	id := conn.transactionsIdx.Add(1)
+	res := &transaction{
+		backend: tx,
+	}
+	conn.transactions.Store(id, res)
+	return idMessage(id)
+}
+
+func convertIsolationLevel(level spannerpb.TransactionOptions_IsolationLevel) sql.IsolationLevel {
+	switch level {
+	case spannerpb.TransactionOptions_SERIALIZABLE:
+		return sql.LevelSerializable
+	case spannerpb.TransactionOptions_REPEATABLE_READ:
+		return sql.LevelRepeatableRead
+	}
+	return sql.LevelDefault
+}
+
 func (conn *Connection) Execute(statement *spannerpb.ExecuteBatchDmlRequest_Statement) *Message {
 	paramsLen := 1
 	if statement.Params != nil {

spannerlib/exported/connection.go

@@ -56,8 +56,9 @@ func CreateConnection(poolId int64, project, instance, database string) *Message
 	}
 	id := poolsIdx.Add(1)
 	conn := &Connection{
-		backend: &backend.SpannerConnection{Conn: sqlConn},
-		results: &sync.Map{},
+		backend:      &backend.SpannerConnection{Conn: sqlConn},
+		results:      &sync.Map{},
+		transactions: &sync.Map{},
 	}
 	pool.connections.Store(id, conn)
 
@@ -90,3 +91,16 @@ func findRows(poolId, connId, rowsId int64) (*rows, error) {
 	res := r.(*rows)
 	return res, nil
 }
+
+func findTx(poolId, connId, txId int64) (*transaction, error) {
+	conn, err := findConnection(poolId, connId)
+	if err != nil {
+		return nil, err
+	}
+	r, ok := conn.transactions.Load(txId)
+	if !ok {
+		return nil, fmt.Errorf("tx %v not found", txId)
+	}
+	res := r.(*transaction)
+	return res, nil
+}

spannerlib/exported/pool.go

@@ -1,11 +1,12 @@
 package exported
 
 import (
-	"cloud.google.com/go/spanner/apiv1/spannerpb"
 	"fmt"
+	"testing"
+
+	"cloud.google.com/go/spanner/apiv1/spannerpb"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/structpb"
-	"testing"
 )
 
 func TestExecute(t *testing.T) {
@@ -19,6 +20,57 @@ func TestExecute(t *testing.T) {
 	}
 	stmtBytes, _ := proto.Marshal(&stmt)
 	rows := Execute(pool.ObjectId, conn.ObjectId, stmtBytes)
+	metadata := Metadata(pool.ObjectId, conn.ObjectId, rows.ObjectId)
+	metadataValue := spannerpb.ResultSetMetadata{}
+	_ = proto.Unmarshal(metadata.Res, &metadataValue)
+	fmt.Printf("Row type: %v\n", metadataValue.RowType)
+	for {
+		row := Next(pool.ObjectId, conn.ObjectId, rows.ObjectId)
+		rowValue := structpb.ListValue{}
+		_ = proto.Unmarshal(row.Res, &rowValue)
+		if row.Length() == 0 {
+			break
+		}
+		fmt.Printf("row: %v\n", rowValue.Values)
+	}
+	CloseRows(pool.ObjectId, conn.ObjectId, rows.ObjectId)
+	CloseConnection(pool.ObjectId, conn.ObjectId)
+	ClosePool(pool.ObjectId)
+}
+
+func TestExecuteDml(t *testing.T) {
+	pool := CreatePool()
+	conn := CreateConnection(pool.ObjectId, "appdev-soda-spanner-staging", "knut-test-ycsb", "knut-test-db")
+	txOpts := &spannerpb.TransactionOptions{
+		Mode: &spannerpb.TransactionOptions_ReadOnly_{
+			ReadOnly: &spannerpb.TransactionOptions_ReadOnly{},
+		},
+	}
+	txOptsBytes, _ := proto.Marshal(txOpts)
+	BeginTransaction(pool.ObjectId, conn.ObjectId, txOptsBytes)
+	stmt := spannerpb.ExecuteBatchDmlRequest_Statement{
+		Sql: "update all_types set col_float8=$1 where col_varchar=$2",
+		Params: &structpb.Struct{
+			Fields: map[string]*structpb.Value{
+				"p1": {Kind: &structpb.Value_NumberValue{NumberValue: 3.14}},
+				"p2": {Kind: &structpb.Value_StringValue{StringValue: "61763b0e7feb3ea8fc9e734a6700f6a4"}},
+			},
+		},
+	}
+	stmtBytes, _ := proto.Marshal(&stmt)
+	rows := Execute(pool.ObjectId, conn.ObjectId, stmtBytes)
+	if rows.Code != 0 {
+		t.Fatalf("failed to execute statement: %s", string(rows.Res))
+	}
+	metadata := Metadata(pool.ObjectId, conn.ObjectId, rows.ObjectId)
+	metadataValue := spannerpb.ResultSetMetadata{}
+	_ = proto.Unmarshal(metadata.Res, &metadataValue)
+	if len(metadataValue.RowType.Fields) > 0 {
+		fmt.Printf("Row type: %v\n", metadataValue.RowType)
+	} else {
+		rowCount := UpdateCount(pool.ObjectId, conn.ObjectId, rows.ObjectId)
+		fmt.Printf("Update count: %v\n", string(rowCount.Res))
+	}
 	for {
 		row := Next(pool.ObjectId, conn.ObjectId, rows.ObjectId)
 		rowValue := structpb.ListValue{}