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