Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Browse files
Browse the repository at this point in the history
feat!: add support for the Query Partition API (#202)
- Loading branch information
1 parent
15d68cd
commit 3996548
Showing
11 changed files
with
334 additions
and
91 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
86 changes: 86 additions & 0 deletions
86
google-cloud-firestore/src/main/java/com/google/cloud/firestore/CollectionGroup.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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(); | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
97 changes: 97 additions & 0 deletions
97
google-cloud-firestore/src/main/java/com/google/cloud/firestore/QueryPartition.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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; | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.