feat: add option to return metadata and stats (#456)

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

@@ -798,7 +798,13 @@ func (c *conn) queryContext(ctx context.Context, query string, execOptions ExecO
 			return nil, err
 		}
 	}
-	res := &rows{it: iter, decodeOption: execOptions.DecodeOption, decodeToNativeArrays: execOptions.DecodeToNativeArrays}
+	res := &rows{
+		it:                      iter,
+		decodeOption:            execOptions.DecodeOption,
+		decodeToNativeArrays:    execOptions.DecodeToNativeArrays,
+		returnResultSetMetadata: execOptions.ReturnResultSetMetadata,
+		returnResultSetStats:    execOptions.ReturnResultSetStats,
+	}
 	if execOptions.DirectExecuteQuery {
 		// This call to res.getColumns() triggers the execution of the statement, as it needs to fetch the metadata.
 		res.getColumns()

driver.go

@@ -176,6 +176,28 @@ type ExecOptions struct {
 	// 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
+
 	// DirectExecute determines whether a query is executed directly when the
 	// [sql.DB.QueryContext] method is called, or whether the actual query execution
 	// is delayed until the first call to [sql.Rows.Next]. The default is to delay

driver_with_mockserver_test.go

@@ -4995,6 +4995,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{}
+}