Allowing DocumentReferences for Query cursors

Author: schmidt-sebastian

google-cloud-firestore/src/main/java/com/google/cloud/firestore/BasePath.java

@@ -73,6 +73,26 @@ B append(BasePath<B> path) {
     return createPathWithSegments(components.build());
   }
 
+  /**
+   * Checks to see if this path is a prefix of (or equals) another path.
+   *
+   * @param path the path to check against
+   * @return true if current path is a prefix of the other path.
+   */
+  boolean isPrefixOf(B path) {
+    ImmutableList<String> prefixSegments = getSegments();
+    ImmutableList<String> childSegments = path.getSegments();
+    if (prefixSegments.size() > path.getSegments().size()) {
+      return false;
+    }
+    for (int i = 0; i < prefixSegments.size(); i++) {
+      if (!prefixSegments.get(i).equals(childSegments.get(i))) {
+        return false;
+      }
+    }
+    return true;
+  }
+
   abstract String[] splitChildPath(String path);
 
   abstract B createPathWithSegments(ImmutableList<String> segments);

google-cloud-firestore/src/main/java/com/google/cloud/firestore/DocumentReference.java

@@ -381,6 +381,10 @@ public void remove() {
     };
   }
 
+  ResourcePath getResourcePath() {
+    return path;
+  }
+
   @Override
   public String toString() {
     return String.format("DocumentReference{path=%s}", path);

google-cloud-firestore/src/main/java/com/google/cloud/firestore/Query.java

@@ -55,25 +55,6 @@ public class Query {
   final FirestoreImpl firestore;
   final QueryOptions options;
 
-  @Override
-  public boolean equals(Object o) {
-    if (this == o) {
-      return true;
-    }
-    if (o == null || getClass() != o.getClass()) {
-      return false;
-    }
-    Query query = (Query) o;
-    return Objects.equals(path, query.path)
-        && Objects.equals(firestore, query.firestore)
-        && Objects.equals(options, query.options);
-  }
-
-  @Override
-  public int hashCode() {
-    return Objects.hash(path, firestore, options);
-  }
-
   /** The direction of a sort. */
   public enum Direction {
     ASCENDING(StructuredQuery.Direction.ASCENDING),
@@ -254,9 +235,31 @@ private Cursor createCursor(Object[] fieldValues, boolean before) {
           .getField()
           .getFieldPath()
           .equals(FieldPath.DOCUMENT_ID.getEncodedPath())) {
-        Preconditions.checkArgument(
-            fieldValues[i] instanceof String, "Argument at index %d must be a String", i);
-        sanitizedValue = new DocumentReference(firestore, path.append((String) fieldValues[i]));
+        DocumentReference cursorDocument;
+        if (fieldValues[i] instanceof String) {
+          cursorDocument = new DocumentReference(firestore, path.append((String) fieldValues[i]));
+        } else if (fieldValues[i] instanceof DocumentReference) {
+          cursorDocument = (DocumentReference) fieldValues[i];
+        } else {
+          throw new IllegalArgumentException(
+              "The corresponding value for FieldPath.documentId() must be a String or a "
+                  + "DocumentReference.");
+        }
+
+        if (!this.path.isPrefixOf(cursorDocument.getResourcePath())) {
+          throw new IllegalArgumentException(
+              String.format(
+                  "'%s' is not part of the query result set and cannot be used as a query boundary.",
+                  cursorDocument.getPath()));
+        }
+        if (!cursorDocument.getParent().getResourcePath().equals(this.path)) {
+          throw new IllegalArgumentException(
+              String.format(
+                  "Only a direct child can be used as a query boundary. Found: '%s'",
+                  cursorDocument.getPath()));
+        }
+
+        sanitizedValue = cursorDocument;
       } else {
         sanitizedValue = CustomClassMapper.serialize(fieldValues[i]);
       }
@@ -770,4 +773,27 @@ public void onCompleted() {
 
     return result;
   }
+
+  ResourcePath getResourcePath() {
+    return path;
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    Query query = (Query) o;
+    return Objects.equals(path, query.path)
+        && Objects.equals(firestore, query.firestore)
+        && Objects.equals(options, query.options);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(path, firestore, options);
+  }
 }

google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java

@@ -36,8 +36,8 @@
 import javax.annotation.Nullable;
 
 /**
- * Abstract class that collects and bundles all write operations for {@link Transaction} and
- * {@link WriteBatch}.
+ * Abstract class that collects and bundles all write operations for {@link Transaction} and {@link
+ * WriteBatch}.
  */
 abstract class UpdateBuilder<T extends UpdateBuilder> {
 

google-cloud-firestore/src/test/java/com/google/cloud/firestore/LocalFirestoreHelper.java

@@ -410,14 +410,22 @@ public static StructuredQuery select(FieldPath fieldPath) {
   }
 
   public static StructuredQuery startAt(boolean before) {
+    return startAt(string(), before);
+  }
+
+  public static StructuredQuery startAt(Value value, boolean before) {
     StructuredQuery.Builder structuredQuery = StructuredQuery.newBuilder();
-    structuredQuery.setStartAt(Cursor.newBuilder().setBefore(before).addValues(string()));
+    structuredQuery.setStartAt(Cursor.newBuilder().setBefore(before).addValues(value));
     return structuredQuery.build();
   }
 
   public static StructuredQuery endAt(boolean before) {
+    return endAt(string(), before);
+  }
+
+  public static StructuredQuery endAt(Value value, boolean before) {
     StructuredQuery.Builder structuredQuery = StructuredQuery.newBuilder();
-    structuredQuery.setEndAt(Cursor.newBuilder().setBefore(before).addValues(string()));
+    structuredQuery.setEndAt(Cursor.newBuilder().setBefore(before).addValues(value));
     return structuredQuery.build();
   }
 

google-cloud-firestore/src/test/java/com/google/cloud/firestore/QueryTest.java

@@ -36,6 +36,7 @@
 import com.google.cloud.firestore.spi.v1beta1.FirestoreRpc;
 import com.google.firestore.v1beta1.RunQueryRequest;
 import com.google.firestore.v1beta1.StructuredQuery;
+import com.google.firestore.v1beta1.Value;
 import java.util.Arrays;
 import java.util.Iterator;
 import java.util.concurrent.Semaphore;
@@ -255,14 +256,54 @@ public void withStartAt() throws Exception {
             streamObserverCapture.capture(),
             Matchers.<ServerStreamingCallable>any());
 
-    query.orderBy("foo").startAt("bar").get().get();
+    query.orderBy("foo").orderBy(FieldPath.documentId()).startAt("bar", "foo").get().get();
+
+    Value documentBoundary =
+        Value.newBuilder().setReferenceValue(query.getResourcePath().toString() + "/foo").build();
 
     RunQueryRequest queryRequest =
-        query(order("foo", StructuredQuery.Direction.ASCENDING), startAt(true));
+        query(
+            order("foo", StructuredQuery.Direction.ASCENDING),
+            order("__name__", StructuredQuery.Direction.ASCENDING),
+            startAt(true),
+            startAt(documentBoundary, true));
 
     assertEquals(queryRequest, runQuery.getValue());
   }
 
+  @Test
+  public void withInvalidStartAt() throws Exception {
+    try {
+      query.orderBy(FieldPath.documentId()).startAt(42).get();
+      fail();
+    } catch (IllegalArgumentException e) {
+      assertEquals(
+          "The corresponding value for FieldPath.documentId() must be a String or a "
+              + "DocumentReference.",
+          e.getMessage());
+    }
+
+    try {
+      query.orderBy(FieldPath.documentId()).startAt("coll/doc/coll").get();
+      fail();
+    } catch (IllegalArgumentException e) {
+      assertEquals(
+          "Only a direct child can be used as a query boundary. Found: "
+              + "'projects/test-project/databases/(default)/documents/coll/coll/doc/coll'",
+          e.getMessage());
+    }
+
+    try {
+      query.orderBy(FieldPath.documentId()).startAt(firestoreMock.document("foo/bar")).get();
+      fail();
+    } catch (IllegalArgumentException e) {
+      assertEquals(
+          "'projects/test-project/databases/(default)/documents/foo/bar' is not part of "
+              + "the query result set and cannot be used as a query boundary.",
+          e.getMessage());
+    }
+  }
+
   @Test
   public void withStartAfter() throws Exception {
     doAnswer(queryResponse())