Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
feat: add Query.limitToLast() (#151)
  • Loading branch information
schmidt-sebastian committed Mar 27, 2020
1 parent 8ca0ea8 commit c104615
Show file tree
Hide file tree
Showing 4 changed files with 298 additions and 94 deletions.
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;
}
}

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;
}
}

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.");

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

0 comments on commit c104615

Please sign in to comment.