spanner/spannertest: implement LIMIT as a rowIter
This further simplifies the implementation of the query executor.
Change-Id: I323f00bc168616530f1021afd3865fbd5a6962f3
Reviewed-on: https://code-review.googlesource.com/c/gocloud/+/52732
Reviewed-by: kokoro <noreply+kokoro@google.com>
Reviewed-by: Shanika Kuruppu <skuruppu@google.com>
diff --git a/spanner/spannertest/db_query.go b/spanner/spannertest/db_query.go
index 3ca4537..4e0029d 100644
--- a/spanner/spannertest/db_query.go
+++ b/spanner/spannertest/db_query.go
@@ -208,6 +208,25 @@
}
}
+// limitIter applies a LIMIT clause.
+type limitIter struct {
+ ri rowIter
+ rem int64
+}
+
+func (li *limitIter) Cols() []colInfo { return li.ri.Cols() }
+func (li *limitIter) Next() (row, error) {
+ if li.rem <= 0 {
+ return nil, io.EOF
+ }
+ row, err := li.ri.Next()
+ if err != nil {
+ return nil, err
+ }
+ li.rem--
+ return row, nil
+}
+
type queryParams map[string]interface{}
func (d *database) Query(q spansql.Query, params queryParams) (rowIter, error) {
@@ -258,20 +277,12 @@
// TODO: OFFSET
// Apply LIMIT.
- // TODO: this can be an iter too.
if q.Limit != nil {
lim, err := evalLimit(q.Limit, params)
if err != nil {
return nil, err
}
- raw, err := toRawIter(ri)
- if err != nil {
- return nil, err
- }
- if n := int(lim); n < len(raw.rows) {
- raw.rows = raw.rows[:n]
- }
- ri = raw
+ ri = &limitIter{ri: ri, rem: lim}
}
return ri, nil
