iterator.go 8.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231
  1. // Copyright 2016 Google Inc. All Rights Reserved.
  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. // Package iterator provides support for standard Google API iterators.
  15. // See https://github.com/GoogleCloudPlatform/gcloud-golang/wiki/Iterator-Guidelines.
  16. package iterator
  17. import (
  18. "errors"
  19. "fmt"
  20. "reflect"
  21. )
  22. // Done is returned by an iterator's Next method when the iteration is
  23. // complete; when there are no more items to return.
  24. var Done = errors.New("no more items in iterator")
  25. // We don't support mixed calls to Next and NextPage because they play
  26. // with the paging state in incompatible ways.
  27. var errMixed = errors.New("iterator: Next and NextPage called on same iterator")
  28. // PageInfo contains information about an iterator's paging state.
  29. type PageInfo struct {
  30. // Token is the token used to retrieve the next page of items from the
  31. // API. You may set Token immediately after creating an iterator to
  32. // begin iteration at a particular point. If Token is the empty string,
  33. // the iterator will begin with the first eligible item.
  34. //
  35. // The result of setting Token after the first call to Next is undefined.
  36. //
  37. // After the underlying API method is called to retrieve a page of items,
  38. // Token is set to the next-page token in the response.
  39. Token string
  40. // MaxSize is the maximum number of items returned by a call to the API.
  41. // Set MaxSize as a hint to optimize the buffering behavior of the iterator.
  42. // If zero, the page size is determined by the underlying service.
  43. //
  44. // Use Pager to retrieve a page of a specific, exact size.
  45. MaxSize int
  46. // The error state of the iterator. Manipulated by PageInfo.next and Pager.
  47. // This is a latch: it starts as nil, and once set should never change.
  48. err error
  49. // If true, no more calls to fetch should be made. Set to true when fetch
  50. // returns an empty page token. The iterator is Done when this is true AND
  51. // the buffer is empty.
  52. atEnd bool
  53. // Function that fetches a page from the underlying service. It should pass
  54. // the pageSize and pageToken arguments to the service, fill the buffer
  55. // with the results from the call, and return the next-page token returned
  56. // by the service. The function must not remove any existing items from the
  57. // buffer. If the underlying RPC takes an int32 page size, pageSize should
  58. // be silently truncated.
  59. fetch func(pageSize int, pageToken string) (nextPageToken string, err error)
  60. // Function that returns the number of currently buffered items.
  61. bufLen func() int
  62. // Function that returns the buffer, after setting the buffer variable to nil.
  63. takeBuf func() interface{}
  64. // Set to true on first call to PageInfo.next or Pager.NextPage. Used to check
  65. // for calls to both Next and NextPage with the same iterator.
  66. nextCalled, nextPageCalled bool
  67. }
  68. // NewPageInfo exposes internals for iterator implementations.
  69. // It is not a stable interface.
  70. var NewPageInfo = newPageInfo
  71. // If an iterator can support paging, its iterator-creating method should call
  72. // this (via the NewPageInfo variable above).
  73. //
  74. // The fetch, bufLen and takeBuf arguments provide access to the
  75. // iterator's internal slice of buffered items. They behave as described in
  76. // PageInfo, above.
  77. //
  78. // The return value is the PageInfo.next method bound to the returned PageInfo value.
  79. // (Returning it avoids exporting PageInfo.next.)
  80. func newPageInfo(fetch func(int, string) (string, error), bufLen func() int, takeBuf func() interface{}) (*PageInfo, func() error) {
  81. pi := &PageInfo{
  82. fetch: fetch,
  83. bufLen: bufLen,
  84. takeBuf: takeBuf,
  85. }
  86. return pi, pi.next
  87. }
  88. // Remaining returns the number of items available before the iterator makes another API call.
  89. func (pi *PageInfo) Remaining() int { return pi.bufLen() }
  90. // next provides support for an iterator's Next function. An iterator's Next
  91. // should return the error returned by next if non-nil; else it can assume
  92. // there is at least one item in its buffer, and it should return that item and
  93. // remove it from the buffer.
  94. func (pi *PageInfo) next() error {
  95. pi.nextCalled = true
  96. if pi.err != nil { // Once we get an error, always return it.
  97. // TODO(jba): fix so users can retry on transient errors? Probably not worth it.
  98. return pi.err
  99. }
  100. if pi.nextPageCalled {
  101. pi.err = errMixed
  102. return pi.err
  103. }
  104. // Loop until we get some items or reach the end.
  105. for pi.bufLen() == 0 && !pi.atEnd {
  106. if err := pi.fill(pi.MaxSize); err != nil {
  107. pi.err = err
  108. return pi.err
  109. }
  110. if pi.Token == "" {
  111. pi.atEnd = true
  112. }
  113. }
  114. // Either the buffer is non-empty or pi.atEnd is true (or both).
  115. if pi.bufLen() == 0 {
  116. // The buffer is empty and pi.atEnd is true, i.e. the service has no
  117. // more items.
  118. pi.err = Done
  119. }
  120. return pi.err
  121. }
  122. // Call the service to fill the buffer, using size and pi.Token. Set pi.Token to the
  123. // next-page token returned by the call.
  124. // If fill returns a non-nil error, the buffer will be empty.
  125. func (pi *PageInfo) fill(size int) error {
  126. tok, err := pi.fetch(size, pi.Token)
  127. if err != nil {
  128. pi.takeBuf() // clear the buffer
  129. return err
  130. }
  131. pi.Token = tok
  132. return nil
  133. }
  134. // Pageable is implemented by iterators that support paging.
  135. type Pageable interface {
  136. // PageInfo returns paging information associated with the iterator.
  137. PageInfo() *PageInfo
  138. }
  139. // Pager supports retrieving iterator items a page at a time.
  140. type Pager struct {
  141. pageInfo *PageInfo
  142. pageSize int
  143. }
  144. // NewPager returns a pager that uses iter. Calls to its NextPage method will
  145. // obtain exactly pageSize items, unless fewer remain. The pageToken argument
  146. // indicates where to start the iteration. Pass the empty string to start at
  147. // the beginning, or pass a token retrieved from a call to Pager.NextPage.
  148. //
  149. // If you use an iterator with a Pager, you must not call Next on the iterator.
  150. func NewPager(iter Pageable, pageSize int, pageToken string) *Pager {
  151. p := &Pager{
  152. pageInfo: iter.PageInfo(),
  153. pageSize: pageSize,
  154. }
  155. p.pageInfo.Token = pageToken
  156. if pageSize <= 0 {
  157. p.pageInfo.err = errors.New("iterator: page size must be positive")
  158. }
  159. return p
  160. }
  161. // NextPage retrieves a sequence of items from the iterator and appends them
  162. // to slicep, which must be a pointer to a slice of the iterator's item type.
  163. // Exactly p.pageSize items will be appended, unless fewer remain.
  164. //
  165. // The first return value is the page token to use for the next page of items.
  166. // If empty, there are no more pages. Aside from checking for the end of the
  167. // iteration, the returned page token is only needed if the iteration is to be
  168. // resumed a later time, in another context (possibly another process).
  169. //
  170. // The second return value is non-nil if an error occurred. It will never be
  171. // the special iterator sentinel value Done. To recognize the end of the
  172. // iteration, compare nextPageToken to the empty string.
  173. //
  174. // It is possible for NextPage to return a single zero-length page along with
  175. // an empty page token when there are no more items in the iteration.
  176. func (p *Pager) NextPage(slicep interface{}) (nextPageToken string, err error) {
  177. p.pageInfo.nextPageCalled = true
  178. if p.pageInfo.err != nil {
  179. return "", p.pageInfo.err
  180. }
  181. if p.pageInfo.nextCalled {
  182. p.pageInfo.err = errMixed
  183. return "", p.pageInfo.err
  184. }
  185. if p.pageInfo.bufLen() > 0 {
  186. return "", errors.New("must call NextPage with an empty buffer")
  187. }
  188. // The buffer must be empty here, so takeBuf is a no-op. We call it just to get
  189. // the buffer's type.
  190. wantSliceType := reflect.PtrTo(reflect.ValueOf(p.pageInfo.takeBuf()).Type())
  191. if slicep == nil {
  192. return "", errors.New("nil passed to Pager.NextPage")
  193. }
  194. vslicep := reflect.ValueOf(slicep)
  195. if vslicep.Type() != wantSliceType {
  196. return "", fmt.Errorf("slicep should be of type %s, got %T", wantSliceType, slicep)
  197. }
  198. for p.pageInfo.bufLen() < p.pageSize {
  199. if err := p.pageInfo.fill(p.pageSize - p.pageInfo.bufLen()); err != nil {
  200. p.pageInfo.err = err
  201. return "", p.pageInfo.err
  202. }
  203. if p.pageInfo.Token == "" {
  204. break
  205. }
  206. }
  207. e := vslicep.Elem()
  208. e.Set(reflect.AppendSlice(e, reflect.ValueOf(p.pageInfo.takeBuf())))
  209. return p.pageInfo.Token, nil
  210. }