Include cursor workaround, fix distinct in StructuredQuery, and remove extra checks

Ajay Kannan · Ajay Kannan · commit acf2c4207f38 · 2015-10-14T18:55:30.000-07:00
diff --git a/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/GqlQuery.java b/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/GqlQuery.java
@@ -20,7 +20,6 @@
 import static com.google.gcloud.datastore.Validator.validateNamespace;
 
 import com.google.common.base.MoreObjects;
-import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSortedMap;
@@ -138,8 +137,10 @@ static Binding fromPb(com.google.datastore.v1beta3.GqlQueryParameter argPb) {
       switch (argPb.getParameterTypeCase()) {
         case CURSOR:
           return new Binding(new Cursor(argPb.getCursor()));
-        default:
+        case VALUE:
           return new Binding(Value.fromPb(argPb.getValue()));
+        default:
+          throw new AssertionError("Unexpected enum value " + argPb.getParameterTypeCase());
       }
     }
   }
@@ -397,13 +398,11 @@ private static <V> GqlQuery<V> fromPb(
     for (Map.Entry<String, com.google.datastore.v1beta3.GqlQueryParameter> nameArg
          : queryPb.getNamedBindings().entrySet()) {
       Binding currBinding = Binding.fromPb(nameArg.getValue());
-      Preconditions.checkState(currBinding != null);
       builder.namedBindings.put(nameArg.getKey(), currBinding);
     }
     for (com.google.datastore.v1beta3.GqlQueryParameter numberArg
          : queryPb.getPositionalBindingsList()) {
       Binding currBinding = Binding.fromPb(numberArg);
-      Preconditions.checkState(currBinding != null);
       builder.positionalBindings.add(currBinding);
     }
     return builder.build();
diff --git a/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/QueryResultsImpl.java b/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/QueryResultsImpl.java
@@ -16,9 +16,9 @@
 
 package com.google.gcloud.datastore;
 
-import com.google.datastore.v1beta3.QueryResultBatch.MoreResultsType;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.AbstractIterator;
+import com.google.datastore.v1beta3.QueryResultBatch.MoreResultsType;
 import com.google.gcloud.datastore.Query.ResultType;
 import com.google.protobuf.ByteString;
 
@@ -55,6 +55,11 @@ class QueryResultsImpl<T> extends AbstractIterator<T> implements QueryResults<T>
     }
     partitionIdPb = pbBuilder.build();
     sendRequest();
+    if (queryResultBatchPb.getSkippedResults() > 0) {
+      cursor = queryResultBatchPb.getSkippedCursor();
+    } else {
+      cursor = ByteString.EMPTY;
+    }
   }
 
   private void sendRequest() {
@@ -68,13 +73,6 @@ private void sendRequest() {
     queryResultBatchPb = datastore.runQuery(requestPb.build()).getBatch();
     lastBatch = queryResultBatchPb.getMoreResults() != MoreResultsType.NOT_FINISHED;
     entityResultPbIter = queryResultBatchPb.getEntityResultsList().iterator();
-    if (queryResultBatchPb.getSkippedResults() > 0) {
-      cursor = queryResultBatchPb.getSkippedCursor();
-    } else if (entityResultPbIter.hasNext()) {
-      cursor = queryResultBatchPb.getEntityResults(0).getCursor();
-    } else {
-      cursor = queryResultBatchPb.getEndCursor();
-    }
     actualResultType = ResultType.fromPb(queryResultBatchPb.getEntityResultType());
     if (Objects.equals(queryResultType, ResultType.PROJECTION_ENTITY)) {
       // projection entity can represent all type of results
diff --git a/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/StructuredQuery.java b/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/StructuredQuery.java
@@ -108,8 +108,10 @@ static Filter fromPb(com.google.datastore.v1beta3.Filter filterPb) {
       switch (filterPb.getFilterTypeCase()) {
         case COMPOSITE_FILTER:
           return CompositeFilter.fromPb(filterPb.getCompositeFilter());
-        default:
+        case PROPERTY_FILTER:
           return PropertyFilter.fromPb(filterPb.getPropertyFilter());
+        default:
+          throw new AssertionError("Unexpected enum value " + filterPb.getFilterTypeCase());
       }
     }
   }
@@ -530,7 +532,7 @@ static class BaseBuilder<V, B extends BaseBuilder<V, B>> {
     private String kind;
     private final List<String> projection = new LinkedList<>();
     private Filter filter;
-    private boolean distinctOnAll = false;
+    private boolean distinct = false;
     private final List<String> distinctOn = new LinkedList<>();
     private final List<OrderBy> orderBy = new LinkedList<>();
     private Cursor startCursor;
@@ -620,14 +622,16 @@ B addProjection(String projection, String... others) {
 
     B clearDistinct() {
       distinctOn.clear();
-      distinctOnAll = false;
+      distinct = false;
       return self();
     }
 
     B distinct(String... properties) {
       clearDistinct();
       if (properties.length == 0) {
-        this.distinctOnAll = true;
+        clearDistinct();
+        this.distinct = true;
+        this.distinctOn.addAll(this.projection);
       } else if (properties.length == 1) {
         addDistinct(properties[0]);
       } else {
@@ -637,6 +641,9 @@ B distinct(String... properties) {
     }
 
     B addDistinct(String property, String... others) {
+      if (this.distinct) {
+        throw new IllegalStateException("\"distinct()\" is currently set.");
+      }
       this.distinctOn.add(property);
       Collections.addAll(this.distinctOn, others);
       return self();
@@ -672,15 +679,10 @@ B mergeFrom(com.google.datastore.v1beta3.Query queryPb) {
       for (com.google.datastore.v1beta3.PropertyReference distinctOnPb : queryPb.getDistinctOnList()) {
         addDistinct(distinctOnPb.getName());
       }
-      distinctOnAll = false;
       return self();
     }
 
     public StructuredQuery<V> build() {
-      if (distinctOnAll) {
-        clearDistinct();
-        this.distinctOn.addAll(this.projection);
-      }
       return new StructuredQuery<>(this);
     }
   }
