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