feat: add option to return metadata and stats

Adds an option to return the full ResultSetMetadata and
ResultSetStats as separate result sets. This gives the caller
access to more details than can be provided through the standard
database/sql interfaces.

Author: olavloite

checksum_row_iterator.go

@@ -41,6 +41,8 @@ func init() {
 	gob.Register(structpb.Value_StructValue{})
 }
 
+var _ rowIterator = &checksumRowIterator{}
+
 // 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,12 @@ func (it *checksumRowIterator) Stop() {
 func (it *checksumRowIterator) Metadata() (*sppb.ResultSetMetadata, error) {
 	return it.metadata, nil
 }
+
+func (it *checksumRowIterator) ResultSetStats() *sppb.ResultSetStats {
+	// TODO: The Spanner client library should offer an option to get the full
+	//       ResultSetStats, instead of only the RowCount and QueryPlan.
+	return &sppb.ResultSetStats{
+		RowCount:  &sppb.ResultSetStats_RowCountExact{RowCountExact: it.RowIterator.RowCount},
+		QueryPlan: it.RowIterator.QueryPlan,
+	}
+}

client_side_statement.go

@@ -414,6 +414,8 @@ func createSingleValueIterator(column string, value interface{}, code spannerpb.
 	}, nil
 }
 
+var _ rowIterator = &clientSideIterator{}
+
 // 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) ResultSetStats() *spannerpb.ResultSetStats {
+	return &spannerpb.ResultSetStats{}
+}

conn.go

