chore: add multi-statement support to SpannerLib

Adds support for multi-statement SQL strings to SpannerLib. This adds a new exported
function NextResultSet that moves the pointer to the next result for a multi-statement
SQL string.

It also updates the return type for the gRPC server, so this returns an indication
whether the result contains more result sets.

These changes are backwards-compatible with existing implementations for Rows. Wrappers
are free to ignore any result sets after the first one.

Author: olavloite

spannerlib/api/connection.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"database/sql"
 	"database/sql/driver"
-	"fmt"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -322,27 +321,14 @@ func execute(ctx, directExecuteContext context.Context, conn *Connection, execut
 	if err != nil {
 		return 0, err
 	}
-	// The first result set should contain the metadata.
-	if !it.Next() {
-		_ = it.Close()
-		return 0, fmt.Errorf("query returned no metadata")
+	res := &rows{
+		backend: it,
 	}
-	metadata := &spannerpb.ResultSetMetadata{}
-	if err := it.Scan(&metadata); err != nil {
-		_ = it.Close()
+	if err := res.readMetadata(ctx); err != nil {
 		return 0, err
 	}
-	// Move to the next result set, which contains the normal data.
-	if !it.NextResultSet() {
-		_ = it.Close()
-		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 {
+	if len(res.metadata.RowType.Fields) == 0 {
 		// No rows returned. Read the stats now.
 		_ = res.readStats(ctx)
 	}

spannerlib/api/rows.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"database/sql"
 	"errors"
+	"fmt"
 
 	"cloud.google.com/go/spanner"
 	"cloud.google.com/go/spanner/apiv1/spannerpb"
@@ -54,6 +55,16 @@ func ResultSetStats(ctx context.Context, poolId, connId, rowsId int64) (*spanner
 	return res.ResultSetStats(ctx)
 }
 
+// NextResultSet returns the ResultSetMetadata of the next result set in the given rows or nil if there are no
+// more result sets in the given Rows object.
+func NextResultSet(ctx context.Context, poolId, connId, rowsId int64) (*spannerpb.ResultSetMetadata, error) {
+	res, err := findRows(poolId, connId, rowsId)
+	if err != nil {
+		return nil, err
+	}
+	return res.NextResultSet(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.
@@ -146,6 +157,32 @@ func (rows *rows) ResultSetStats(ctx context.Context) (*spannerpb.ResultSetStats
 	return rows.stats, nil
 }
 
+func (rows *rows) NextResultSet(ctx context.Context) (*spannerpb.ResultSetMetadata, error) {
+	if !rows.done {
+		// The current result set has not been read to the end.
+		// We therefore need to move the cursor to the next result set contains the stats for the current result set,
+		// so we can skip those.
+		rows.backend.NextResultSet()
+	}
+	if rows.backend.NextResultSet() {
+		rows.done = false
+		rows.stats = nil
+		if err := rows.readMetadata(ctx); err != nil {
+			return nil, err
+		}
+		if len(rows.metadata.RowType.Fields) == 0 {
+			if err := rows.readStats(ctx); err != nil {
+				return nil, err
+			}
+		}
+		return rows.metadata, nil
+	}
+	if err := rows.backend.Err(); err != nil {
+		return nil, err
+	}
+	return nil, nil
+}
+
 type genericValue struct {
 	v *structpb.Value
 }
@@ -199,6 +236,25 @@ func (rows *rows) Next(ctx context.Context) (*structpb.ListValue, error) {
 	return rows.values, nil
 }
 
+func (rows *rows) readMetadata(ctx context.Context) error {
+	// The first result set should contain the metadata.
+	if !rows.backend.Next() {
+		_ = rows.backend.Close()
+		return fmt.Errorf("query returned no metadata")
+	}
+	rows.metadata = &spannerpb.ResultSetMetadata{}
+	if err := rows.backend.Scan(&rows.metadata); err != nil {
+		_ = rows.backend.Close()
+		return err
+	}
+	// Move to the next result set, which contains the normal data.
+	if !rows.backend.NextResultSet() {
+		_ = rows.backend.Close()
+		return fmt.Errorf("no results found after metadata")
+	}
+	return nil
+}
+
 func (rows *rows) readStats(ctx context.Context) error {
 	rows.stats = &spannerpb.ResultSetStats{}
 	if !rows.backend.NextResultSet() {

spannerlib/api/rows_test.go

@@ -82,6 +82,142 @@ func TestExecute(t *testing.T) {
 	}
 }
 
+func TestExecuteMultiStatement(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	server, teardown := setupMockServer(t)
+	defer teardown()
+	dsn := fmt.Sprintf("%s/projects/p/instances/i/databases/d?useplaintext=true", server.Address)
+
+	poolId, err := CreatePool(ctx, dsn)
+	if err != nil {
+		t.Fatalf("CreatePool returned unexpected error: %v", err)
+	}
+	connId, err := CreateConnection(ctx, poolId)
+	if err != nil {
+		t.Fatalf("CreateConnection returned unexpected error: %v", err)
+	}
+	rowsId, err := Execute(ctx, poolId, connId, &spannerpb.ExecuteSqlRequest{
+		Sql: fmt.Sprintf("%s;%s", testutil.SelectFooFromBar, testutil.SelectFooFromBar),
+	})
+	if rowsId == 0 {
+		t.Fatal("Execute returned unexpected zero id")
+	}
+
+	totalRowCount := 0
+	for {
+		rowCount := 0
+		for {
+			row, err := Next(ctx, poolId, connId, rowsId)
+			if err != nil {
+				t.Fatalf("Next returned unexpected error: %v", err)
+			}
+			if row == nil {
+				break
+			}
+			rowCount++
+			totalRowCount++
+		}
+		if g, w := rowCount, 2; g != w {
+			t.Fatalf("row count mismatch\n Got: %d\nWant: %d", g, w)
+		}
+		metadata, err := NextResultSet(ctx, poolId, connId, rowsId)
+		if err != nil {
+			t.Fatalf("NextResultSet returned unexpected error: %v", err)
+		}
+		if metadata == nil {
+			break
+		}
+	}
+	if g, w := totalRowCount, 4; g != w {
+		t.Fatalf("total row count mismatch\n Got: %d\nWant: %d", g, w)
+	}
+
+	if err := CloseRows(ctx, poolId, connId, rowsId); err != nil {
+		t.Fatalf("CloseRows returned unexpected error: %v", err)
+	}
+	if err := CloseConnection(ctx, poolId, connId); err != nil {
+		t.Fatalf("CloseConnection returned unexpected error: %v", err)
+	}
+	if err := ClosePool(ctx, poolId); err != nil {
+		t.Fatalf("ClosePool returned unexpected error: %v", err)
+	}
+
+	requests := server.TestSpanner.DrainRequestsFromServer()
+	executeRequests := testutil.RequestsOfType(requests, reflect.TypeOf(&spannerpb.ExecuteSqlRequest{}))
+	if g, w := len(executeRequests), 2; g != w {
+		t.Fatalf("num ExecuteSql requests mismatch\n Got: %d\nWant: %d", g, w)
+	}
+}
+
+func TestExecuteMultiStatement_MoveToNextResultSetHalfway(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	server, teardown := setupMockServer(t)
+	defer teardown()
+	dsn := fmt.Sprintf("%s/projects/p/instances/i/databases/d?useplaintext=true", server.Address)
+
+	poolId, err := CreatePool(ctx, dsn)
+	if err != nil {
+		t.Fatalf("CreatePool returned unexpected error: %v", err)
+	}
+	connId, err := CreateConnection(ctx, poolId)
+	if err != nil {
+		t.Fatalf("CreateConnection returned unexpected error: %v", err)
+	}
+	rowsId, err := Execute(ctx, poolId, connId, &spannerpb.ExecuteSqlRequest{
+		Sql: fmt.Sprintf("%s;%s", testutil.SelectFooFromBar, testutil.SelectFooFromBar),
+	})
+	if rowsId == 0 {
+		t.Fatal("Execute returned unexpected zero id")
+	}
+
+	// Read one row from the first result set.
+	row, err := Next(ctx, poolId, connId, rowsId)
+	if err != nil {
+		t.Fatalf("Next returned unexpected error: %v", err)
+	}
+	if row == nil {
+		t.Fatal("Next returned unexpected nil row")
+	}
+	// Then move to the next result set.
+	metadata, err := NextResultSet(ctx, poolId, connId, rowsId)
+	if err != nil {
+		t.Fatalf("NextResultSet returned unexpected error: %v", err)
+	}
+	if metadata == nil {
+		t.Fatal("NextResultSet returned unexpected nil metadata")
+	}
+	// Try to read two rows from the second result set.
+	for range 2 {
+		row, err := Next(ctx, poolId, connId, rowsId)
+		if err != nil {
+			t.Fatalf("Next returned unexpected error: %v", err)
+		}
+		if row == nil {
+			t.Fatal("Next returned unexpected nil row")
+		}
+	}
+
+	if err := CloseRows(ctx, poolId, connId, rowsId); err != nil {
+		t.Fatalf("CloseRows returned unexpected error: %v", err)
+	}
+	if err := CloseConnection(ctx, poolId, connId); err != nil {
+		t.Fatalf("CloseConnection returned unexpected error: %v", err)
+	}
+	if err := ClosePool(ctx, poolId); err != nil {
+		t.Fatalf("ClosePool returned unexpected error: %v", err)
+	}
+
+	requests := server.TestSpanner.DrainRequestsFromServer()
+	executeRequests := testutil.RequestsOfType(requests, reflect.TypeOf(&spannerpb.ExecuteSqlRequest{}))
+	if g, w := len(executeRequests), 2; g != w {
+		t.Fatalf("num ExecuteSql requests mismatch\n Got: %d\nWant: %d", g, w)
+	}
+}
+
 func TestExecuteUnknownConnection(t *testing.T) {
 	t.Parallel()
 

spannerlib/grpc-server/build-protos.sh

@@ -26,5 +26,5 @@ protoc \
   google/spannerlib/v1/spannerlib.proto
 cd .. || exit 1
 rm googleapis/google/spannerlib
-git rm googleapis
+git rm googleapis -f
 rm ../../.gitmodules

spannerlib/grpc-server/google/spannerlib/v1/spannerlib.pb.go

@@ -143,6 +143,7 @@ message RowData {
     (google.api.field_behavior) = REQUIRED
   ];
   google.spanner.v1.ResultSetStats stats = 4;
+  bool has_more_results = 5;
 }
 
 message MetadataRequest {

spannerlib/grpc-server/google/spannerlib/v1/spannerlib.proto

@@ -133,6 +133,7 @@ func (s *spannerLibServer) ExecuteStreaming(request *pb.ExecuteRequest, stream g
 	if err != nil {
 		return err
 	}
+
 	first := true
 	for {
 		if row, err := api.Next(queryContext, request.Connection.Pool.Id, request.Connection.Id, id); err != nil {
@@ -143,14 +144,23 @@ func (s *spannerLibServer) ExecuteStreaming(request *pb.ExecuteRequest, stream g
 				if err != nil {
 					return err
 				}
-				res := &pb.RowData{Rows: rows, Stats: stats}
+				nextMetadata, err := api.NextResultSet(queryContext, request.Connection.Pool.Id, request.Connection.Id, id)
+				if err != nil {
+					return err
+				}
+				res := &pb.RowData{Rows: rows, Stats: stats, HasMoreResults: nextMetadata != nil}
 				if first {
 					res.Metadata = metadata
 					first = false
 				}
 				if err := stream.Send(res); err != nil {
 					return err
 				}
+				if res.HasMoreResults {
+					metadata = nextMetadata
+					first = true
+					continue
+				}
 				break
 			}
 			res := &pb.RowData{Rows: rows, Data: []*structpb.ListValue{row}}

spannerlib/grpc-server/server.go

@@ -264,13 +264,18 @@ func TestExecuteStreamingMultiStatement(t *testing.T) {
 	if err != nil {
 		t.Fatalf("failed to execute: %v", err)
 	}
+	numResultSets := 1
 	numRows := 0
 	for {
 		row, err := stream.Recv()
 		if err != nil {
 			t.Fatalf("failed to receive row: %v", err)
 		}
 		if len(row.Data) == 0 {
+			if row.HasMoreResults {
+				numResultSets++
+				continue
+			}
 			break
 		}
 		if g, w := len(row.Data), 1; g != w {
@@ -281,9 +286,12 @@ func TestExecuteStreamingMultiStatement(t *testing.T) {
 		}
 		numRows++
 	}
-	if g, w := numRows, 2; g != w {
+	if g, w := numRows, 4; g != w {
 		t.Fatalf("num rows mismatch\n Got: %v\nWant: %v", g, w)
 	}
+	if g, w := numResultSets, 2; g != w {
+		t.Fatalf("num result sets mismatch\n Got: %v\nWant: %v", g, w)
+	}
 
 	if _, err := client.ClosePool(ctx, pool); err != nil {
 		t.Fatalf("failed to close pool: %v", err)