Adopt spaniter v0.1.0 row iterator adapters

Adopt `github.com/apstndb/spaniter` for Cloud Spanner row iterator lifecycle handling and update the dependency to `v0.1.0`.

- Replace local row-iterator sequencing/draining paths with spaniter helpers.
- Preserve metadata, query plan, query stats, and row-count assembly through `BuildResultSetFromParts`.
- Add eager `PROFILE` + redacted rows coverage so stats remain available even when row values are omitted.

Validation:
- `go test ./...`
- `make test`
- `git diff --check`
- GitHub checks `build` and `lint`

Author: apstndb

go.mod

@@ -9,6 +9,7 @@ require (
 	github.com/apstndb/gsqlsep v0.0.0-20240823174243-432be37d515a
 	github.com/apstndb/memebridge v0.6.0
 	github.com/apstndb/spanemuboost v0.0.0-20241212215948-d4eb945c8f10
+	github.com/apstndb/spaniter v0.1.0
 	github.com/apstndb/spannerotel v0.0.0-20220113012943-45b1c9cc5afb
 	github.com/apstndb/spanvalue v0.8.0
 	github.com/cloudspannerecosystem/memefish v0.6.2

go.sum

@@ -670,6 +670,8 @@ github.com/apstndb/memebridge v0.6.0 h1:45uu/RFF63u5vJ+d8HT4PvOfotm1cV4eEveWFwGv
 github.com/apstndb/memebridge v0.6.0/go.mod h1:E/HVP4iaSgiPXMIQTPT1u1uKkjmRtRECkVKE2TrF3bc=
 github.com/apstndb/spanemuboost v0.0.0-20241212215948-d4eb945c8f10 h1:Ea5al1Su3ZNw2HltOc1fjLADUp16jV9ns03D9Xvpep4=
 github.com/apstndb/spanemuboost v0.0.0-20241212215948-d4eb945c8f10/go.mod h1:JtIpejiunpnBufWCwhEun+ayfd0JWnPp9kqbcNJkU+E=
+github.com/apstndb/spaniter v0.1.0 h1:OpTpePq0bN/EWS1mIXAk0lfQ98PAI4xRcBQ4pd9Ijk0=
+github.com/apstndb/spaniter v0.1.0/go.mod h1:vl6zVOnA2m35fAeF1WsGD7AuaWUK5aCVG6aS/35njFE=
 github.com/apstndb/spannerotel v0.0.0-20220113012943-45b1c9cc5afb h1:wop6M8ZQWhVLYtiBKXvkNd6TFUSpiGya86Az57KiTl0=
 github.com/apstndb/spannerotel v0.0.0-20220113012943-45b1c9cc5afb/go.mod h1:H9relGptAtWqSNrJy+V1jiYPeBIk8KYh6lACV3VZrqU=
 github.com/apstndb/spantype v0.3.11 h1:wKue4WLYGT82MH3B3TRSFn8tWIbk1Geczs+5BAEtnK4=

integration_test.go

@@ -3,15 +3,16 @@ package main
 import (
 	"bytes"
 	"context"
-	"encoding/csv"
 	_ "embed"
+	"encoding/csv"
 	"strings"
 	"testing"
 
 	"cloud.google.com/go/spanner"
 	sppb "cloud.google.com/go/spanner/apiv1/spannerpb"
 	"github.com/apstndb/gsqlsep"
 	"github.com/apstndb/spanemuboost"
+	"github.com/apstndb/spaniter"
 	"github.com/google/go-cmp/cmp"
 	"google.golang.org/protobuf/testing/protocmp"
 	"google.golang.org/protobuf/types/known/structpb"
@@ -84,14 +85,13 @@ func TestWithCloudSpannerEmulator(t *testing.T) {
 		rowIter := clients.Client.Single().QueryWithOptions(ctx,
 			spanner.Statement{SQL: "SELECT @i64, @f32, @f64, @s, @bs, @bl, @dt, @ts, @n, @a_str, @a_s, @j", Params: params},
 			spanner.QueryOptions{Mode: sppb.ExecuteSqlRequest_PLAN.Enum()})
-		defer rowIter.Stop()
 
-		err = skipRowIter(rowIter)
+		result, err := spaniter.DrainRowIterator(rowIter)
 		if err != nil {
 			t.Fatal(err)
 		}
 
-		got := rowIter.Metadata.GetRowType()
+		got := result.Metadata.GetRowType()
 		if diff := cmp.Diff(want, got, protocmp.Transform()); diff != "" {
 			t.Error(diff)
 			return
@@ -258,9 +258,8 @@ func TestWithCloudSpannerEmulator(t *testing.T) {
 				rowIter := client.Single().QueryWithOptions(ctx,
 					spanner.Statement{SQL: "SELECT @v AS v", Params: params},
 					spanner.QueryOptions{Mode: sppb.ExecuteSqlRequest_NORMAL.Enum()})
-				defer rowIter.Stop()
 
-				rows, err := xiter.TryCollect(rowIterSeq(rowIter))
+				rows, err := xiter.TryCollect(spaniter.RowIteratorSeq(rowIter))
 				if err != nil {
 					t.Error(err)
 					return
@@ -407,6 +406,31 @@ func TestWithCloudSpannerEmulator(t *testing.T) {
 		}
 	})
 
+	t.Run("eager PROFILE redact rows preserves stats", func(t *testing.T) {
+		rowIter := client.Single().QueryWithOptions(ctx,
+			spanner.Statement{SQL: "SELECT SingerId FROM Singers ORDER BY SingerId LIMIT 3"},
+			spanner.QueryOptions{Mode: sppb.ExecuteSqlRequest_PROFILE.Enum()},
+		)
+
+		rs, err := consumeRowIterIntoResultSet(rowIter, true)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if len(rs.GetRows()) != 0 {
+			t.Fatalf("rows: got %d, want 0", len(rs.GetRows()))
+		}
+		fields := rs.GetMetadata().GetRowType().GetFields()
+		if len(fields) != 1 || fields[0].GetName() != "SingerId" {
+			t.Fatalf("metadata fields: got %v", fields)
+		}
+		if rs.GetStats() == nil {
+			t.Fatal("stats: got nil, want PROFILE stats")
+		}
+		if rs.GetStats().GetQueryPlan() == nil && rs.GetStats().GetQueryStats() == nil {
+			t.Fatal("stats: got no query plan or query stats, want PROFILE stats")
+		}
+	})
+
 	t.Run("experimental_csv writeCsvFromRowIter", func(t *testing.T) {
 		readCSV := func(t *testing.T, out string) [][]string {
 			t.Helper()

jqresult/protojson.go

@@ -61,30 +61,37 @@ func ResultSetMap(rs *sppb.ResultSet) (map[string]any, error) {
 
 // BuildResultSet constructs ResultSet stats from a consumed row iterator.
 func BuildResultSet(rows []*structpb.ListValue, rowIter *spanner.RowIterator) (*sppb.ResultSet, error) {
+	return BuildResultSetFromParts(rows, rowIter.Metadata, rowIter.QueryPlan, rowIter.QueryStats, rowIter.RowCount)
+}
+
+// BuildResultSetFromParts constructs a ResultSet from materialized rows and
+// consumed iterator metadata. rows may be nil when row values are intentionally
+// redacted; metadata and stats are still preserved from the drained iterator.
+func BuildResultSetFromParts(rows []*structpb.ListValue, metadata *sppb.ResultSetMetadata, queryPlan *sppb.QueryPlan, queryStatsMap map[string]any, rowCount int64) (*sppb.ResultSet, error) {
 	out := &sppb.ResultSet{
 		Rows:     rows,
-		Metadata: rowIter.Metadata,
+		Metadata: metadata,
 	}
 
 	var queryStats *structpb.Struct
-	if rowIter.QueryStats != nil {
-		qs, err := structpb.NewStruct(rowIter.QueryStats)
+	if queryStatsMap != nil {
+		qs, err := structpb.NewStruct(queryStatsMap)
 		if err != nil {
 			return nil, err
 		}
 		queryStats = qs
 	}
 
-	if rowIter.QueryPlan == nil && queryStats == nil && rowIter.RowCount == 0 {
+	if queryPlan == nil && queryStats == nil && rowCount == 0 {
 		return out, nil
 	}
 
 	out.Stats = &sppb.ResultSetStats{
-		QueryPlan:  rowIter.QueryPlan,
+		QueryPlan:  queryPlan,
 		QueryStats: queryStats,
 	}
-	if rowIter.RowCount != 0 {
-		out.Stats.RowCount = &sppb.ResultSetStats_RowCountExact{RowCountExact: rowIter.RowCount}
+	if rowCount != 0 {
+		out.Stats.RowCount = &sppb.ResultSetStats_RowCountExact{RowCountExact: rowCount}
 	}
 	return out, nil
 }

main.go

@@ -6,9 +6,7 @@ import (
 	"errors"
 	"github.com/apstndb/execspansql/internal"
 	"github.com/apstndb/execspansql/params"
-	"google.golang.org/api/iterator"
 	"io"
-	"iter"
 	"regexp"
 	"slices"
 	"spheric.cloud/xiter"
@@ -36,9 +34,10 @@ import (
 
 	"cloud.google.com/go/spanner"
 	sppb "cloud.google.com/go/spanner/apiv1/spannerpb"
+	"github.com/apstndb/execspansql/jqresult"
+	"github.com/apstndb/spaniter"
 	"github.com/apstndb/spannerotel/interceptor"
 	svwriter "github.com/apstndb/spanvalue/writer"
-	"github.com/apstndb/execspansql/jqresult"
 	"github.com/jessevdk/go-flags"
 	"github.com/wader/gojq"
 	"google.golang.org/protobuf/types/known/structpb"
@@ -568,56 +567,28 @@ func newEncoder(writer io.Writer, format string, compactOutput bool, rawOutput b
 // consumeRowIterIntoResultSet construct *spannerpb.ResultSet using *spanner.RowIterator.
 // rowIter must be passed without calling any method, and it will be closed by this function.
 func consumeRowIterIntoResultSet(rowIter *spanner.RowIterator, redactRows bool) (*sppb.ResultSet, error) {
-	defer rowIter.Stop()
-
-	var rows []*structpb.ListValue
-	var err error
-
 	if redactRows {
-		err = skipRowIter(rowIter)
-	} else {
-		rows, err = consumeRowIterIntoListValues(rowIter)
+		result, err := spaniter.DrainRowIterator(rowIter)
+		if err != nil {
+			return nil, err
+		}
+		return jqresult.BuildResultSetFromParts(nil, result.Metadata, result.Stats.QueryPlan, result.Stats.QueryStats, result.Stats.RowCount)
 	}
 
+	var result spaniter.RowIteratorResult
+	rows, err := consumeRowIterIntoListValues(rowIter,
+		spaniter.WithResult(&result),
+	)
 	if err != nil {
 		return nil, err
 	}
-
-	return convertToResultSet(rows, rowIter)
-}
-
-// convertToResultSet convert rows and rowIter into sppb.ResultSet.
-// rowIter must be consumed with iterator.Done, and not Stop()-ed.
-func convertToResultSet(rows []*structpb.ListValue, rowIter *spanner.RowIterator) (*sppb.ResultSet, error) {
-	return jqresult.BuildResultSet(rows, rowIter)
-}
-
-func rowIterSeq(rowIter *spanner.RowIterator) iter.Seq2[*spanner.Row, error] {
-	return func(yield func(*spanner.Row, error) bool) {
-		for {
-			r, err := rowIter.Next()
-			if errors.Is(err, iterator.Done) {
-				return
-			}
-			if err != nil {
-				_ = yield(nil, err)
-				return
-			}
-			if !yield(r, nil) {
-				return
-			}
-		}
-	}
+	return jqresult.BuildResultSetFromParts(rows, result.Metadata, result.Stats.QueryPlan, result.Stats.QueryStats, result.Stats.RowCount)
 }
 
 func rowToListValue(r *spanner.Row) *structpb.ListValue {
 	return &structpb.ListValue{Values: slices.Collect(xiter.Map(xiter.Range(0, r.Size()), r.ColumnValue))}
 }
 
-func consumeRowIterIntoListValues(rowIter *spanner.RowIterator) ([]*structpb.ListValue, error) {
-	return xiter.TryCollect(internal.MapNonError(rowIterSeq(rowIter), rowToListValue))
-}
-
-func skipRowIter(rowIter *spanner.RowIterator) error {
-	return rowIter.Do(func(r *spanner.Row) error { return nil })
+func consumeRowIterIntoListValues(rowIter *spanner.RowIterator, opts ...spaniter.Option) ([]*structpb.ListValue, error) {
+	return xiter.TryCollect(internal.MapNonError(spaniter.RowIteratorSeq(rowIter, opts...), rowToListValue))
 }