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