@@ -793,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
+	return &rows{
+		it:                      iter,
+		decodeOption:            execOptions.DecodeOption,
+		decodeToNativeArrays:    execOptions.DecodeToNativeArrays,
+		returnResultSetMetadata: execOptions.ReturnResultSetMetadata,
+		returnResultSetStats:    execOptions.ReturnResultSetStats,
+	}, nil
 }
 
 func (c *conn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {

driver.go

@@ -163,6 +163,28 @@ type ExecOptions struct {
 	// AutoCommitDMLMode determines the type of transaction that DML statements
 	// that are executed outside explicit transactions use.
 	AutocommitDMLMode AutocommitDMLMode
+
+	// ReturnResultSetMetadata instructs the driver to return an additional result
+	// set with the full spannerpb.ResultSetMetadata of the query. This result set
+	// contains one row and one column, and the value in that cell is the
+	// spannerpb.ResultSetMetadata that was returned by Spanner when executing the
+	// query. This result set will be the first result set in the sql.Rows object
+	// that is returned.
+	//
+	// You have to call [sql.Rows.NextResultSet] to move to the result set that
+	// contains the actual query data.
+	ReturnResultSetMetadata bool
+
+	// ReturnResultSetStats instructs the driver to return an additional result
+	// set with the full spannerpb.ResultSetStats of the query. This result set
+	// contains one row and one column, and the value in that cell is the
+	// spannerpb.ResultSetStats that was returned by Spanner when executing the
+	// query. This result set will be the last result set in the sql.Rows object
+	// that is returned.
+	//
+	// You have to call [sql.Rows.NextResultSet] after fetching all query data in
+	// order to move to the result set that contains the spannerpb.ResultSetStats.
+	ReturnResultSetStats bool
 }
 
 type DecodeOption int

driver_with_mockserver_test.go

@@ -4747,6 +4747,258 @@ func TestPostgreSQLDialect(t *testing.T) {
 	}
 }
 
+func TestReturnResultSetMetadata(t *testing.T) {
+	t.Parallel()
+
+	db, _, teardown := setupTestDBConnection(t)
+	defer teardown()
+	rows, err := db.QueryContext(context.Background(), testutil.SelectFooFromBar, ExecOptions{ReturnResultSetMetadata: true})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = rows.Close() }()
+
+	// Verify that the first result set contains the ResultSetMetadata.
+	if !rows.Next() {
+		t.Fatal("no rows")
+	}
+	var meta *sppb.ResultSetMetadata
+	if err := rows.Scan(&meta); err != nil {
+		t.Fatalf("failed to scan metadata: %v", err)
+	}
+	if g, w := len(meta.RowType.Fields), 1; g != w {
+		t.Fatalf("cols count mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	if rows.Next() {
+		t.Fatal("more rows than expected")
+	}
+
+	// Move to the next result set, which should contain the data.
+	if !rows.NextResultSet() {
+		t.Fatal("no more result sets found")
+	}
+
+	for want := int64(1); rows.Next(); want++ {
+		cols, err := rows.Columns()
+		if err != nil {
+			t.Fatal(err)
+		}
+		if !cmp.Equal(cols, []string{"FOO"}) {
+			t.Fatalf("cols mismatch\nGot: %v\nWant: %v", cols, []string{"FOO"})
+		}
+		var got int64
+		err = rows.Scan(&got)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got != want {
+			t.Fatalf("value mismatch\nGot: %v\nWant: %v", got, want)
+		}
+	}
+	if rows.Err() != nil {
+		t.Fatal(rows.Err())
+	}
+
+	// There should be no more result sets.
+	if rows.NextResultSet() {
+		t.Fatal("more result sets than expected")
+	}
+}
+
+func TestReturnResultSetMetadataError(t *testing.T) {
+	t.Parallel()
+
+	db, server, teardown := setupTestDBConnection(t)
+	defer teardown()
+	query := "select * from non_existing_table"
+	_ = server.TestSpanner.PutStatementResult(query, &testutil.StatementResult{
+		Type: testutil.StatementResultError,
+		Err:  gstatus.Error(codes.NotFound, "Table not found"),
+	})
+	rows, err := db.QueryContext(context.Background(), query, ExecOptions{ReturnResultSetMetadata: true})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = rows.Close() }()
+
+	if rows.Next() {
+		t.Fatal("Next should fail")
+	}
+	var meta *sppb.ResultSetMetadata
+	if err := rows.Scan(&meta); err == nil {
+		t.Fatal("missing error when scanning metadata")
+	} else {
+		if g, w := spanner.ErrCode(err), codes.NotFound; g != w {
+			t.Fatalf("error code mismatch\n Got: %v\nWant: %v", g, w)
+		}
+	}
+
+	// Moving to the next result set fails because the query failed.
+	if rows.NextResultSet() {
+		t.Fatal("got unexpected next result set")
+	}
+}
+
+func TestReturnResultSetStats(t *testing.T) {
+	t.Parallel()
+
+	db, server, teardown := setupTestDBConnection(t)
+	defer teardown()
+	query := "insert into singers (name) values ('test') then return id"
+	resultSet := testutil.CreateSingleColumnInt64ResultSet([]int64{42598}, "id")
+	_ = server.TestSpanner.PutStatementResult(query, &testutil.StatementResult{
+		Type:        testutil.StatementResultResultSet,
+		ResultSet:   resultSet,
+		UpdateCount: 1,
+	})
+
+	rows, err := db.QueryContext(context.Background(), query, ExecOptions{ReturnResultSetStats: true})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = rows.Close() }()
+
+	// The first result set should contain the data.
+	for want := int64(42598); rows.Next(); want++ {
+		cols, err := rows.Columns()
+		if err != nil {
+			t.Fatal(err)
+		}
+		if !cmp.Equal(cols, []string{"id"}) {
+			t.Fatalf("cols mismatch\nGot: %v\nWant: %v", cols, []string{"id"})
+		}
+		var got int64
+		err = rows.Scan(&got)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got != want {
+			t.Fatalf("value mismatch\nGot: %v\nWant: %v", got, want)
+		}
+	}
+	if rows.Err() != nil {
+		t.Fatal(rows.Err())
+	}
+
+	// The next result set should contain the stats.
+	if !rows.NextResultSet() {
+		t.Fatal("missing stats result set")
+	}
+
+	// Get the stats.
+	if !rows.Next() {
+		t.Fatal("no stats rows")
+	}
+	var stats *sppb.ResultSetStats
+	if err := rows.Scan(&stats); err != nil {
+		t.Fatalf("failed to scan stats: %v", err)
+	}
+	if g, w := stats.GetRowCountExact(), int64(1); g != w {
+		t.Fatalf("row count mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	if rows.Next() {
+		t.Fatal("more rows than expected")
+	}
+
+	// There should be no more result sets.
+	if rows.NextResultSet() {
+		t.Fatal("more result sets than expected")
+	}
+}
+
+func TestReturnResultSetMetadataAndStats(t *testing.T) {
+	t.Parallel()
+
+	db, server, teardown := setupTestDBConnection(t)
+	defer teardown()
+
+	query := "insert into singers (name) values ('test') then return id"
+	resultSet := testutil.CreateSingleColumnInt64ResultSet([]int64{42598}, "id")
+	_ = server.TestSpanner.PutStatementResult(query, &testutil.StatementResult{
+		Type:        testutil.StatementResultResultSet,
+		ResultSet:   resultSet,
+		UpdateCount: 1,
+	})
+
+	rows, err := db.QueryContext(context.Background(), query, ExecOptions{
+		ReturnResultSetMetadata: true,
+		ReturnResultSetStats:    true,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = rows.Close() }()
+
+	// Verify that the first result set contains the ResultSetMetadata.
+	if !rows.Next() {
+		t.Fatal("no rows")
+	}
+	var meta *sppb.ResultSetMetadata
+	if err := rows.Scan(&meta); err != nil {
+		t.Fatalf("failed to scan metadata: %v", err)
+	}
+	if g, w := len(meta.RowType.Fields), 1; g != w {
+		t.Fatalf("cols count mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	if g, w := meta.RowType.Fields[0].Name, "id"; g != w {
+		t.Fatalf("column name mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	if rows.Next() {
+		t.Fatal("more rows than expected")
+	}
+
+	// Move to the next result set, which should contain the data.
+	if !rows.NextResultSet() {
+		t.Fatal("no more result sets found")
+	}
+
+	for want := int64(42598); rows.Next(); want++ {
+		cols, err := rows.Columns()
+		if err != nil {
+			t.Fatal(err)
+		}
+		if !cmp.Equal(cols, []string{"id"}) {
+			t.Fatalf("cols mismatch\nGot: %v\nWant: %v", cols, []string{"id"})
+		}
+		var got int64
+		err = rows.Scan(&got)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got != want {
+			t.Fatalf("value mismatch\n Got: %v\nWant: %v", got, want)
+		}
+	}
+	if rows.Err() != nil {
+		t.Fatal(rows.Err())
+	}
+
+	// The next result set should contain the stats.
+	if !rows.NextResultSet() {
+		t.Fatal("missing stats result set")
+	}
+
+	// Get the stats.
+	if !rows.Next() {
+		t.Fatal("no stats rows")
+	}
+	var stats *sppb.ResultSetStats
+	if err := rows.Scan(&stats); err != nil {
+		t.Fatalf("failed to scan stats: %v", err)
+	}
+	if g, w := stats.GetRowCountExact(), int64(1); g != w {
+		t.Fatalf("row count mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	if rows.Next() {
+		t.Fatal("more rows than expected")
+	}
+
+	// There should be no more result sets.
+	if rows.NextResultSet() {
+		t.Fatal("more result sets than expected")
+	}
+}
+
 func numeric(v string) big.Rat {
 	res, _ := big.NewRat(1, 1).SetString(v)
 	return *res

merged_row_iterator.go

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