Skip to content

Commit c104615

Browse files
schmidt-sebastianschmidt-sebastian
authored
feat: add Query.limitToLast() (#151)
1 parent 8ca0ea8 commit c104615

File tree

4 files changed

+298
-94
lines changed

4 files changed

+298
-94
lines changed

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

Lines changed: 96 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
package com.google.cloud.firestore;
1818

19+
import static com.google.common.collect.Lists.reverse;
1920
import static com.google.firestore.v1.StructuredQuery.FieldFilter.Operator.ARRAY_CONTAINS;
2021
import static com.google.firestore.v1.StructuredQuery.FieldFilter.Operator.ARRAY_CONTAINS_ANY;
2122
import static com.google.firestore.v1.StructuredQuery.FieldFilter.Operator.EQUAL;
@@ -182,6 +183,12 @@ Order toProto() {
182183
}
183184
}
184185

186+
/** Denotes whether a provided limit is applied to the beginning or the end of the result set. */
187+
enum LimitType {
188+
First,
189+
Last
190+
}
191+
185192
/** Options that define a Firestore Query. */
186193
@AutoValue
187194
abstract static class QueryOptions {
@@ -194,6 +201,8 @@ abstract static class QueryOptions {
194201

195202
abstract @Nullable Integer getLimit();
196203

204+
abstract LimitType getLimitType();
205+
197206
abstract @Nullable Integer getOffset();
198207

199208
abstract @Nullable Cursor getStartCursor();
@@ -209,6 +218,7 @@ abstract static class QueryOptions {
209218
static Builder builder() {
210219
return new AutoValue_Query_QueryOptions.Builder()
211220
.setAllDescendants(false)
221+
.setLimitType(LimitType.First)
212222
.setFieldOrders(ImmutableList.<FieldOrder>of())
213223
.setFieldFilters(ImmutableList.<FieldFilter>of())
214224
.setFieldProjections(ImmutableList.<FieldReference>of());
@@ -226,6 +236,8 @@ abstract static class Builder {
226236

227237
abstract Builder setLimit(Integer value);
228238

239+
abstract Builder setLimitType(LimitType value);
240+
229241
abstract Builder setOffset(Integer value);
230242

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

766778
/**
767-
* Creates and returns a new Query that's additionally limited to only return up to the specified
768-
* number of documents.
779+
* Creates and returns a new Query that only returns the first matching documents.
769780
*
770781
* @param limit The maximum number of items to return.
771782
* @return The created Query.
772783
*/
773784
@Nonnull
774785
public Query limit(int limit) {
775-
return new Query(firestore, options.toBuilder().setLimit(limit).build());
786+
return new Query(
787+
firestore, options.toBuilder().setLimit(limit).setLimitType(LimitType.First).build());
788+
}
789+
790+
/**
791+
* Creates and returns a new Query that only returns the last matching documents.
792+
*
793+
* <p>You must specify at least one orderBy clause for limitToLast queries. Otherwise, an {@link
794+
* java.lang.IllegalStateException} is thrown during execution.
795+
*
796+
* <p>Results for limitToLast() queries are only available once all documents are received. Hence,
797+
* limitToLast() queries cannot be streamed via the {@link #stream(ApiStreamObserver)} API.
798+
*
799+
* @param limit the maximum number of items to return
800+
* @return the created Query
801+
*/
802+
@Nonnull
803+
public Query limitToLast(int limit) {
804+
return new Query(
805+
firestore, options.toBuilder().setLimit(limit).setLimitType(LimitType.Last).build());
776806
}
777807

778808
/**
@@ -1004,8 +1034,25 @@ StructuredQuery.Builder buildQuery() {
10041034

10051035
if (!options.getFieldOrders().isEmpty()) {
10061036
for (FieldOrder order : options.getFieldOrders()) {
1007-
structuredQuery.addOrderBy(order.toProto());
1037+
switch (options.getLimitType()) {
1038+
case First:
1039+
structuredQuery.addOrderBy(order.toProto());
1040+
break;
1041+
case Last:
1042+
// Flip the orderBy directions since we want the last results
1043+
order =
1044+
new FieldOrder(
1045+
order.fieldPath,
1046+
order.direction.equals(Direction.ASCENDING)
1047+
? Direction.DESCENDING
1048+
: Direction.ASCENDING);
1049+
structuredQuery.addOrderBy(order.toProto());
1050+
break;
1051+
}
10081052
}
1053+
} else if (LimitType.Last.equals(options.getLimitType())) {
1054+
throw new IllegalStateException(
1055+
"limitToLast() queries require specifying at least one orderBy() clause.");
10091056
}
10101057

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

10231070
if (options.getStartCursor() != null) {
1024-
structuredQuery.setStartAt(options.getStartCursor());
1071+
switch (options.getLimitType()) {
1072+
case First:
1073+
structuredQuery.setStartAt(options.getStartCursor());
1074+
break;
1075+
case Last:
1076+
// Swap the cursors to match the flipped query ordering.
1077+
Cursor cursor =
1078+
options
1079+
.getStartCursor()
1080+
.toBuilder()
1081+
.setBefore(!options.getStartCursor().getBefore())
1082+
.build();
1083+
structuredQuery.setEndAt(cursor);
1084+
break;
1085+
}
10251086
}
10261087

10271088
if (options.getEndCursor() != null) {
1028-
structuredQuery.setEndAt(options.getEndCursor());
1089+
switch (options.getLimitType()) {
1090+
case First:
1091+
structuredQuery.setEndAt(options.getEndCursor());
1092+
break;
1093+
case Last:
1094+
// Swap the cursors to match the flipped query ordering.
1095+
Cursor cursor =
1096+
options
1097+
.getEndCursor()
1098+
.toBuilder()
1099+
.setBefore(!options.getEndCursor().getBefore())
1100+
.build();
1101+
structuredQuery.setStartAt(cursor);
1102+
break;
1103+
}
10291104
}
10301105

10311106
return structuredQuery;
@@ -1037,7 +1112,12 @@ StructuredQuery.Builder buildQuery() {
10371112
* @param responseObserver The observer to be notified when results arrive.
10381113
*/
10391114
public void stream(@Nonnull final ApiStreamObserver<DocumentSnapshot> responseObserver) {
1040-
stream(
1115+
Preconditions.checkState(
1116+
!LimitType.Last.equals(Query.this.options.getLimitType()),
1117+
"Query results for queries that include limitToLast() constraints cannot be streamed. "
1118+
+ "Use Query.get() instead.");
1119+
1120+
internalStream(
10411121
new QuerySnapshotObserver() {
10421122
@Override
10431123
public void onNext(QueryDocumentSnapshot documentSnapshot) {
@@ -1073,7 +1153,7 @@ Timestamp getReadTime() {
10731153
}
10741154
}
10751155

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

1181-
stream(
1261+
internalStream(
11821262
new QuerySnapshotObserver() {
11831263
List<QueryDocumentSnapshot> documentSnapshots = new ArrayList<>();
11841264

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

11951275
@Override
11961276
public void onCompleted() {
1277+
// The results for limitToLast queries need to be flipped since we reversed the
1278+
// ordering constraints before sending the query to the backend.
1279+
List<QueryDocumentSnapshot> resultView =
1280+
LimitType.Last.equals(Query.this.options.getLimitType())
1281+
? reverse(documentSnapshots)
1282+
: documentSnapshots;
11971283
QuerySnapshot querySnapshot =
1198-
QuerySnapshot.withDocuments(Query.this, this.getReadTime(), documentSnapshots);
1284+
QuerySnapshot.withDocuments(Query.this, this.getReadTime(), resultView);
11991285
result.set(querySnapshot);
12001286
}
12011287
},

‎google-cloud-firestore/src/test/java/com/google/cloud/firestore/QueryTest.java‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import static com.google.cloud.firestore.LocalFirestoreHelper.reference;
3030
import static com.google.cloud.firestore.LocalFirestoreHelper.select;
3131
import static com.google.cloud.firestore.LocalFirestoreHelper.startAt;
32+
import static com.google.cloud.firestore.LocalFirestoreHelper.string;
3233
import static com.google.cloud.firestore.LocalFirestoreHelper.unaryFilter;
3334
import static org.junit.Assert.assertEquals;
3435
import static org.junit.Assert.assertFalse;
@@ -92,6 +93,96 @@ public void withLimit() throws Exception {
9293
assertEquals(query(limit(42)), runQuery.getValue());
9394
}
9495

96+
@Test
97+
public void limitToLastReversesOrderingConstraints() throws Exception {
98+
doAnswer(queryResponse())
99+
.when(firestoreMock)
100+
.streamRequest(
101+
runQuery.capture(),
102+
streamObserverCapture.capture(),
103+
Matchers.<ServerStreamingCallable>any());
104+
105+
query.orderBy("foo").limitToLast(42).get().get();
106+
107+
assertEquals(
108+
query(limit(42), order("foo", StructuredQuery.Direction.DESCENDING)), runQuery.getValue());
109+
}
110+
111+
@Test
112+
public void limitToLastReversesCursors() throws Exception {
113+
doAnswer(queryResponse())
114+
.when(firestoreMock)
115+
.streamRequest(
116+
runQuery.capture(),
117+
streamObserverCapture.capture(),
118+
Matchers.<ServerStreamingCallable>any());
119+
120+
query.orderBy("foo").startAt("foo").endAt("bar").limitToLast(42).get().get();
121+
122+
assertEquals(
123+
query(
124+
limit(42),
125+
order("foo", StructuredQuery.Direction.DESCENDING),
126+
endAt(string("foo"), false),
127+
startAt(string("bar"), true)),
128+
runQuery.getValue());
129+
}
130+
131+
@Test
132+
public void limitToLastReversesResults() throws Exception {
133+
doAnswer(queryResponse(DOCUMENT_NAME + "2", DOCUMENT_NAME + "1"))
134+
.when(firestoreMock)
135+
.streamRequest(
136+
runQuery.capture(),
137+
streamObserverCapture.capture(),
138+
Matchers.<ServerStreamingCallable>any());
139+
140+
QuerySnapshot querySnapshot = query.orderBy("foo").limitToLast(2).get().get();
141+
142+
Iterator<QueryDocumentSnapshot> docIterator = querySnapshot.iterator();
143+
assertEquals("doc1", docIterator.next().getId());
144+
assertEquals("doc2", docIterator.next().getId());
145+
}
146+
147+
@Test
148+
public void limitToLastRequiresAtLeastOneOrderingConstraint() throws Exception {
149+
doAnswer(queryResponse())
150+
.when(firestoreMock)
151+
.streamRequest(
152+
runQuery.capture(),
153+
streamObserverCapture.capture(),
154+
Matchers.<ServerStreamingCallable>any());
155+
156+
try {
157+
query.limitToLast(1).get().get();
158+
fail("Expected exception");
159+
} catch (IllegalStateException e) {
160+
assertEquals(
161+
e.getMessage(),
162+
"limitToLast() queries require specifying at least one orderBy() clause.");
163+
}
164+
}
165+
166+
@Test
167+
public void limitToLastRejectsStream() throws Exception {
168+
doAnswer(queryResponse())
169+
.when(firestoreMock)
170+
.streamRequest(
171+
runQuery.capture(),
172+
streamObserverCapture.capture(),
173+
Matchers.<ServerStreamingCallable>any());
174+
175+
try {
176+
query.orderBy("foo").limitToLast(1).stream(null);
177+
fail("Expected exception");
178+
} catch (IllegalStateException e) {
179+
assertEquals(
180+
e.getMessage(),
181+
"Query results for queries that include limitToLast() constraints cannot be streamed. "
182+
+ "Use Query.get() instead.");
183+
}
184+
}
185+
95186
@Test
96187
public void withOffset() throws Exception {
97188
doAnswer(queryResponse())

0 commit comments

Comments
 (0)