1// Copyright 2018 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
20	vkit "cloud.google.com/go/firestore/apiv1"
21	"cloud.google.com/go/internal/trace"
22	"google.golang.org/api/iterator"
23	pb "google.golang.org/genproto/googleapis/firestore/v1"
24)
25
26// DocumentRefIterator is an iterator over DocumentRefs.
27type DocumentRefIterator struct {
28	client   *Client
29	it       *vkit.DocumentIterator
30	pageInfo *iterator.PageInfo
31	nextFunc func() error
32	items    []*DocumentRef
33	err      error
34}
35
36func newDocumentRefIterator(ctx context.Context, cr *CollectionRef, tid []byte) *DocumentRefIterator {
37	ctx = trace.StartSpan(ctx, "cloud.google.com/go/firestore.ListDocuments")
38	defer func() { trace.EndSpan(ctx, nil) }()
39
40	client := cr.c
41	req := &pb.ListDocumentsRequest{
42		Parent:       cr.parentPath,
43		CollectionId: cr.ID,
44		ShowMissing:  true,
45		Mask:         &pb.DocumentMask{}, // empty mask: we want only the ref
46	}
47	if tid != nil {
48		req.ConsistencySelector = &pb.ListDocumentsRequest_Transaction{tid}
49	}
50	it := &DocumentRefIterator{
51		client: client,
52		it:     client.c.ListDocuments(withResourceHeader(ctx, client.path()), req),
53	}
54	it.pageInfo, it.nextFunc = iterator.NewPageInfo(
55		it.fetch,
56		func() int { return len(it.items) },
57		func() interface{} { b := it.items; it.items = nil; return b })
58	return it
59}
60
61// PageInfo supports pagination. See the google.golang.org/api/iterator package for details.
62func (it *DocumentRefIterator) PageInfo() *iterator.PageInfo { return it.pageInfo }
63
64// Next returns the next result. Its second return value is iterator.Done if there
65// are no more results. Once Next returns Done, all subsequent calls will return
66// Done.
67func (it *DocumentRefIterator) Next() (*DocumentRef, error) {
68	if err := it.nextFunc(); err != nil {
69		return nil, err
70	}
71	item := it.items[0]
72	it.items = it.items[1:]
73	return item, nil
74}
75
76func (it *DocumentRefIterator) fetch(pageSize int, pageToken string) (string, error) {
77	if it.err != nil {
78		return "", it.err
79	}
80	return iterFetch(pageSize, pageToken, it.it.PageInfo(), func() error {
81		docProto, err := it.it.Next()
82		if err != nil {
83			return err
84		}
85		docRef, err := pathToDoc(docProto.Name, it.client)
86		if err != nil {
87			return err
88		}
89		it.items = append(it.items, docRef)
90		return nil
91	})
92}
93
94// GetAll returns all the DocumentRefs remaining from the iterator.
95func (it *DocumentRefIterator) GetAll() ([]*DocumentRef, error) {
96	var drs []*DocumentRef
97	for {
98		dr, err := it.Next()
99		if err == iterator.Done {
100			break
101		}
102		if err != nil {
103			return nil, err
104		}
105		drs = append(drs, dr)
106	}
107	return drs, nil
108}
109