Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add Query.limitToLast() #151

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
Expand Up @@ -16,6 +16,7 @@

package com.google.cloud.firestore;

import static com.google.common.collect.Lists.reverse;
import static com.google.firestore.v1.StructuredQuery.FieldFilter.Operator.ARRAY_CONTAINS;
import static com.google.firestore.v1.StructuredQuery.FieldFilter.Operator.ARRAY_CONTAINS_ANY;
import static com.google.firestore.v1.StructuredQuery.FieldFilter.Operator.EQUAL;
Expand Down Expand Up @@ -182,6 +183,12 @@ Order toProto() {
}
}

/** Denotes whether a provided limit is applied to the beginning or the end of the result set. */
enum LimitType {
First,
Last
}

/** Options that define a Firestore Query. */
@AutoValue
abstract static class QueryOptions {
Expand All @@ -194,6 +201,8 @@ abstract static class QueryOptions {

abstract @Nullable Integer getLimit();

abstract LimitType getLimitType();

abstract @Nullable Integer getOffset();

abstract @Nullable Cursor getStartCursor();
Expand All @@ -209,6 +218,7 @@ abstract static class QueryOptions {
static Builder builder() {
return new AutoValue_Query_QueryOptions.Builder()
.setAllDescendants(false)
.setLimitType(LimitType.First)
.setFieldOrders(ImmutableList.<FieldOrder>of())
.setFieldFilters(ImmutableList.<FieldFilter>of())
.setFieldProjections(ImmutableList.<FieldReference>of());
Expand All @@ -226,6 +236,8 @@ abstract static class Builder {

abstract Builder setLimit(Integer value);

abstract Builder setLimitType(LimitType value);

abstract Builder setOffset(Integer value);

abstract Builder setStartCursor(@Nullable Cursor value);
Expand Down Expand Up @@ -764,15 +776,33 @@ public Query orderBy(@Nonnull FieldPath fieldPath, @Nonnull Direction direction)
}

/**
* Creates and returns a new Query that's additionally limited to only return up to the specified
* number of documents.
* Creates and returns a new Query that only returns the first matching documents.
*
* @param limit The maximum number of items to return.
* @return The created Query.
*/
@Nonnull
public Query limit(int limit) {
return new Query(firestore, options.toBuilder().setLimit(limit).build());
return new Query(
firestore, options.toBuilder().setLimit(limit).setLimitType(LimitType.First).build());
}

/**
* Creates and returns a new Query that only returns the last matching documents.
*
* <p>You must specify at least one orderBy clause for limitToLast queries. Otherwise, an {@link
* java.lang.IllegalStateException} is thrown during execution.
*
* <p>Results for limitToLast() queries are only available once all documents are received. Hence,
* limitToLast() queries cannot be streamed via the {@link #stream(ApiStreamObserver)} API.
*
* @param limit the maximum number of items to return
* @return the created Query
*/
@Nonnull
public Query limitToLast(int limit) {
return new Query(
firestore, options.toBuilder().setLimit(limit).setLimitType(LimitType.Last).build());
}

/**
Expand Down Expand Up @@ -1004,8 +1034,25 @@ StructuredQuery.Builder buildQuery() {

if (!options.getFieldOrders().isEmpty()) {
for (FieldOrder order : options.getFieldOrders()) {
structuredQuery.addOrderBy(order.toProto());
switch (options.getLimitType()) {
case First:
structuredQuery.addOrderBy(order.toProto());
break;
case Last:
// Flip the orderBy directions since we want the last results
order =
new FieldOrder(
order.fieldPath,
order.direction.equals(Direction.ASCENDING)
? Direction.DESCENDING
: Direction.ASCENDING);
structuredQuery.addOrderBy(order.toProto());
break;
}
}
} else if (LimitType.Last.equals(options.getLimitType())) {
throw new IllegalStateException(
"limitToLast() queries require specifying at least one orderBy() clause.");
}

if (!options.getFieldProjections().isEmpty()) {
Expand All @@ -1021,11 +1068,39 @@ StructuredQuery.Builder buildQuery() {
}

if (options.getStartCursor() != null) {
structuredQuery.setStartAt(options.getStartCursor());
switch (options.getLimitType()) {
case First:
structuredQuery.setStartAt(options.getStartCursor());
break;
case Last:
// Swap the cursors to match the flipped query ordering.
Cursor cursor =
options
.getStartCursor()
.toBuilder()
.setBefore(!options.getStartCursor().getBefore())
.build();
structuredQuery.setEndAt(cursor);
break;
}
BenWhitehead marked this conversation as resolved.
Show resolved Hide resolved
}

if (options.getEndCursor() != null) {
structuredQuery.setEndAt(options.getEndCursor());
switch (options.getLimitType()) {
case First:
structuredQuery.setEndAt(options.getEndCursor());
break;
case Last:
// Swap the cursors to match the flipped query ordering.
Cursor cursor =
options
.getEndCursor()
.toBuilder()
.setBefore(!options.getEndCursor().getBefore())
.build();
structuredQuery.setStartAt(cursor);
break;
}
BenWhitehead marked this conversation as resolved.
Show resolved Hide resolved
}

return structuredQuery;
Expand All @@ -1037,7 +1112,12 @@ StructuredQuery.Builder buildQuery() {
* @param responseObserver The observer to be notified when results arrive.
*/
public void stream(@Nonnull final ApiStreamObserver<DocumentSnapshot> responseObserver) {
stream(
Preconditions.checkState(
!LimitType.Last.equals(Query.this.options.getLimitType()),
"Query results for queries that include limitToLast() constraints cannot be streamed. "
+ "Use Query.get() instead.");
BenWhitehead marked this conversation as resolved.
Show resolved Hide resolved

internalStream(
new QuerySnapshotObserver() {
@Override
public void onNext(QueryDocumentSnapshot documentSnapshot) {
Expand Down Expand Up @@ -1073,7 +1153,7 @@ Timestamp getReadTime() {
}
}

private void stream(
private void internalStream(
final QuerySnapshotObserver documentObserver, @Nullable ByteString transactionId) {
RunQueryRequest.Builder request = RunQueryRequest.newBuilder();
request.setStructuredQuery(buildQuery()).setParent(options.getParentPath().toString());
Expand Down Expand Up @@ -1178,7 +1258,7 @@ public ListenerRegistration addSnapshotListener(
ApiFuture<QuerySnapshot> get(@Nullable ByteString transactionId) {
final SettableApiFuture<QuerySnapshot> result = SettableApiFuture.create();

stream(
internalStream(
new QuerySnapshotObserver() {
List<QueryDocumentSnapshot> documentSnapshots = new ArrayList<>();

Expand All @@ -1194,8 +1274,14 @@ public void onError(Throwable throwable) {

@Override
public void onCompleted() {
// The results for limitToLast queries need to be flipped since we reversed the
// ordering constraints before sending the query to the backend.
List<QueryDocumentSnapshot> resultView =
LimitType.Last.equals(Query.this.options.getLimitType())
? reverse(documentSnapshots)
: documentSnapshots;
QuerySnapshot querySnapshot =
QuerySnapshot.withDocuments(Query.this, this.getReadTime(), documentSnapshots);
QuerySnapshot.withDocuments(Query.this, this.getReadTime(), resultView);
result.set(querySnapshot);
}
},
Expand Down
Expand Up @@ -29,6 +29,7 @@
import static com.google.cloud.firestore.LocalFirestoreHelper.reference;
import static com.google.cloud.firestore.LocalFirestoreHelper.select;
import static com.google.cloud.firestore.LocalFirestoreHelper.startAt;
import static com.google.cloud.firestore.LocalFirestoreHelper.string;
import static com.google.cloud.firestore.LocalFirestoreHelper.unaryFilter;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
Expand Down Expand Up @@ -92,6 +93,96 @@ public void withLimit() throws Exception {
assertEquals(query(limit(42)), runQuery.getValue());
}

@Test
public void limitToLastReversesOrderingConstraints() throws Exception {
doAnswer(queryResponse())
.when(firestoreMock)
.streamRequest(
runQuery.capture(),
streamObserverCapture.capture(),
Matchers.<ServerStreamingCallable>any());

query.orderBy("foo").limitToLast(42).get().get();

assertEquals(
query(limit(42), order("foo", StructuredQuery.Direction.DESCENDING)), runQuery.getValue());
}

@Test
public void limitToLastReversesCursors() throws Exception {
doAnswer(queryResponse())
.when(firestoreMock)
.streamRequest(
runQuery.capture(),
streamObserverCapture.capture(),
Matchers.<ServerStreamingCallable>any());

query.orderBy("foo").startAt("foo").endAt("bar").limitToLast(42).get().get();

assertEquals(
query(
limit(42),
order("foo", StructuredQuery.Direction.DESCENDING),
endAt(string("foo"), false),
startAt(string("bar"), true)),
runQuery.getValue());
}

@Test
public void limitToLastReversesResults() throws Exception {
doAnswer(queryResponse(DOCUMENT_NAME + "2", DOCUMENT_NAME + "1"))
.when(firestoreMock)
.streamRequest(
runQuery.capture(),
streamObserverCapture.capture(),
Matchers.<ServerStreamingCallable>any());

QuerySnapshot querySnapshot = query.orderBy("foo").limitToLast(2).get().get();

Iterator<QueryDocumentSnapshot> docIterator = querySnapshot.iterator();
assertEquals("doc1", docIterator.next().getId());
assertEquals("doc2", docIterator.next().getId());
}

@Test
public void limitToLastRequiresAtLeastOneOrderingConstraint() throws Exception {
doAnswer(queryResponse())
.when(firestoreMock)
.streamRequest(
runQuery.capture(),
streamObserverCapture.capture(),
Matchers.<ServerStreamingCallable>any());

try {
query.limitToLast(1).get().get();
fail("Expected exception");
} catch (IllegalStateException e) {
assertEquals(
e.getMessage(),
"limitToLast() queries require specifying at least one orderBy() clause.");
}
}

@Test
public void limitToLastRejectsStream() throws Exception {
doAnswer(queryResponse())
.when(firestoreMock)
.streamRequest(
runQuery.capture(),
streamObserverCapture.capture(),
Matchers.<ServerStreamingCallable>any());

try {
query.orderBy("foo").limitToLast(1).stream(null);
fail("Expected exception");
} catch (IllegalStateException e) {
assertEquals(
e.getMessage(),
"Query results for queries that include limitToLast() constraints cannot be streamed. "
+ "Use Query.get() instead.");
}
}

@Test
public void withOffset() throws Exception {
doAnswer(queryResponse())
Expand Down