fix: support array of references for IN queries (#211)

googleapis · May 11, 2020 · b376003 · b376003
1 parent cd19eae
commit b376003
Show file tree

Hide file tree

Showing 3 changed files with 113 additions and 4 deletions.
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Query.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Query.java
@@ -407,8 +407,10 @@ private Object convertReference(Object fieldValue) {
       reference = (DocumentReference) fieldValue;
     } else {
       throw new IllegalArgumentException(
-          "The corresponding value for FieldPath.documentId() must be a String or a "
-              + "DocumentReference.");
+          String.format(
+              "The corresponding value for FieldPath.documentId() must be a String or a "
+                  + "DocumentReference, but was: %s.",
+              fieldValue.toString()));
     }
 
     if (!basePath.isPrefixOf(reference.getResourcePath())) {
@@ -709,7 +711,27 @@ public Query whereIn(@Nonnull FieldPath fieldPath, @Nonnull List<? extends Objec
   private Query whereHelper(
       FieldPath fieldPath, StructuredQuery.FieldFilter.Operator operator, Object value) {
     if (fieldPath.equals(FieldPath.DOCUMENT_ID)) {
-      value = this.convertReference(value);
+      if (operator == ARRAY_CONTAINS || operator == ARRAY_CONTAINS_ANY) {
+        throw new IllegalArgumentException(
+            String.format(
+                "Invalid query. You cannot perform '%s' queries on FieldPath.documentId().",
+                operator.toString()));
+      } else if (operator == IN) {
+        if (!(value instanceof List) || ((List<?>) value).isEmpty()) {
+          throw new IllegalArgumentException(
+              String.format(
+                  "Invalid Query. A non-empty array is required for '%s' filters.",
+                  operator.toString()));
+        }
+        List<Object> referenceList = new ArrayList<>();
+        for (Object arrayValue : (List<Object>) value) {
+          Object convertedValue = this.convertReference(arrayValue);
+          referenceList.add(convertedValue);
+        }
+        value = referenceList;
+      } else {
+        value = this.convertReference(value);
+      }
     }
 
     Builder newOptions = options.toBuilder();

diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/QueryTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/QueryTest.java
@@ -40,6 +40,7 @@
 import com.google.api.gax.rpc.ServerStreamingCallable;
 import com.google.cloud.Timestamp;
 import com.google.cloud.firestore.spi.v1.FirestoreRpc;
+import com.google.firestore.v1.ArrayValue;
 import com.google.firestore.v1.RunQueryRequest;
 import com.google.firestore.v1.StructuredQuery;
 import com.google.firestore.v1.StructuredQuery.Direction;
@@ -273,6 +274,77 @@ public void withFieldPathFilter() throws Exception {
     }
   }
 
+  @Test
+  public void inQueriesWithReferenceArray() throws Exception {
+    doAnswer(queryResponse())
+        .when(firestoreMock)
+        .streamRequest(
+            runQuery.capture(),
+            streamObserverCapture.capture(),
+            Matchers.<ServerStreamingCallable>any());
+
+    query
+        .whereIn(
+            FieldPath.documentId(),
+            Arrays.<Object>asList("doc", firestoreMock.document("coll/doc")))
+        .get()
+        .get();
+
+    Value value =
+        Value.newBuilder()
+            .setArrayValue(
+                ArrayValue.newBuilder()
+                    .addValues(reference(DOCUMENT_NAME))
+                    .addValues(reference(DOCUMENT_NAME))
+                    .build())
+            .build();
+    RunQueryRequest expectedRequest = query(filter(Operator.IN, "__name__", value));
+
+    assertEquals(expectedRequest, runQuery.getValue());
+  }
+
+  @Test
+  public void validatesInQueries() throws Exception {
+    try {
+      query.whereIn(FieldPath.documentId(), Arrays.<Object>asList("foo", 42)).get();
+      fail();
+    } catch (IllegalArgumentException e) {
+      assertEquals(
+          "The corresponding value for FieldPath.documentId() must be a String or a "
+              + "DocumentReference, but was: 42.",
+          e.getMessage());
+    }
+
+    try {
+      query.whereIn(FieldPath.documentId(), Arrays.<Object>asList()).get();
+      fail();
+    } catch (IllegalArgumentException e) {
+      assertEquals(
+          "Invalid Query. A non-empty array is required for 'IN' filters.", e.getMessage());
+    }
+  }
+
+  @Test
+  public void validatesQueryOperatorForFieldPathDocumentId() throws Exception {
+    try {
+      query.whereArrayContains(FieldPath.documentId(), "bar");
+      fail();
+    } catch (IllegalArgumentException e) {
+      assertEquals(
+          "Invalid query. You cannot perform 'ARRAY_CONTAINS' queries on FieldPath.documentId().",
+          e.getMessage());
+    }
+
+    try {
+      query.whereArrayContainsAny(FieldPath.documentId(), Collections.<Object>singletonList("bar"));
+      fail();
+    } catch (IllegalArgumentException e) {
+      assertEquals(
+          "Invalid query. You cannot perform 'ARRAY_CONTAINS_ANY' queries on FieldPath.documentId().",
+          e.getMessage());
+    }
+  }
+
   @Test
   public void withDocumentIdFilter() throws Exception {
     doAnswer(queryResponse())
@@ -523,7 +595,7 @@ public void withInvalidStartAt() throws Exception {
     } catch (IllegalArgumentException e) {
       assertEquals(
           "The corresponding value for FieldPath.documentId() must be a String or a "
-              + "DocumentReference.",
+              + "DocumentReference, but was: 42.",
           e.getMessage());
     }
 

diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITSystemTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITSystemTest.java
@@ -1141,6 +1141,21 @@ public void inQueries() throws Exception {
     assertEquals(asList("a", "c"), querySnapshotToIds(querySnapshot));
   }
 
+  @Test
+  public void inQueriesWithDocumentId() throws Exception {
+    DocumentReference doc1 = setDocument("a", map("count", 1));
+    DocumentReference doc2 = setDocument("b", map("count", 2));
+    setDocument("c", map("count", 3));
+
+    QuerySnapshot querySnapshot =
+        randomColl
+            .whereIn(FieldPath.documentId(), Arrays.<Object>asList(doc1.getId(), doc2))
+            .get()
+            .get();
+
+    assertEquals(asList("a", "b"), querySnapshotToIds(querySnapshot));
+  }
+
   @Test
   public void arrayContainsAnyQueries() throws Exception {
     setDocument("a", map("array", asList(42)));