use of com.google.spanner.v1.PartitionQueryRequest in project java-firestore by googleapis.
the class CollectionGroup method getPartitions.
/**
* 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.
*
* @deprecated Please use {@link #getPartitions(long)} instead. All cursors will be loaded before
* any value will be provided to {@code observer}.
* @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, final ApiStreamObserver<QueryPartition> observer) {
if (desiredPartitionCount == 1) {
// Short circuit if the user only requested a single partition.
observer.onNext(new QueryPartition(partitionQuery, null, null));
} else {
PartitionQueryRequest request = buildRequest(desiredPartitionCount);
final PartitionQueryPagedResponse response;
final TraceUtil traceUtil = TraceUtil.getInstance();
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_PARTITIONQUERY);
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
response = ApiExceptions.callAndTranslateApiException(rpcContext.sendRequest(request, rpcContext.getClient().partitionQueryPagedCallable()));
consumePartitions(response, queryPartition -> {
observer.onNext(queryPartition);
return null;
});
observer.onCompleted();
} catch (ApiException exception) {
span.setStatus(Status.UNKNOWN.withDescription(exception.getMessage()));
throw FirestoreException.forApiException(exception);
} finally {
span.end(TraceUtil.END_SPAN_OPTIONS);
}
}
}
use of com.google.spanner.v1.PartitionQueryRequest in project java-firestore by googleapis.
the class PartitionQuery method convertsPartitionsToQueries.
@Test
public void convertsPartitionsToQueries() throws Exception {
int desiredPartitionsCount = 3;
PartitionQueryResponse response = PartitionQueryResponse.newBuilder().addPartitions(CURSOR1).build();
when(pagedResponse.iterateAll()).thenReturn(ImmutableList.of(CURSOR1, CURSOR2));
when(queryPage.getResponse()).thenReturn(response);
doReturn(ApiFutures.immediateFuture(pagedResponse)).when(firestoreMock).sendRequest(requestCaptor.capture(), Matchers.<UnaryCallable<PartitionQueryRequest, PartitionQueryPagedResponse>>any());
doAnswer(queryResponse()).when(firestoreMock).streamRequest(runQuery.capture(), streamObserverCapture.capture(), Matchers.<ServerStreamingCallable>any());
List<QueryPartition> partitions = firestoreMock.collectionGroup("collectionId").getPartitions(desiredPartitionsCount).get();
assertEquals(partitions.size(), 3);
for (QueryPartition partition : partitions) {
partition.createQuery().get();
}
assertEquals(runQuery.getAllValues().size(), 3);
assertFalse(runQuery.getAllValues().get(0).getStructuredQuery().hasStartAt());
assertEquals(runQuery.getAllValues().get(0).getStructuredQuery().getEndAt(), PARTITION1);
assertEquals(runQuery.getAllValues().get(1).getStructuredQuery().getStartAt(), PARTITION1);
assertEquals(runQuery.getAllValues().get(1).getStructuredQuery().getEndAt(), PARTITION2);
assertEquals(runQuery.getAllValues().get(2).getStructuredQuery().getStartAt(), PARTITION2);
assertFalse(runQuery.getAllValues().get(2).getStructuredQuery().hasEndAt());
}
use of com.google.spanner.v1.PartitionQueryRequest in project java-firestore by googleapis.
the class PartitionQuery method sortsPartitions.
@Test
public void sortsPartitions() throws Exception {
int desiredPartitionsCount = 3;
PartitionQueryResponse response = PartitionQueryResponse.newBuilder().addPartitions(CURSOR1).build();
when(pagedResponse.iterateAll()).thenReturn(ImmutableList.of(CURSOR2, CURSOR1));
when(queryPage.getResponse()).thenReturn(response);
doReturn(ApiFutures.immediateFuture(pagedResponse)).when(firestoreMock).sendRequest(requestCaptor.capture(), Matchers.<UnaryCallable<PartitionQueryRequest, PartitionQueryPagedResponse>>any());
List<QueryPartition> partitions = firestoreMock.collectionGroup("collectionId").getPartitions(desiredPartitionsCount).get();
assertEquals(((DocumentReference) partitions.get(0).getEndBefore()[0]).getId(), "doc1");
assertEquals(((DocumentReference) partitions.get(1).getEndBefore()[0]).getId(), "doc2");
}
use of com.google.spanner.v1.PartitionQueryRequest in project java-firestore by googleapis.
the class FirestoreClientTest method partitionQueryTest.
@Test
public void partitionQueryTest() throws Exception {
Cursor responsesElement = Cursor.newBuilder().build();
PartitionQueryResponse expectedResponse = PartitionQueryResponse.newBuilder().setNextPageToken("").addAllPartitions(Arrays.asList(responsesElement)).build();
mockFirestore.addResponse(expectedResponse);
PartitionQueryRequest request = PartitionQueryRequest.newBuilder().setParent("parent-995424086").setPartitionCount(-1738969222).setPageToken("pageToken873572522").setPageSize(883849137).build();
PartitionQueryPagedResponse pagedListResponse = client.partitionQuery(request);
List<Cursor> resources = Lists.newArrayList(pagedListResponse.iterateAll());
Assert.assertEquals(1, resources.size());
Assert.assertEquals(expectedResponse.getPartitionsList().get(0), resources.get(0));
List<AbstractMessage> actualRequests = mockFirestore.getRequests();
Assert.assertEquals(1, actualRequests.size());
PartitionQueryRequest actualRequest = ((PartitionQueryRequest) actualRequests.get(0));
Assert.assertEquals(request.getParent(), actualRequest.getParent());
Assert.assertEquals(request.getStructuredQuery(), actualRequest.getStructuredQuery());
Assert.assertEquals(request.getPartitionCount(), actualRequest.getPartitionCount());
Assert.assertEquals(request.getPageToken(), actualRequest.getPageToken());
Assert.assertEquals(request.getPageSize(), actualRequest.getPageSize());
Assert.assertTrue(channelProvider.isHeaderSent(ApiClientHeaderProvider.getDefaultApiClientHeaderKey(), GaxGrpcProperties.getDefaultApiClientHeaderPattern()));
}
use of com.google.spanner.v1.PartitionQueryRequest in project beam by apache.
the class FirestoreV1FnPartitionQueryTest method resumeFromLastReadValue.
@Override
public void resumeFromLastReadValue() throws Exception {
when(ff.getFirestoreStub(any())).thenReturn(stub);
when(ff.getRpcQos(any())).thenReturn(rpcQos);
when(rpcQos.newReadAttempt(any())).thenReturn(attempt);
when(attempt.awaitSafeToProceed(any())).thenReturn(true);
// First page of the response
PartitionQueryRequest request1 = PartitionQueryRequest.newBuilder().setParent(String.format("projects/%s/databases/(default)/document", projectId)).build();
PartitionQueryResponse response1 = PartitionQueryResponse.newBuilder().addPartitions(Cursor.newBuilder().addValues(Value.newBuilder().setReferenceValue("doc-100"))).addPartitions(Cursor.newBuilder().addValues(Value.newBuilder().setReferenceValue("doc-200"))).addPartitions(Cursor.newBuilder().addValues(Value.newBuilder().setReferenceValue("doc-300"))).setNextPageToken("page2").build();
when(page1.getResponse()).thenReturn(response1);
when(callable.call(request1)).thenReturn(pagedResponse1);
doNothing().when(attempt).checkCanRetry(any(), eq(RETRYABLE_ERROR));
when(pagedResponse1.iteratePages()).thenAnswer(invocation -> new Iterable<PartitionQueryPage>() {
@Override
public Iterator<PartitionQueryPage> iterator() {
return new AbstractIterator<PartitionQueryPage>() {
private boolean first = true;
@Override
protected PartitionQueryPage computeNext() {
if (first) {
first = false;
return page1;
} else {
throw RETRYABLE_ERROR;
}
}
};
}
});
// Second page of the response
PartitionQueryRequest request2 = PartitionQueryRequest.newBuilder().setParent(String.format("projects/%s/databases/(default)/document", projectId)).setPageToken("page2").build();
PartitionQueryResponse response2 = PartitionQueryResponse.newBuilder().addPartitions(Cursor.newBuilder().addValues(Value.newBuilder().setReferenceValue("doc-400"))).build();
PartitionQueryResponse expectedResponse = response1.toBuilder().clearNextPageToken().addAllPartitions(response2.getPartitionsList()).build();
when(page2.getResponse()).thenReturn(response2);
when(page2.hasNextPage()).thenReturn(false);
when(callable.call(request2)).thenReturn(pagedResponse2);
when(pagedResponse2.iteratePages()).thenReturn(ImmutableList.of(page2));
when(stub.partitionQueryPagedCallable()).thenReturn(callable);
when(ff.getFirestoreStub(any())).thenReturn(stub);
ArgumentCaptor<PartitionQueryPair> responses = ArgumentCaptor.forClass(PartitionQueryPair.class);
doNothing().when(processContext).output(responses.capture());
when(processContext.element()).thenReturn(request1);
PartitionQueryFn fn = new PartitionQueryFn(clock, ff, rpcQosOptions);
runFunction(fn);
List<PartitionQueryPair> expected = newArrayList(new PartitionQueryPair(request1, expectedResponse));
List<PartitionQueryPair> allValues = responses.getAllValues();
assertEquals(expected, allValues);
}
Aggregations