feat: Make count queries publicly available for use (#1042)

dconeybe · web-flow · commit 1c8d2428d94a · 2022-10-03T13:33:37.000-04:00
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/AggregateQuery.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/AggregateQuery.java
@@ -33,9 +33,9 @@
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
-// TODO(count) Make this class public
+/** A query that calculates aggregations over an underlying query. */
 @InternalExtensionOnly
-class AggregateQuery {
+public class AggregateQuery {
 
   /**
    * The "alias" to specify in the {@link RunAggregationQueryRequest} proto when running a count
@@ -50,11 +50,17 @@ class AggregateQuery {
     this.query = query;
   }
 
+  /** Returns the query whose aggregations will be calculated by this object. */
   @Nonnull
   public Query getQuery() {
     return query;
   }
 
+  /**
+   * Executes this query.
+   *
+   * @return An {@link ApiFuture} that will be resolved with the results of the query.
+   */
   @Nonnull
   public ApiFuture<AggregateQuerySnapshot> get() {
     return get(null);
@@ -131,6 +137,12 @@ public void onError(Throwable throwable) {
     public void onComplete() {}
   }
 
+  /**
+   * Returns the {@link RunAggregationQueryRequest} that this AggregateQuery instance represents.
+   * The request contain the serialized form of all aggregations and Query constraints.
+   *
+   * @return the serialized RunAggregationQueryRequest
+   */
   @Nonnull
   public RunAggregationQueryRequest toProto() {
     return toProto(null);
@@ -159,6 +171,17 @@ RunAggregationQueryRequest toProto(@Nullable final ByteString transactionId) {
     return request.build();
   }
 
+  /**
+   * Returns an AggregateQuery instance that can be used to execute the provided {@link
+   * RunAggregationQueryRequest}.
+   *
+   * <p>Only RunAggregationQueryRequests that pertain to the same project as the Firestore instance
+   * can be deserialized.
+   *
+   * @param firestore a Firestore instance to apply the query to.
+   * @param proto the serialized RunAggregationQueryRequest.
+   * @return a AggregateQuery instance that can be used to execute the RunAggregationQueryRequest.
+   */
   @Nonnull
   public static AggregateQuery fromProto(Firestore firestore, RunAggregationQueryRequest proto) {
     RunQueryRequest runQueryRequest =
@@ -170,19 +193,40 @@ public static AggregateQuery fromProto(Firestore firestore, RunAggregationQueryR
     return new AggregateQuery(query);
   }
 
+  /**
+   * Calculates and returns the hash code for this object.
+   *
+   * @return the hash code for this object.
+   */
   @Override
   public int hashCode() {
     return query.hashCode();
   }
 
+  /**
+   * Compares this object with the given object for equality.
+   *
+   * <p>This object is considered "equal" to the other object if and only if all of the following
+   * conditions are satisfied:
+   *
+   * <ol>
+   *   <li>{@code object} is a non-null instance of {@link AggregateQuery}.
+   *   <li>{@code object} performs the same aggregations as this {@link AggregateQuery}.
+   *   <li>The underlying {@link Query} of {@code object} compares equal to that of this object.
+   * </ol>
+   *
+   * @param object The object to compare to this object for equality.
+   * @return {@code true} if this object is "equal" to the given object, as defined above, or {@code
+   *     false} otherwise.
+   */
   @Override
-  public boolean equals(Object obj) {
-    if (obj == this) {
+  public boolean equals(Object object) {
+    if (object == this) {
       return true;
-    } else if (!(obj instanceof AggregateQuery)) {
+    } else if (!(object instanceof AggregateQuery)) {
       return false;
     }
-    AggregateQuery other = (AggregateQuery) obj;
+    AggregateQuery other = (AggregateQuery) object;
     return query.equals(other.query);
   }
 }
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/AggregateQuerySnapshot.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/AggregateQuerySnapshot.java
@@ -21,9 +21,9 @@
 import java.util.Objects;
 import javax.annotation.Nonnull;
 
-// TODO(count) Make this class public
+/** The results of executing an {@link AggregateQuery}. */
 @InternalExtensionOnly
-class AggregateQuerySnapshot {
+public class AggregateQuerySnapshot {
 
   @Nonnull private final AggregateQuery query;
   @Nonnull private final Timestamp readTime;
@@ -35,34 +35,60 @@ class AggregateQuerySnapshot {
     this.count = count;
   }
 
+  /** Returns the query that was executed to produce this result. */
   @Nonnull
   public AggregateQuery getQuery() {
     return query;
   }
 
+  /** Returns the time at which this snapshot was read. */
   @Nonnull
   public Timestamp getReadTime() {
     return readTime;
   }
 
+  /** Returns the number of documents in the result set of the underlying query. */
   public long getCount() {
     return count;
   }
 
+  /**
+   * Compares this object with the given object for equality.
+   *
+   * <p>This object is considered "equal" to the other object if and only if all of the following
+   * conditions are satisfied:
+   *
+   * <ol>
+   *   <li>{@code object} is a non-null instance of {@link AggregateQuerySnapshot}.
+   *   <li>The {@link AggregateQuery} of {@code object} compares equal to that of this object.
+   *   <li>{@code object} has the same results as this object.
+   * </ol>
+   *
+   * @param object The object to compare to this object for equality.
+   * @return {@code true} if this object is "equal" to the given object, as defined above, or {@code
+   *     false} otherwise.
+   */
   @Override
-  public boolean equals(Object obj) {
-    if (obj == this) {
+  public boolean equals(Object object) {
+    if (object == this) {
       return true;
-    } else if (!(obj instanceof AggregateQuerySnapshot)) {
+    } else if (!(object instanceof AggregateQuerySnapshot)) {
       return false;
     }
 
-    AggregateQuerySnapshot other = (AggregateQuerySnapshot) obj;
-    return query.equals(other.query) && readTime.equals(other.readTime) && count == other.count;
+    AggregateQuerySnapshot other = (AggregateQuerySnapshot) object;
+
+    // Don't check `readTime`, because `DocumentSnapshot.equals()` doesn't either.
+    return query.equals(other.query) && count == other.count;
   }
 
+  /**
+   * Calculates and returns the hash code for this object.
+   *
+   * @return the hash code for this object.
+   */
   @Override
   public int hashCode() {
-    return Objects.hash(query, readTime, count);
+    return Objects.hash(query, count);
   }
 }
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
@@ -1846,9 +1846,20 @@ private boolean isRetryableError(Throwable throwable) {
     return false;
   }
 
-  // TODO(count) Make this method public
+  /**
+   * Returns a query that counts the documents in the result set of this query.
+   *
+   * <p>The returned query, when executed, counts the documents in the result set of this query
+   * <em>without actually downloading the documents</em>.
+   *
+   * <p>Using the returned query to count the documents is efficient because only the final count,
+   * not the documents' data, is downloaded. The returned query can even count the documents if the
+   * result set would be prohibitively large to download entirely (e.g. thousands of documents).
+   *
+   * @return a query that counts the documents in the result set of this query.
+   */
   @Nonnull
-  AggregateQuery count() {
+  public AggregateQuery count() {
     return new AggregateQuery(this);
   }
 
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Transaction.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Transaction.java
@@ -191,15 +191,14 @@ public ApiFuture<QuerySnapshot> get(@Nonnull Query query) {
     return query.get(transactionId);
   }
 
-  // TODO(count) Make this method public
   /**
    * Returns the result from the provided aggregate query. Holds a pessimistic lock on all accessed
    * documents.
    *
    * @return The result of the aggregation.
    */
   @Nonnull
-  ApiFuture<AggregateQuerySnapshot> get(@Nonnull AggregateQuery query) {
+  public ApiFuture<AggregateQuerySnapshot> get(@Nonnull AggregateQuery query) {
     Preconditions.checkState(isEmpty(), READ_BEFORE_WRITE_ERROR_MSG);
 
     return query.get(transactionId);
diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/AggregateQuerySnapshotTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/AggregateQuerySnapshotTest.java
@@ -84,12 +84,12 @@ public void hashCodeShouldReturnDifferentHashCodeWhenConstructedDifferentAggrega
   }
 
   @Test
-  public void hashCodeShouldReturnDifferentHashCodeWhenConstructedDifferentTimestamp() {
+  public void hashCodeShouldReturnSameHashCodeWhenConstructedDifferentTimestamp() {
     AggregateQuerySnapshot snapshot1 =
         new AggregateQuerySnapshot(sampleAggregateQuery, sampleTimestamp, 42);
     AggregateQuerySnapshot snapshot2 =
         new AggregateQuerySnapshot(sampleAggregateQuery, sampleTimestamp2, 42);
-    assertThat(snapshot1.hashCode()).isNotEqualTo(snapshot2.hashCode());
+    assertThat(snapshot1.hashCode()).isEqualTo(snapshot2.hashCode());
   }
 
   @Test
@@ -125,12 +125,12 @@ public void equalsShouldReturnFalseWhenGivenAnAggregateQuerySnapshotWithADiffere
   }
 
   @Test
-  public void equalsShouldReturnFalseWhenGivenAnAggregateQuerySnapshotWithADifferentReadTime() {
+  public void equalsShouldReturnTrueWhenGivenAnAggregateQuerySnapshotWithADifferentReadTime() {
     AggregateQuerySnapshot snapshot1 =
         new AggregateQuerySnapshot(sampleAggregateQuery, sampleTimestamp, 42);
     AggregateQuerySnapshot snapshot2 =
         new AggregateQuerySnapshot(sampleAggregateQuery, sampleTimestamp2, 42);
-    assertThat(snapshot1.equals(snapshot2)).isFalse();
+    assertThat(snapshot1.equals(snapshot2)).isTrue();
   }
 
   @Test
diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITQueryCountTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITQueryCountTest.java
@@ -14,18 +14,29 @@
  * limitations under the License.
  */
 
-package com.google.cloud.firestore;
+package com.google.cloud.firestore.it;
 
 import static com.google.cloud.firestore.LocalFirestoreHelper.autoId;
 import static com.google.common.truth.Truth.assertThat;
 import static java.util.Collections.singletonMap;
 import static org.junit.Assert.assertThrows;
-import static org.junit.Assume.assumeFalse;
 import static org.junit.Assume.assumeTrue;
 
 import com.google.api.core.ApiFuture;
 import com.google.auto.value.AutoValue;
 import com.google.cloud.Timestamp;
+import com.google.cloud.firestore.AggregateQuery;
+import com.google.cloud.firestore.AggregateQuerySnapshot;
+import com.google.cloud.firestore.CollectionGroup;
+import com.google.cloud.firestore.CollectionReference;
+import com.google.cloud.firestore.DocumentReference;
+import com.google.cloud.firestore.Firestore;
+import com.google.cloud.firestore.FirestoreOptions;
+import com.google.cloud.firestore.Query;
+import com.google.cloud.firestore.QueryDocumentSnapshot;
+import com.google.cloud.firestore.TransactionOptions;
+import com.google.cloud.firestore.WriteBatch;
+import com.google.cloud.firestore.WriteResult;
 import com.google.common.base.Preconditions;
 import java.util.ArrayList;
 import java.util.List;
@@ -42,7 +53,6 @@
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
-// TODO(count) Move this class back into the "it" subdirectory.
 @RunWith(JUnit4.class)
 public class ITQueryCountTest {
 
@@ -289,16 +299,6 @@ public void aggregateQueryInATransactionShouldRespectReadTime() throws Exception
                     .build());
 
     Long transactionCount = transactionFuture.get();
-
-    // Put this assumption as close to the end of this method as possible, so that at least the code
-    // above can be _executed_, even if we end up skipping the assertThat() check below.
-    assumeFalse(
-        "Snapshot reads are not yet implemented in the Firestore emulator (b/220918135). "
-            + "As a result, this test will fail when run against the Firestore emulator "
-            + "because it will incorrectly ignore the read time and return the count "
-            + "of the documents at the current time.",
-        isRunningAgainstFirestoreEmulator());
-
     assertThat(transactionCount).isEqualTo(5);
   }
 
