googleapis
diff --git a/‎src/main/java/com/google/gcloud/datastore/QueryResult.java‎
Lines changed: 7 additions & 13 deletions b/‎src/main/java/com/google/gcloud/datastore/QueryResult.java‎
Lines changed: 7 additions & 13 deletions
diff --git a/‎src/main/java/com/google/gcloud/datastore/QueryResultImpl.java‎
Lines changed: 11 additions & 27 deletions b/‎src/main/java/com/google/gcloud/datastore/QueryResultImpl.java‎
Lines changed: 11 additions & 27 deletions
@@ -4,26 +4,20 @@
 
 /**
  * The result of a Google Cloud Datastore query submission.
- * When the result is not typed it is possible to cast it to its appropriate type according to
- * the {@link #getType} result.
+ * When result is not typed it is possible to cast it to its appropriate type according to
+ * the {@link #resultClass} value.
  *
- * @param V the type of values the result holds.
+ * @param V the type of the results value.
  */
 public interface QueryResult<T> extends Iterator<T> {
 
   /**
-   * Returns the actual type of the result's values.
-   * When needed the result could be casted accordingly:
-   * <pre> {@code
-   * Type.FULL -> (QueryResult<Entity>)
-   * Type.PROJECTION -> (QueryResult<PartialEntity>)
-   * Type.KEY_ONLY -> (QueryResult<Key>)
-   * } </pre>
+   * Returns the actual class of the result's values.
    */
-  Query.Type getType();
+  Class<?> resultClass();
 
   /**
-   * Return the Cursor for the next result. Not currently implemented (depends on v1beta3).
+   * Returns the Cursor for the next result. Not currently implemented (depends on v1beta3).
    */
-  Cursor getCursor();
+  Cursor cursor();
 }
@@ -2,12 +2,12 @@
 
 import com.google.api.client.repackaged.com.google.common.base.Preconditions;
 import com.google.api.services.datastore.DatastoreV1;
+import com.google.common.collect.AbstractIterator;
 import com.google.common.collect.ImmutableMap;
-import com.google.protobuf.ByteString;
 
 import java.util.Iterator;
 
-class QueryResultImpl<T> implements QueryResult<T> {
+class QueryResultImpl<T> extends AbstractIterator<T> implements QueryResult<T> {
 
   private static final ImmutableMap<DatastoreV1.EntityResult.ResultType, Query.Type>
       RESULT_TYPE_CONVERTER;
@@ -19,8 +19,6 @@ class QueryResultImpl<T> implements QueryResult<T> {
   private Query.Type type;
   private DatastoreV1.QueryResultBatch resultPb;
   private Iterator<DatastoreV1.EntityResult> entityResultPbIter;
-  private ByteString endCursor;
-  private int count;
 
   static {
     ImmutableMap.Builder<DatastoreV1.EntityResult.ResultType, Query.Type> builder =
@@ -57,45 +55,31 @@ private DatastoreV1.QueryResultBatch sendRequest() {
     query.populatePb(requestPb);
     resultPb = datastore.runQuery(requestPb.build()).getBatch();
     entityResultPbIter = resultPb.getEntityResultList().iterator();
-    if (DatastoreV1.QueryResultBatch.MoreResultsType.NOT_FINISHED == resultPb.getMoreResults()) {
-      endCursor = resultPb.getEndCursor();
-    } else {
-      endCursor = null;
-    }
     type = RESULT_TYPE_CONVERTER.get(resultPb.getEntityResultType());
     Preconditions.checkState(resultClass.isAssignableFrom(type.resultClass()),
         "Unexpected result type");
     return resultPb;
   }
 
   @Override
-  public boolean hasNext() {
-    return entityResultPbIter.hasNext() || endCursor != null;
-  }
-
-  @Override
-  public T next() {
-    if (!hasNext() && endCursor != null) {
+  protected T computeNext() {
+    while (!entityResultPbIter.hasNext()
+        && resultPb.getMoreResults() == DatastoreV1.QueryResultBatch.MoreResultsType.NOT_FINISHED) {
       query = query.nextQuery(resultPb);
       sendRequest();
     }
-    DatastoreV1.Entity entity = entityResultPbIter.next().getEntity();
-    count++;
-    return type.convert(entity);
-  }
-
-  @Override
-  public void remove() {
-    throw new UnsupportedOperationException("QueryResult is read-only");
+    return entityResultPbIter.hasNext()
+        ? type.<T>convert(entityResultPbIter.next().getEntity())
+        : endOfData();
   }
 
   @Override
-  public Query.Type getType() {
-    return type;
+  public Class<?> resultClass() {
+    return type.resultClass().value();
   }
 
   @Override
-  public Cursor getCursor() {
+  public Cursor cursor() {
     // TODO(ozarov): implement when v1beta3 is available
     return null;
   }