@@ -25,11 +25,13 @@ public class QueryIterable<T extends Resource> implements Iterable<T> {
2525 private boolean hasStarted = false ;
2626 private List <T > items = new ArrayList <T >();
2727 private Map <String , String > responseHeaders ;
28+ private int currentIndex = 0 ;
29+ private boolean hasNext = true ;
2830
2931 QueryIterable (DocumentClient client ,
30- DocumentServiceRequest request ,
31- ReadType readType ,
32- Class <T > classT ) {
32+ DocumentServiceRequest request ,
33+ ReadType readType ,
34+ Class <T > classT ) {
3335 this .client = client ;
3436 this .retryPolicy = new ResourceThrottleRetryPolicy (
3537 client .getRetryPolicy ().getMaxRetryAttemptsOnQuery ());
@@ -39,15 +41,6 @@ public class QueryIterable<T extends Resource> implements Iterable<T> {
3941 this .reset ();
4042 }
4143
42- private void reset () {
43- if (this .request != null && this .request .getHeaders () != null ) {
44- String continuationToken = this .request .getHeaders ().get (HttpConstants .HttpHeaders .CONTINUATION );
45- if (!StringUtils .isBlank (continuationToken )) {
46- this .continuation = continuationToken ;
47- }
48- }
49- }
50-
5144 /**
5245 * Gets the response headers.
5346 *
@@ -75,14 +68,12 @@ String getContinuation() {
7568 public Iterator <T > iterator () {
7669 Iterator <T > it = new Iterator <T >() {
7770
78- private int currentIndex = 0 ;
79- private boolean hasNext = true ;
80-
8171 private BackoffRetryUtilityDelegate delegate = new BackoffRetryUtilityDelegate () {
8272
8373 @ Override
8474 public void apply () throws Exception {
85- if (fetchNextBlock () <= 0 ) {
75+ List <T > results = fetchNextBlock ();
76+ if (results == null || results .size () <= 0 ) {
8677 hasNext = false ;
8778 }
8879 }
@@ -95,11 +86,11 @@ public void apply() throws Exception {
9586 */
9687 @ Override
9788 public boolean hasNext () {
98- if (this . currentIndex >= items .size () && this . hasNext ) {
89+ if (currentIndex >= items .size () && hasNext ) {
9990 BackoffRetryUtility .execute (this .delegate , retryPolicy );
10091 }
10192
102- return this . hasNext ;
93+ return hasNext ;
10394 }
10495
10596 /**
@@ -109,12 +100,12 @@ public boolean hasNext() {
109100 */
110101 @ Override
111102 public T next () {
112- if (this . currentIndex >= items .size () && this . hasNext ) {
103+ if (currentIndex >= items .size () && hasNext ) {
113104 BackoffRetryUtility .execute (this .delegate , retryPolicy );
114105 }
115106
116- if (!this . hasNext ) return null ;
117- return items .get (this . currentIndex ++);
107+ if (!hasNext ) return null ;
108+ return items .get (currentIndex ++);
118109 }
119110
120111 /**
@@ -123,8 +114,8 @@ public T next() {
123114 @ Override
124115 public void remove () {
125116 if (!hasNext ()) throw new NoSuchElementException ();
126- if (this . currentIndex < items .size () - 1 ) {
127- items .remove (this . currentIndex );
117+ if (currentIndex < items .size ()) {
118+ items .remove (currentIndex );
128119 }
129120 }
130121
@@ -145,16 +136,39 @@ public List<T> toList() {
145136 return list ;
146137 }
147138
148- private int fetchNextBlock ()
139+ /**
140+ * Resets the iterable.
141+ */
142+ public void reset () {
143+ this .hasStarted = false ;
144+ this .continuation = null ;
145+ this .items = new ArrayList <T >();
146+ this .currentIndex = 0 ;
147+ this .hasNext = true ;
148+ if (this .request != null && this .request .getHeaders () != null ) {
149+ String continuationToken = this .request .getHeaders ().get (HttpConstants .HttpHeaders .CONTINUATION );
150+ if (!StringUtils .isBlank (continuationToken )) {
151+ this .continuation = continuationToken ;
152+ }
153+ }
154+ }
155+
156+ /**
157+ * Fetch the next block of query results.
158+ *
159+ * @return the list of fetched resources.
160+ * @throws DocumentClientException
161+ */
162+ public List <T > fetchNextBlock ()
149163 throws DocumentClientException {
150164 DocumentServiceResponse response = null ;
151165 List <T > fetchedItems = null ;
152166
153- while (!this .isNullEmptyOrFalse (this .continuation ) ||
154- !this .hasStarted ) {
167+ while (!this .isNullEmptyOrFalse (this .continuation ) || !this .hasStarted ) {
155168 if (!this .isNullEmptyOrFalse (this .continuation )) {
156- request .getHeaders ().put (HttpConstants .HttpHeaders .CONTINUATION ,
157- this .continuation );
169+ request .getHeaders ().put (HttpConstants .HttpHeaders .CONTINUATION , this .continuation );
170+ } else {
171+ request .getHeaders ().remove (HttpConstants .HttpHeaders .CONTINUATION );
158172 }
159173
160174 if (this .readType == ReadType .Feed ) {
@@ -181,11 +195,10 @@ private int fetchNextBlock()
181195 }
182196 }
183197
184- return fetchedItems != null ? fetchedItems . size () : 0 ;
198+ return fetchedItems ;
185199 }
186200
187201 private boolean isNullEmptyOrFalse (String s ) {
188202 return s == null || s .isEmpty () || s == "false" || s == "False" ;
189203 }
190-
191204}
0 commit comments