feat!: add support for the Query Partition API (#202)

Author: schmidt-sebastian

google-cloud-firestore/clirr-ignored-differences.xml

@@ -142,4 +142,19 @@
  <to>*</to>
  </difference>
 
+ <!--
+ Query Partition API Feature
+ -->
+ <difference>
+ <differenceType>7012</differenceType>
+ <className>com/google/cloud/firestore/spi/v1/FirestoreRpc</className>
+ <method>com.google.api.gax.rpc.UnaryCallable partitionQueryPagedCallable()</method>
+ </difference>
+ <difference>
+ <differenceType>7006</differenceType>
+ <className>com/google/cloud/firestore/Firestore</className>
+ <method>com.google.cloud.firestore.Query collectionGroup(java.lang.String)</method>
+ <to>com.google.cloud.firestore.CollectionGroup</to>
+ </difference>
+
 </differences>

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

@@ -0,0 +1,86 @@
+/*
+ * Copyright 2020 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.firestore;
+
+import com.google.api.gax.rpc.ApiException;
+import com.google.api.gax.rpc.ApiExceptions;
+import com.google.api.gax.rpc.ApiStreamObserver;
+import com.google.cloud.firestore.v1.FirestoreClient;
+import com.google.firestore.v1.Cursor;
+import com.google.firestore.v1.PartitionQueryRequest;
+import javax.annotation.Nullable;
+
+/**
+ * A Collection Group query matches all documents that are contained in a collection or
+ * subcollection with a specific collection ID.
+ */
+public class CollectionGroup extends Query {
+ CollectionGroup(FirestoreRpcContext<?> rpcContext, String collectionId) {
+ super(
+ rpcContext,
+ QueryOptions.builder()
+ .setParentPath(rpcContext.getResourcePath())
+ .setCollectionId(collectionId)
+ .setAllDescendants(true)
+ .build());
+ }
+
+ /**
+ * Partitions a query by returning partition cursors that can be used to run the query in
+ * parallel. The returned partition cursors are split points that can be used as starting/end
+ * points for the query results.
+ *
+ * @param desiredPartitionCount The desired maximum number of partition points. The number must be
+ * strictly positive. The actual number of partitions returned may be fewer.
+ * @param observer a stream observer that receives the result of the Partition request.
+ */
+ public void getPartitions(
+ long desiredPartitionCount, ApiStreamObserver<QueryPartition> observer) {
+ // Partition queries require explicit ordering by __name__.
+ Query queryWithDefaultOrder = orderBy(FieldPath.DOCUMENT_ID);
+
+ PartitionQueryRequest.Builder request = PartitionQueryRequest.newBuilder();
+ request.setStructuredQuery(queryWithDefaultOrder.buildQuery());
+ request.setParent(options.getParentPath().toString());
+
+ // Since we are always returning an extra partition (with en empty endBefore cursor), we
+ // reduce the desired partition count by one.
+ request.setPartitionCount(desiredPartitionCount - 1);
+
+ final FirestoreClient.PartitionQueryPagedResponse response;
+ try {
+ response =
+ ApiExceptions.callAndTranslateApiException(
+ rpcContext.sendRequest(
+ request.build(), rpcContext.getClient().partitionQueryPagedCallable()));
+ } catch (ApiException exception) {
+ throw FirestoreException.apiException(exception);
+ }
+
+ @Nullable Object[] lastCursor = null;
+ for (Cursor cursor : response.iterateAll()) {
+ Object[] decodedCursorValue = new Object[cursor.getValuesCount()];
+ for (int i = 0; i < cursor.getValuesCount(); ++i) {
+ decodedCursorValue[i] = UserDataConverter.decodeValue(rpcContext, cursor.getValues(i));
+ }
+ observer.onNext(new QueryPartition(queryWithDefaultOrder, lastCursor, decodedCursorValue));
+ lastCursor = decodedCursorValue;
+ }
+ observer.onNext(new QueryPartition(queryWithDefaultOrder, lastCursor, null));
+ observer.onCompleted();
+ }
+}

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

@@ -23,11 +23,9 @@
 import com.google.firestore.v1.Document;
 import com.google.firestore.v1.Value;
 import com.google.firestore.v1.Write;
-import java.util.ArrayList;
 import java.util.Date;
 import java.util.HashMap;
 import java.util.Iterator;
-import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import javax.annotation.Nonnull;
@@ -115,48 +113,6 @@ static DocumentSnapshot fromMissing(
  return new DocumentSnapshot(rpcContext, documentReference, null, readTime, null, null);
  }
 
- private Object decodeValue(Value v) {
- Value.ValueTypeCase typeCase = v.getValueTypeCase();
- switch (typeCase) {
- case NULL_VALUE:
- return null;
- case BOOLEAN_VALUE:
- return v.getBooleanValue();
- case INTEGER_VALUE:
- return v.getIntegerValue();
- case DOUBLE_VALUE:
- return v.getDoubleValue();
- case TIMESTAMP_VALUE:
- return Timestamp.fromProto(v.getTimestampValue());
- case STRING_VALUE:
- return v.getStringValue();
- case BYTES_VALUE:
- return Blob.fromByteString(v.getBytesValue());
- case REFERENCE_VALUE:
- String pathName = v.getReferenceValue();
- return new DocumentReference(rpcContext, ResourcePath.create(pathName));
- case GEO_POINT_VALUE:
- return new GeoPoint(
- v.getGeoPointValue().getLatitude(), v.getGeoPointValue().getLongitude());
- case ARRAY_VALUE:
- List<Object> list = new ArrayList<>();
- List<Value> lv = v.getArrayValue().getValuesList();
- for (Value iv : lv) {
- list.add(decodeValue(iv));
- }
- return list;
- case MAP_VALUE:
- Map<String, Object> outputMap = new HashMap<>();
- Map<String, Value> inputMap = v.getMapValue().getFieldsMap();
- for (Map.Entry<String, Value> entry : inputMap.entrySet()) {
- outputMap.put(entry.getKey(), decodeValue(entry.getValue()));
- }
- return outputMap;
- default:
- throw FirestoreException.invalidState(String.format("Unknown Value Type: %s", typeCase));
- }
- }
-
  /**
  * Returns the time at which this snapshot was read.
  *
@@ -214,7 +170,7 @@ public Map<String, Object> getData() {
 
  Map<String, Object> decodedFields = new HashMap<>();
  for (Map.Entry<String, Value> entry : fields.entrySet()) {
- Object decodedValue = decodeValue(entry.getValue());
+ Object decodedValue = UserDataConverter.decodeValue(rpcContext, entry.getValue());
  decodedFields.put(entry.getKey(), decodedValue);
  }
  return decodedFields;
@@ -293,7 +249,7 @@ public Object get(@Nonnull FieldPath fieldPath) {
  return null;
  }
 
- return decodeValue(value);
+ return UserDataConverter.decodeValue(rpcContext, value);
  }
 
  /**

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

@@ -56,14 +56,14 @@ public interface Firestore extends Service<FirestoreOptions>, AutoCloseable {
  Iterable<CollectionReference> listCollections();
 
  /**
- * Creates and returns a new @link{Query} that includes all documents in the database that are
- * contained in a collection or subcollection with the given @code{collectionId}.
+ * Creates and returns a new {@link CollectionGroup} that includes all documents in the database
+ * that are contained in a collection or subcollection with the given @code{collectionId}.
  *
  * @param collectionId Identifies the collections to query over. Every collection or subcollection
  * with this ID as the last segment of its path will be included. Cannot contain a slash.
  * @return The created Query.
  */
- Query collectionGroup(@Nonnull String collectionId);
+ CollectionGroup collectionGroup(@Nonnull String collectionId);
 
  /**
  * Executes the given updateFunction and then attempts to commit the changes applied within the

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

@@ -263,12 +263,12 @@ public void onCompleted() {
 
  @Nonnull
  @Override
- public Query collectionGroup(@Nonnull final String collectionId) {
+ public CollectionGroup collectionGroup(@Nonnull final String collectionId) {
  Preconditions.checkArgument(
  !collectionId.contains("/"),
  String.format(
  "Invalid collectionId '%s'. Collection IDs must not contain '/'.", collectionId));
- return new Query(this, collectionId);
+ return new CollectionGroup(this, collectionId);
  }
 
  @Nonnull

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

@@ -296,21 +296,7 @@ abstract static class Builder {
  .build());
  }
 
- /**
- * Creates a Collection Group query that matches all documents directly nested under a
- * specifically named collection
- */
- Query(FirestoreRpcContext<?> rpcContext, String collectionId) {
- this(
- rpcContext,
- QueryOptions.builder()
- .setParentPath(rpcContext.getResourcePath())
- .setCollectionId(collectionId)
- .setAllDescendants(true)
- .build());
- }
-
- private Query(FirestoreRpcContext<?> rpcContext, QueryOptions queryOptions) {
+ protected Query(FirestoreRpcContext<?> rpcContext, QueryOptions queryOptions) {
  this.rpcContext = rpcContext;
  this.options = queryOptions;
  }

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

@@ -0,0 +1,97 @@
+/*
+ * Copyright 2020 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.firestore;
+
+import java.util.Arrays;
+import java.util.Objects;
+import javax.annotation.Nullable;
+
+/**
+ * A split point that can be used in a query as a starting and/or end point for the query results.
+ * The cursors returned by {@link #getStartAt()} and {@link #getEndBefore()} can only be used in a
+ * query that matches the constraint of query that produced this partition.
+ */
+public class QueryPartition {
+ private final Query query;
+ @Nullable private final Object[] startAt;
+ @Nullable private final Object[] endBefore;
+
+ public QueryPartition(Query query, @Nullable Object[] startAt, @Nullable Object[] endBefore) {
+ this.query = query;
+ this.startAt = startAt;
+ this.endBefore = endBefore;
+ }
+
+ /**
+ * The cursor that defines the first result for this partition. {@code null} if this is the first
+ * partition.
+ *
+ * @return a cursor value that can be used with {@link Query#startAt(Object...)} or {@code null}
+ * if this is the first partition.
+ */
+ @Nullable
+ public Object[] getStartAt() {
+ return startAt;
+ }
+
+ /**
+ * The cursor that defines the first result after this partition. {@code null} if this is the last
+ * partition.
+ *
+ * @return a cursor value that can be used with {@link Query#endBefore(Object...)} or {@code null}
+ * if this is the last partition.
+ */
+ @Nullable
+ public Object[] getEndBefore() {
+ return endBefore;
+ }
+
+ /**
+ * Returns a query that only returns the documents for this partition.
+ *
+ * @return a query partitioned by a {@link Query#startAt(Object...)} and {@link
+ * Query#endBefore(Object...)} cursor.
+ */
+ public Query createQuery() {
+ Query baseQuery = query;
+ if (startAt != null) {
+ baseQuery = baseQuery.startAt(startAt);
+ }
+ if (endBefore != null) {
+ baseQuery = baseQuery.endBefore(endBefore);
+ }
+ return baseQuery;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (!(o instanceof QueryPartition)) return false;
+ QueryPartition partition = (QueryPartition) o;
+ return query.equals(partition.query)
+ && Arrays.equals(startAt, partition.startAt)
+ && Arrays.equals(endBefore, partition.endBefore);
+ }
+
+ @Override
+ public int hashCode() {
+ int result = Objects.hash(query);
+ result = 31 * result + Arrays.hashCode(startAt);
+ result = 31 * result + Arrays.hashCode(endBefore);
+ return result;
+ }
+}

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

@@ -22,7 +22,9 @@
 import com.google.firestore.v1.MapValue;
 import com.google.firestore.v1.Value;
 import com.google.protobuf.NullValue;
+import java.util.ArrayList;
 import java.util.Date;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
@@ -181,4 +183,46 @@ static Value encodeValue(
 
  throw FirestoreException.invalidState("Cannot convert %s to Firestore Value", sanitizedObject);
  }
+
+ static Object decodeValue(FirestoreRpcContext<?> rpcContext, Value v) {
+ Value.ValueTypeCase typeCase = v.getValueTypeCase();
+ switch (typeCase) {
+ case NULL_VALUE:
+ return null;
+ case BOOLEAN_VALUE:
+ return v.getBooleanValue();
+ case INTEGER_VALUE:
+ return v.getIntegerValue();
+ case DOUBLE_VALUE:
+ return v.getDoubleValue();
+ case TIMESTAMP_VALUE:
+ return Timestamp.fromProto(v.getTimestampValue());
+ case STRING_VALUE:
+ return v.getStringValue();
+ case BYTES_VALUE:
+ return Blob.fromByteString(v.getBytesValue());
+ case REFERENCE_VALUE:
+ String pathName = v.getReferenceValue();
+ return new DocumentReference(rpcContext, ResourcePath.create(pathName));
+ case GEO_POINT_VALUE:
+ return new GeoPoint(
+ v.getGeoPointValue().getLatitude(), v.getGeoPointValue().getLongitude());
+ case ARRAY_VALUE:
+ List<Object> list = new ArrayList<>();
+ List<Value> lv = v.getArrayValue().getValuesList();
+ for (Value iv : lv) {
+ list.add(decodeValue(rpcContext, iv));
+ }
+ return list;
+ case MAP_VALUE:
+ Map<String, Object> outputMap = new HashMap<>();
+ Map<String, Value> inputMap = v.getMapValue().getFieldsMap();
+ for (Map.Entry<String, Value> entry : inputMap.entrySet()) {
+ outputMap.put(entry.getKey(), decodeValue(rpcContext, entry.getValue()));
+ }
+ return outputMap;
+ default:
+ throw FirestoreException.invalidState(String.format("Unknown Value Type: %s", typeCase));
+ }
+ }
 }