1// Copyright 2017 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//      http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15package firestore
16
17import (
18	"context"
19	"errors"
20	"fmt"
21	"io"
22	"math"
23	"reflect"
24	"time"
25
26	"cloud.google.com/go/internal/btree"
27	"cloud.google.com/go/internal/trace"
28	"github.com/golang/protobuf/ptypes/wrappers"
29	"google.golang.org/api/iterator"
30	pb "google.golang.org/genproto/googleapis/firestore/v1"
31)
32
33// Query represents a Firestore query.
34//
35// Query values are immutable. Each Query method creates
36// a new Query; it does not modify the old.
37type Query struct {
38	c                      *Client
39	path                   string // path to query (collection)
40	parentPath             string // path of the collection's parent (document)
41	collectionID           string
42	selection              []FieldPath
43	filters                []filter
44	orders                 []order
45	offset                 int32
46	limit                  *wrappers.Int32Value
47	limitToLast            bool
48	startVals, endVals     []interface{}
49	startDoc, endDoc       *DocumentSnapshot
50	startBefore, endBefore bool
51	err                    error
52
53	// allDescendants indicates whether this query is for all collections
54	// that match the ID under the specified parentPath.
55	allDescendants bool
56}
57
58// DocumentID is the special field name representing the ID of a document
59// in queries.
60const DocumentID = "__name__"
61
62// Select returns a new Query that specifies the paths
63// to return from the result documents.
64// Each path argument can be a single field or a dot-separated sequence of
65// fields, and must not contain any of the runes "˜*/[]".
66//
67// An empty Select call will produce a query that returns only document IDs.
68func (q Query) Select(paths ...string) Query {
69	var fps []FieldPath
70	for _, s := range paths {
71		fp, err := parseDotSeparatedString(s)
72		if err != nil {
73			q.err = err
74			return q
75		}
76		fps = append(fps, fp)
77	}
78	return q.SelectPaths(fps...)
79}
80
81// SelectPaths returns a new Query that specifies the field paths
82// to return from the result documents.
83//
84// An empty SelectPaths call will produce a query that returns only document IDs.
85func (q Query) SelectPaths(fieldPaths ...FieldPath) Query {
86	if len(fieldPaths) == 0 {
87		q.selection = []FieldPath{{DocumentID}}
88	} else {
89		q.selection = fieldPaths
90	}
91	return q
92}
93
94// Where returns a new Query that filters the set of results.
95// A Query can have multiple filters.
96// The path argument can be a single field or a dot-separated sequence of
97// fields, and must not contain any of the runes "˜*/[]".
98// The op argument must be one of "==", "!=", "<", "<=", ">", ">=",
99// "array-contains", "array-contains-any", "in" or "not-in".
100func (q Query) Where(path, op string, value interface{}) Query {
101	fp, err := parseDotSeparatedString(path)
102	if err != nil {
103		q.err = err
104		return q
105	}
106	q.filters = append(append([]filter(nil), q.filters...), filter{fp, op, value})
107	return q
108}
109
110// WherePath returns a new Query that filters the set of results.
111// A Query can have multiple filters.
112// The op argument must be one of "==", "!=", "<", "<=", ">", ">=",
113// "array-contains", "array-contains-any", "in" or "not-in".
114func (q Query) WherePath(fp FieldPath, op string, value interface{}) Query {
115	q.filters = append(append([]filter(nil), q.filters...), filter{fp, op, value})
116	return q
117}
118
119// Direction is the sort direction for result ordering.
120type Direction int32
121
122const (
123	// Asc sorts results from smallest to largest.
124	Asc Direction = Direction(pb.StructuredQuery_ASCENDING)
125
126	// Desc sorts results from largest to smallest.
127	Desc Direction = Direction(pb.StructuredQuery_DESCENDING)
128)
129
130// OrderBy returns a new Query that specifies the order in which results are
131// returned. A Query can have multiple OrderBy/OrderByPath specifications.
132// OrderBy appends the specification to the list of existing ones.
133//
134// The path argument can be a single field or a dot-separated sequence of
135// fields, and must not contain any of the runes "˜*/[]".
136//
137// To order by document name, use the special field path DocumentID.
138func (q Query) OrderBy(path string, dir Direction) Query {
139	fp, err := parseDotSeparatedString(path)
140	if err != nil {
141		q.err = err
142		return q
143	}
144	q.orders = append(q.copyOrders(), order{fp, dir})
145	return q
146}
147
148// OrderByPath returns a new Query that specifies the order in which results are
149// returned. A Query can have multiple OrderBy/OrderByPath specifications.
150// OrderByPath appends the specification to the list of existing ones.
151func (q Query) OrderByPath(fp FieldPath, dir Direction) Query {
152	q.orders = append(q.copyOrders(), order{fp, dir})
153	return q
154}
155
156func (q *Query) copyOrders() []order {
157	return append([]order(nil), q.orders...)
158}
159
160// Offset returns a new Query that specifies the number of initial results to skip.
161// It must not be negative.
162func (q Query) Offset(n int) Query {
163	q.offset = trunc32(n)
164	return q
165}
166
167// Limit returns a new Query that specifies the maximum number of first results
168// to return. It must not be negative.
169func (q Query) Limit(n int) Query {
170	q.limit = &wrappers.Int32Value{Value: trunc32(n)}
171	q.limitToLast = false
172	return q
173}
174
175// LimitToLast returns a new Query that specifies the maximum number of last
176// results to return. It must not be negative.
177func (q Query) LimitToLast(n int) Query {
178	q.limit = &wrappers.Int32Value{Value: trunc32(n)}
179	q.limitToLast = true
180	return q
181}
182
183// StartAt returns a new Query that specifies that results should start at
184// the document with the given field values.
185//
186// StartAt may be called with a single DocumentSnapshot, representing an
187// existing document within the query. The document must be a direct child of
188// the location being queried (not a parent document, or document in a
189// different collection, or a grandchild document, for example).
190//
191// Otherwise, StartAt should be called with one field value for each OrderBy clause,
192// in the order that they appear. For example, in
193//   q.OrderBy("X", Asc).OrderBy("Y", Desc).StartAt(1, 2)
194// results will begin at the first document where X = 1 and Y = 2.
195//
196// If an OrderBy call uses the special DocumentID field path, the corresponding value
197// should be the document ID relative to the query's collection. For example, to
198// start at the document "NewYork" in the "States" collection, write
199//
200//   client.Collection("States").OrderBy(DocumentID, firestore.Asc).StartAt("NewYork")
201//
202// Calling StartAt overrides a previous call to StartAt or StartAfter.
203func (q Query) StartAt(docSnapshotOrFieldValues ...interface{}) Query {
204	q.startBefore = true
205	q.startVals, q.startDoc, q.err = q.processCursorArg("StartAt", docSnapshotOrFieldValues)
206	return q
207}
208
209// StartAfter returns a new Query that specifies that results should start just after
210// the document with the given field values. See Query.StartAt for more information.
211//
212// Calling StartAfter overrides a previous call to StartAt or StartAfter.
213func (q Query) StartAfter(docSnapshotOrFieldValues ...interface{}) Query {
214	q.startBefore = false
215	q.startVals, q.startDoc, q.err = q.processCursorArg("StartAfter", docSnapshotOrFieldValues)
216	return q
217}
218
219// EndAt returns a new Query that specifies that results should end at the
220// document with the given field values. See Query.StartAt for more information.
221//
222// Calling EndAt overrides a previous call to EndAt or EndBefore.
223func (q Query) EndAt(docSnapshotOrFieldValues ...interface{}) Query {
224	q.endBefore = false
225	q.endVals, q.endDoc, q.err = q.processCursorArg("EndAt", docSnapshotOrFieldValues)
226	return q
227}
228
229// EndBefore returns a new Query that specifies that results should end just before
230// the document with the given field values. See Query.StartAt for more information.
231//
232// Calling EndBefore overrides a previous call to EndAt or EndBefore.
233func (q Query) EndBefore(docSnapshotOrFieldValues ...interface{}) Query {
234	q.endBefore = true
235	q.endVals, q.endDoc, q.err = q.processCursorArg("EndBefore", docSnapshotOrFieldValues)
236	return q
237}
238
239func (q *Query) processCursorArg(name string, docSnapshotOrFieldValues []interface{}) ([]interface{}, *DocumentSnapshot, error) {
240	for _, e := range docSnapshotOrFieldValues {
241		if ds, ok := e.(*DocumentSnapshot); ok {
242			if len(docSnapshotOrFieldValues) == 1 {
243				return nil, ds, nil
244			}
245			return nil, nil, fmt.Errorf("firestore: a document snapshot must be the only argument to %s", name)
246		}
247	}
248	return docSnapshotOrFieldValues, nil, nil
249}
250
251func (q Query) query() *Query { return &q }
252
253func (q Query) toProto() (*pb.StructuredQuery, error) {
254	if q.err != nil {
255		return nil, q.err
256	}
257	if q.collectionID == "" {
258		return nil, errors.New("firestore: query created without CollectionRef")
259	}
260	if q.startBefore {
261		if len(q.startVals) == 0 && q.startDoc == nil {
262			return nil, errors.New("firestore: StartAt/StartAfter must be called with at least one value")
263		}
264	}
265	if q.endBefore {
266		if len(q.endVals) == 0 && q.endDoc == nil {
267			return nil, errors.New("firestore: EndAt/EndBefore must be called with at least one value")
268		}
269	}
270	p := &pb.StructuredQuery{
271		From: []*pb.StructuredQuery_CollectionSelector{{
272			CollectionId:   q.collectionID,
273			AllDescendants: q.allDescendants,
274		}},
275		Offset: q.offset,
276		Limit:  q.limit,
277	}
278	if len(q.selection) > 0 {
279		p.Select = &pb.StructuredQuery_Projection{}
280		for _, fp := range q.selection {
281			if err := fp.validate(); err != nil {
282				return nil, err
283			}
284			p.Select.Fields = append(p.Select.Fields, fref(fp))
285		}
286	}
287	// If there is only filter, use it directly. Otherwise, construct
288	// a CompositeFilter.
289	if len(q.filters) == 1 {
290		pf, err := q.filters[0].toProto()
291		if err != nil {
292			return nil, err
293		}
294		p.Where = pf
295	} else if len(q.filters) > 1 {
296		cf := &pb.StructuredQuery_CompositeFilter{
297			Op: pb.StructuredQuery_CompositeFilter_AND,
298		}
299		p.Where = &pb.StructuredQuery_Filter{
300			FilterType: &pb.StructuredQuery_Filter_CompositeFilter{cf},
301		}
302		for _, f := range q.filters {
303			pf, err := f.toProto()
304			if err != nil {
305				return nil, err
306			}
307			cf.Filters = append(cf.Filters, pf)
308		}
309	}
310	orders := q.orders
311	if q.startDoc != nil || q.endDoc != nil {
312		orders = q.adjustOrders()
313	}
314	for _, ord := range orders {
315		po, err := ord.toProto()
316		if err != nil {
317			return nil, err
318		}
319		p.OrderBy = append(p.OrderBy, po)
320	}
321
322	cursor, err := q.toCursor(q.startVals, q.startDoc, q.startBefore, orders)
323	if err != nil {
324		return nil, err
325	}
326	p.StartAt = cursor
327	cursor, err = q.toCursor(q.endVals, q.endDoc, q.endBefore, orders)
328	if err != nil {
329		return nil, err
330	}
331	p.EndAt = cursor
332	return p, nil
333}
334
335// If there is a start/end that uses a Document Snapshot, we may need to adjust the OrderBy
336// clauses that the user provided: we add OrderBy(__name__) if it isn't already present, and
337// we make sure we don't invalidate the original query by adding an OrderBy for inequality filters.
338func (q *Query) adjustOrders() []order {
339	// If the user is already ordering by document ID, don't change anything.
340	for _, ord := range q.orders {
341		if ord.isDocumentID() {
342			return q.orders
343		}
344	}
345	// If there are OrderBy clauses, append an OrderBy(DocumentID), using the direction of the last OrderBy clause.
346	if len(q.orders) > 0 {
347		return append(q.copyOrders(), order{
348			fieldPath: FieldPath{DocumentID},
349			dir:       q.orders[len(q.orders)-1].dir,
350		})
351	}
352	// If there are no OrderBy clauses but there is an inequality, add an OrderBy clause
353	// for the field of the first inequality.
354	var orders []order
355	for _, f := range q.filters {
356		if f.op != "==" {
357			orders = []order{{fieldPath: f.fieldPath, dir: Asc}}
358			break
359		}
360	}
361	// Add an ascending OrderBy(DocumentID).
362	return append(orders, order{fieldPath: FieldPath{DocumentID}, dir: Asc})
363}
364
365func (q *Query) toCursor(fieldValues []interface{}, ds *DocumentSnapshot, before bool, orders []order) (*pb.Cursor, error) {
366	var vals []*pb.Value
367	var err error
368	if ds != nil {
369		vals, err = q.docSnapshotToCursorValues(ds, orders)
370	} else if len(fieldValues) != 0 {
371		vals, err = q.fieldValuesToCursorValues(fieldValues)
372	} else {
373		return nil, nil
374	}
375	if err != nil {
376		return nil, err
377	}
378	return &pb.Cursor{Values: vals, Before: before}, nil
379}
380
381// toPositionValues converts the field values to protos.
382func (q *Query) fieldValuesToCursorValues(fieldValues []interface{}) ([]*pb.Value, error) {
383	if len(fieldValues) != len(q.orders) {
384		return nil, errors.New("firestore: number of field values in StartAt/StartAfter/EndAt/EndBefore does not match number of OrderBy fields")
385	}
386	vals := make([]*pb.Value, len(fieldValues))
387	var err error
388	for i, ord := range q.orders {
389		fval := fieldValues[i]
390		if ord.isDocumentID() {
391			// TODO(jba): support DocumentRefs as well as strings.
392			// TODO(jba): error if document ref does not belong to the right collection.
393			docID, ok := fval.(string)
394			if !ok {
395				return nil, fmt.Errorf("firestore: expected doc ID for DocumentID field, got %T", fval)
396			}
397			vals[i] = &pb.Value{ValueType: &pb.Value_ReferenceValue{q.path + "/" + docID}}
398		} else {
399			var sawTransform bool
400			vals[i], sawTransform, err = toProtoValue(reflect.ValueOf(fval))
401			if err != nil {
402				return nil, err
403			}
404			if sawTransform {
405				return nil, errors.New("firestore: transforms disallowed in query value")
406			}
407		}
408	}
409	return vals, nil
410}
411
412func (q *Query) docSnapshotToCursorValues(ds *DocumentSnapshot, orders []order) ([]*pb.Value, error) {
413	vals := make([]*pb.Value, len(orders))
414	for i, ord := range orders {
415		if ord.isDocumentID() {
416			dp, qp := ds.Ref.Parent.Path, q.path
417			if !q.allDescendants && dp != qp {
418				return nil, fmt.Errorf("firestore: document snapshot for %s passed to query on %s", dp, qp)
419			}
420			vals[i] = &pb.Value{ValueType: &pb.Value_ReferenceValue{ds.Ref.Path}}
421		} else {
422			val, err := valueAtPath(ord.fieldPath, ds.proto.Fields)
423			if err != nil {
424				return nil, err
425			}
426			vals[i] = val
427		}
428	}
429	return vals, nil
430}
431
432// Returns a function that compares DocumentSnapshots according to q's ordering.
433func (q Query) compareFunc() func(d1, d2 *DocumentSnapshot) (int, error) {
434	// Add implicit sorting by name, using the last specified direction.
435	lastDir := Asc
436	if len(q.orders) > 0 {
437		lastDir = q.orders[len(q.orders)-1].dir
438	}
439	orders := append(q.copyOrders(), order{[]string{DocumentID}, lastDir})
440	return func(d1, d2 *DocumentSnapshot) (int, error) {
441		for _, ord := range orders {
442			var cmp int
443			if len(ord.fieldPath) == 1 && ord.fieldPath[0] == DocumentID {
444				cmp = compareReferences(d1.Ref.Path, d2.Ref.Path)
445			} else {
446				v1, err := valueAtPath(ord.fieldPath, d1.proto.Fields)
447				if err != nil {
448					return 0, err
449				}
450				v2, err := valueAtPath(ord.fieldPath, d2.proto.Fields)
451				if err != nil {
452					return 0, err
453				}
454				cmp = compareValues(v1, v2)
455			}
456			if cmp != 0 {
457				if ord.dir == Desc {
458					cmp = -cmp
459				}
460				return cmp, nil
461			}
462		}
463		return 0, nil
464	}
465}
466
467type filter struct {
468	fieldPath FieldPath
469	op        string
470	value     interface{}
471}
472
473func (f filter) toProto() (*pb.StructuredQuery_Filter, error) {
474	if err := f.fieldPath.validate(); err != nil {
475		return nil, err
476	}
477	if uop, ok := unaryOpFor(f.value); ok {
478		if f.op != "==" {
479			return nil, fmt.Errorf("firestore: must use '==' when comparing %v", f.value)
480		}
481		return &pb.StructuredQuery_Filter{
482			FilterType: &pb.StructuredQuery_Filter_UnaryFilter{
483				UnaryFilter: &pb.StructuredQuery_UnaryFilter{
484					OperandType: &pb.StructuredQuery_UnaryFilter_Field{
485						Field: fref(f.fieldPath),
486					},
487					Op: uop,
488				},
489			},
490		}, nil
491	}
492	var op pb.StructuredQuery_FieldFilter_Operator
493	switch f.op {
494	case "<":
495		op = pb.StructuredQuery_FieldFilter_LESS_THAN
496	case "<=":
497		op = pb.StructuredQuery_FieldFilter_LESS_THAN_OR_EQUAL
498	case ">":
499		op = pb.StructuredQuery_FieldFilter_GREATER_THAN
500	case ">=":
501		op = pb.StructuredQuery_FieldFilter_GREATER_THAN_OR_EQUAL
502	case "==":
503		op = pb.StructuredQuery_FieldFilter_EQUAL
504	case "!=":
505		op = pb.StructuredQuery_FieldFilter_NOT_EQUAL
506	case "in":
507		op = pb.StructuredQuery_FieldFilter_IN
508	case "not-in":
509		op = pb.StructuredQuery_FieldFilter_NOT_IN
510	case "array-contains":
511		op = pb.StructuredQuery_FieldFilter_ARRAY_CONTAINS
512	case "array-contains-any":
513		op = pb.StructuredQuery_FieldFilter_ARRAY_CONTAINS_ANY
514	default:
515		return nil, fmt.Errorf("firestore: invalid operator %q", f.op)
516	}
517	val, sawTransform, err := toProtoValue(reflect.ValueOf(f.value))
518	if err != nil {
519		return nil, err
520	}
521	if sawTransform {
522		return nil, errors.New("firestore: transforms disallowed in query value")
523	}
524	return &pb.StructuredQuery_Filter{
525		FilterType: &pb.StructuredQuery_Filter_FieldFilter{
526			FieldFilter: &pb.StructuredQuery_FieldFilter{
527				Field: fref(f.fieldPath),
528				Op:    op,
529				Value: val,
530			},
531		},
532	}, nil
533}
534
535func unaryOpFor(value interface{}) (pb.StructuredQuery_UnaryFilter_Operator, bool) {
536	switch {
537	case value == nil:
538		return pb.StructuredQuery_UnaryFilter_IS_NULL, true
539	case isNaN(value):
540		return pb.StructuredQuery_UnaryFilter_IS_NAN, true
541	default:
542		return pb.StructuredQuery_UnaryFilter_OPERATOR_UNSPECIFIED, false
543	}
544}
545
546func isNaN(x interface{}) bool {
547	switch x := x.(type) {
548	case float32:
549		return math.IsNaN(float64(x))
550	case float64:
551		return math.IsNaN(x)
552	default:
553		return false
554	}
555}
556
557type order struct {
558	fieldPath FieldPath
559	dir       Direction
560}
561
562func (r order) isDocumentID() bool {
563	return len(r.fieldPath) == 1 && r.fieldPath[0] == DocumentID
564}
565
566func (r order) toProto() (*pb.StructuredQuery_Order, error) {
567	if err := r.fieldPath.validate(); err != nil {
568		return nil, err
569	}
570	return &pb.StructuredQuery_Order{
571		Field:     fref(r.fieldPath),
572		Direction: pb.StructuredQuery_Direction(r.dir),
573	}, nil
574}
575
576func fref(fp FieldPath) *pb.StructuredQuery_FieldReference {
577	return &pb.StructuredQuery_FieldReference{FieldPath: fp.toServiceFieldPath()}
578}
579
580func trunc32(i int) int32 {
581	if i > math.MaxInt32 {
582		i = math.MaxInt32
583	}
584	return int32(i)
585}
586
587// Documents returns an iterator over the query's resulting documents.
588func (q Query) Documents(ctx context.Context) *DocumentIterator {
589	return &DocumentIterator{
590		iter: newQueryDocumentIterator(withResourceHeader(ctx, q.c.path()), &q, nil), q: &q,
591	}
592}
593
594// DocumentIterator is an iterator over documents returned by a query.
595type DocumentIterator struct {
596	iter docIterator
597	err  error
598	q    *Query
599}
600
601// Unexported interface so we can have two different kinds of DocumentIterator: one
602// for straight queries, and one for query snapshots. We do it this way instead of
603// making DocumentIterator an interface because in the client libraries, iterators are
604// always concrete types, and the fact that this one has two different implementations
605// is an internal detail.
606type docIterator interface {
607	next() (*DocumentSnapshot, error)
608	stop()
609}
610
611// Next returns the next result. Its second return value is iterator.Done if there
612// are no more results. Once Next returns Done, all subsequent calls will return
613// Done.
614func (it *DocumentIterator) Next() (*DocumentSnapshot, error) {
615	if it.err != nil {
616		return nil, it.err
617	}
618	if it.q.limitToLast {
619		return nil, errors.New("firestore: queries that include limitToLast constraints cannot be streamed. Use DocumentIterator.GetAll() instead")
620	}
621	ds, err := it.iter.next()
622	if err != nil {
623		it.err = err
624	}
625	return ds, err
626}
627
628// Stop stops the iterator, freeing its resources.
629// Always call Stop when you are done with a DocumentIterator.
630// It is not safe to call Stop concurrently with Next.
631func (it *DocumentIterator) Stop() {
632	if it.iter != nil { // possible in error cases
633		it.iter.stop()
634	}
635	if it.err == nil {
636		it.err = iterator.Done
637	}
638}
639
640// GetAll returns all the documents remaining from the iterator.
641// It is not necessary to call Stop on the iterator after calling GetAll.
642func (it *DocumentIterator) GetAll() ([]*DocumentSnapshot, error) {
643	defer it.Stop()
644
645	q := it.q
646	limitedToLast := q.limitToLast
647	if q.limitToLast {
648		// Flip order statements before posting a request.
649		for i := range q.orders {
650			if q.orders[i].dir == Asc {
651				q.orders[i].dir = Desc
652			} else {
653				q.orders[i].dir = Asc
654			}
655		}
656		// Swap cursors.
657		q.startVals, q.endVals = q.endVals, q.startVals
658		q.startDoc, q.endDoc = q.endDoc, q.startDoc
659		q.startBefore, q.endBefore = q.endBefore, q.startBefore
660
661		q.limitToLast = false
662	}
663	var docs []*DocumentSnapshot
664	for {
665		doc, err := it.Next()
666		if err == iterator.Done {
667			break
668		}
669		if err != nil {
670			return nil, err
671		}
672		docs = append(docs, doc)
673	}
674	if limitedToLast {
675		// Flip docs order before return.
676		for i, j := 0, len(docs)-1; i < j; {
677			docs[i], docs[j] = docs[j], docs[i]
678			i++
679			j--
680		}
681	}
682	return docs, nil
683}
684
685type queryDocumentIterator struct {
686	ctx          context.Context
687	cancel       func()
688	q            *Query
689	tid          []byte // transaction ID, if any
690	streamClient pb.Firestore_RunQueryClient
691}
692
693func newQueryDocumentIterator(ctx context.Context, q *Query, tid []byte) *queryDocumentIterator {
694	ctx, cancel := context.WithCancel(ctx)
695	return &queryDocumentIterator{
696		ctx:    ctx,
697		cancel: cancel,
698		q:      q,
699		tid:    tid,
700	}
701}
702
703func (it *queryDocumentIterator) next() (_ *DocumentSnapshot, err error) {
704	it.ctx = trace.StartSpan(it.ctx, "cloud.google.com/go/firestore.Query.RunQuery")
705	defer func() { trace.EndSpan(it.ctx, err) }()
706
707	client := it.q.c
708	if it.streamClient == nil {
709		sq, err := it.q.toProto()
710		if err != nil {
711			return nil, err
712		}
713		req := &pb.RunQueryRequest{
714			Parent:    it.q.parentPath,
715			QueryType: &pb.RunQueryRequest_StructuredQuery{sq},
716		}
717		if it.tid != nil {
718			req.ConsistencySelector = &pb.RunQueryRequest_Transaction{it.tid}
719		}
720		it.streamClient, err = client.c.RunQuery(it.ctx, req)
721		if err != nil {
722			return nil, err
723		}
724	}
725	var res *pb.RunQueryResponse
726	for {
727		res, err = it.streamClient.Recv()
728		if err == io.EOF {
729			return nil, iterator.Done
730		}
731		if err != nil {
732			return nil, err
733		}
734		if res.Document != nil {
735			break
736		}
737		// No document => partial progress; keep receiving.
738	}
739	docRef, err := pathToDoc(res.Document.Name, client)
740	if err != nil {
741		return nil, err
742	}
743	doc, err := newDocumentSnapshot(docRef, res.Document, client, res.ReadTime)
744	if err != nil {
745		return nil, err
746	}
747	return doc, nil
748}
749
750func (it *queryDocumentIterator) stop() {
751	it.cancel()
752}
753
754// Snapshots returns an iterator over snapshots of the query. Each time the query
755// results change, a new snapshot will be generated.
756func (q Query) Snapshots(ctx context.Context) *QuerySnapshotIterator {
757	ws, err := newWatchStreamForQuery(ctx, q)
758	if err != nil {
759		return &QuerySnapshotIterator{err: err}
760	}
761	return &QuerySnapshotIterator{
762		Query: q,
763		ws:    ws,
764	}
765}
766
767// QuerySnapshotIterator is an iterator over snapshots of a query.
768// Call Next on the iterator to get a snapshot of the query's results each time they change.
769// Call Stop on the iterator when done.
770//
771// For an example, see Query.Snapshots.
772type QuerySnapshotIterator struct {
773	// The Query used to construct this iterator.
774	Query Query
775
776	ws  *watchStream
777	err error
778}
779
780// Next blocks until the query's results change, then returns a QuerySnapshot for
781// the current results.
782//
783// Next is not expected to return iterator.Done unless it is called after Stop.
784// Rarely, networking issues may also cause iterator.Done to be returned.
785func (it *QuerySnapshotIterator) Next() (*QuerySnapshot, error) {
786	if it.err != nil {
787		return nil, it.err
788	}
789	btree, changes, readTime, err := it.ws.nextSnapshot()
790	if err != nil {
791		if err == io.EOF {
792			err = iterator.Done
793		}
794		it.err = err
795		return nil, it.err
796	}
797	return &QuerySnapshot{
798		Documents: &DocumentIterator{
799			iter: (*btreeDocumentIterator)(btree.BeforeIndex(0)), q: &it.Query,
800		},
801		Size:     btree.Len(),
802		Changes:  changes,
803		ReadTime: readTime,
804	}, nil
805}
806
807// Stop stops receiving snapshots. You should always call Stop when you are done with
808// a QuerySnapshotIterator, to free up resources. It is not safe to call Stop
809// concurrently with Next.
810func (it *QuerySnapshotIterator) Stop() {
811	if it.ws != nil {
812		it.ws.stop()
813	}
814}
815
816// A QuerySnapshot is a snapshot of query results. It is returned by
817// QuerySnapshotIterator.Next whenever the results of a query change.
818type QuerySnapshot struct {
819	// An iterator over the query results.
820	// It is not necessary to call Stop on this iterator.
821	Documents *DocumentIterator
822
823	// The number of results in this snapshot.
824	Size int
825
826	// The changes since the previous snapshot.
827	Changes []DocumentChange
828
829	// The time at which this snapshot was obtained from Firestore.
830	ReadTime time.Time
831}
832
833type btreeDocumentIterator btree.Iterator
834
835func (it *btreeDocumentIterator) next() (*DocumentSnapshot, error) {
836	if !(*btree.Iterator)(it).Next() {
837		return nil, iterator.Done
838	}
839	return it.Key.(*DocumentSnapshot), nil
840}
841
842func (*btreeDocumentIterator) stop() {}
843