diff --git a/firestore/integration_test.go b/firestore/integration_test.go index 16f832150c33..123465fa1198 100644 --- a/firestore/integration_test.go +++ b/firestore/integration_test.go @@ -632,7 +632,7 @@ func TestIntegration_WriteBatch(t *testing.T) { // TODO(jba): test verify when it is supported. } -func TestIntegration_Query(t *testing.T) { +func TestIntegration_QueryDocuments(t *testing.T) { ctx := context.Background() coll := integrationColl(t) h := testHelper{t} @@ -657,6 +657,7 @@ func TestIntegration_Query(t *testing.T) { {q.StartAfter(1), wants[2:]}, {q.EndAt(1), wants[:2]}, {q.EndBefore(1), wants[:1]}, + {q.LimitToLast(2), wants[1:]}, } { gotDocs, err := test.q.Documents(ctx).GetAll() if err != nil { @@ -681,7 +682,7 @@ func TestIntegration_Query(t *testing.T) { if err != nil { t.Fatal(err) } - seen := map[int64]bool{} // "q" values we see + seen := map[int64]bool{} // "q" values we see. for _, d := range allDocs { data := d.Data() q, ok := data["q"] @@ -705,6 +706,16 @@ func TestIntegration_Query(t *testing.T) { } } +func TestIntegration_QueryDocuments_LimitToLast_Fail(t *testing.T) { + ctx := context.Background() + coll := integrationColl(t) + q := coll.Select("q").OrderBy("q", Asc).LimitToLast(1) + got, err := q.Documents(ctx).Next() + if err == nil { + t.Errorf("got %v doc, want error", got) + } +} + // Test unary filters. func TestIntegration_QueryUnary(t *testing.T) { ctx := context.Background() diff --git a/firestore/query.go b/firestore/query.go index 41a7231fecdb..03b32b7f9215 100644 --- a/firestore/query.go +++ b/firestore/query.go @@ -43,6 +43,7 @@ type Query struct { orders []order offset int32 limit *wrappers.Int32Value + limitToLast bool startVals, endVals []interface{} startDoc, endDoc *DocumentSnapshot startBefore, endBefore bool @@ -162,10 +163,19 @@ func (q Query) Offset(n int) Query { return q } -// Limit returns a new Query that specifies the maximum number of results to return. -// It must not be negative. +// Limit returns a new Query that specifies the maximum number of first results +// to return. It must not be negative. func (q Query) Limit(n int) Query { q.limit = &wrappers.Int32Value{Value: trunc32(n)} + q.limitToLast = false + return q +} + +// LimitToLast returns a new Query that specifies the maximum number of last +// results to return. It must not be negative. +func (q Query) LimitToLast(n int) Query { + q.limit = &wrappers.Int32Value{Value: trunc32(n)} + q.limitToLast = true return q } @@ -572,7 +582,7 @@ func trunc32(i int) int32 { // Documents returns an iterator over the query's resulting documents. func (q Query) Documents(ctx context.Context) *DocumentIterator { return &DocumentIterator{ - iter: newQueryDocumentIterator(withResourceHeader(ctx, q.c.path()), &q, nil), + iter: newQueryDocumentIterator(withResourceHeader(ctx, q.c.path()), &q, nil), q: &q, } } @@ -580,6 +590,7 @@ func (q Query) Documents(ctx context.Context) *DocumentIterator { type DocumentIterator struct { iter docIterator err error + q *Query } // Unexported interface so we can have two different kinds of DocumentIterator: one @@ -599,6 +610,9 @@ func (it *DocumentIterator) Next() (*DocumentSnapshot, error) { if it.err != nil { return nil, it.err } + if it.q.limitToLast { + return nil, errors.New("firestore: queries that include limitToLast constraints cannot be streamed. Use DocumentIterator.GetAll() instead") + } ds, err := it.iter.next() if err != nil { it.err = err @@ -622,6 +636,25 @@ func (it *DocumentIterator) Stop() { // It is not necessary to call Stop on the iterator after calling GetAll. func (it *DocumentIterator) GetAll() ([]*DocumentSnapshot, error) { defer it.Stop() + + q := it.q + limitedToLast := q.limitToLast + if q.limitToLast { + // Flip order statements before posting a request. + for i := range q.orders { + if q.orders[i].dir == Asc { + q.orders[i].dir = Desc + } else { + q.orders[i].dir = Asc + } + } + // Swap cursors. + q.startVals, q.endVals = q.endVals, q.startVals + q.startDoc, q.endDoc = q.endDoc, q.startDoc + q.startBefore, q.endBefore = q.endBefore, q.startBefore + + q.limitToLast = false + } var docs []*DocumentSnapshot for { doc, err := it.Next() @@ -633,6 +666,14 @@ func (it *DocumentIterator) GetAll() ([]*DocumentSnapshot, error) { } docs = append(docs, doc) } + if limitedToLast { + // Flip docs order before return. + for i, j := 0, len(docs)-1; i < j; { + docs[i], docs[j] = docs[j], docs[i] + i++ + j-- + } + } return docs, nil } @@ -747,7 +788,7 @@ func (it *QuerySnapshotIterator) Next() (*QuerySnapshot, error) { } return &QuerySnapshot{ Documents: &DocumentIterator{ - iter: (*btreeDocumentIterator)(btree.BeforeIndex(0)), + iter: (*btreeDocumentIterator)(btree.BeforeIndex(0)), q: &it.Query, }, Size: btree.Len(), Changes: changes, diff --git a/firestore/transaction.go b/firestore/transaction.go index e3ac314b547b..392168c718d3 100644 --- a/firestore/transaction.go +++ b/firestore/transaction.go @@ -228,8 +228,9 @@ func (t *Transaction) Documents(q Queryer) *DocumentIterator { t.readAfterWrite = true return &DocumentIterator{err: errReadAfterWrite} } + query := q.query() return &DocumentIterator{ - iter: newQueryDocumentIterator(t.ctx, q.query(), t.id), + iter: newQueryDocumentIterator(t.ctx, query, t.id), q: query, } }