Examples with PartitionQueryPagedResponse com.google.cloud.firestore.v1.FirestoreClient.PartitionQueryPagedResponse used on opensource projects

Example 1 with PartitionQueryPagedResponse

use of com.google.cloud.firestore.v1.FirestoreClient.PartitionQueryPagedResponse 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);
        }
    }
}

Example 2 with PartitionQueryPagedResponse

use of com.google.cloud.firestore.v1.FirestoreClient.PartitionQueryPagedResponse 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());
}

Example 3 with PartitionQueryPagedResponse

use of com.google.cloud.firestore.v1.FirestoreClient.PartitionQueryPagedResponse 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");
}

Example 4 with PartitionQueryPagedResponse

use of com.google.cloud.firestore.v1.FirestoreClient.PartitionQueryPagedResponse 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()));
}

Example 5 with PartitionQueryPagedResponse

use of com.google.cloud.firestore.v1.FirestoreClient.PartitionQueryPagedResponse in project java-firestore by googleapis.

the class PartitionQuery method requestsOneLessThanDesired.

@Test
public void requestsOneLessThanDesired() throws Exception {
    int desiredPartitionsCount = 2;
    PartitionQueryRequest expectedRequest = PartitionQueryRequest.newBuilder().setParent(DATABASE_NAME).setStructuredQuery(StructuredQuery.newBuilder().addFrom(StructuredQuery.CollectionSelector.newBuilder().setAllDescendants(true).setCollectionId("collectionId")).addOrderBy(StructuredQuery.Order.newBuilder().setField(StructuredQuery.FieldReference.newBuilder().setFieldPath("__name__")).setDirection(StructuredQuery.Direction.ASCENDING))).setPartitionCount(desiredPartitionsCount - 1).build();
    PartitionQueryResponse response = PartitionQueryResponse.newBuilder().addPartitions(CURSOR1).build();
    when(pagedResponse.iterateAll()).thenReturn(ImmutableList.of(CURSOR1));
    when(queryPage.getResponse()).thenReturn(response);
    doReturn(ApiFutures.immediateFuture(pagedResponse)).when(firestoreMock).sendRequest(requestCaptor.capture(), Matchers.<UnaryCallable<PartitionQueryRequest, PartitionQueryPagedResponse>>any());
    firestoreMock.collectionGroup("collectionId").getPartitions(desiredPartitionsCount).get();
    PartitionQueryRequest actualRequest = requestCaptor.getValue();
    assertEquals(actualRequest, expectedRequest);
}