Examples with PartitionQueryPair org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.PartitionQueryPair used on opensource projects

Example 1 with PartitionQueryPair

use of org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.PartitionQueryPair 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);
}

Example 2 with PartitionQueryPair

use of org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.PartitionQueryPair in project beam by apache.

the class FirestoreV1FnPartitionQueryTest method endToEnd_emptyCursors.

@Test
public void endToEnd_emptyCursors() throws Exception {
    // First page of the response
    PartitionQueryRequest request1 = PartitionQueryRequest.newBuilder().setParent(String.format("projects/%s/databases/(default)/document", projectId)).build();
    PartitionQueryResponse response1 = PartitionQueryResponse.newBuilder().build();
    when(callable.call(request1)).thenReturn(pagedResponse1);
    when(page1.getResponse()).thenReturn(response1);
    when(pagedResponse1.iteratePages()).thenReturn(ImmutableList.of(page1));
    when(stub.partitionQueryPagedCallable()).thenReturn(callable);
    when(ff.getFirestoreStub(any())).thenReturn(stub);
    RpcQosOptions options = RpcQosOptions.defaultOptions();
    when(ff.getRpcQos(any())).thenReturn(FirestoreStatefulComponentFactory.INSTANCE.getRpcQos(options));
    ArgumentCaptor<PartitionQueryPair> responses = ArgumentCaptor.forClass(PartitionQueryPair.class);
    doNothing().when(processContext).output(responses.capture());
    when(processContext.element()).thenReturn(request1);
    PartitionQueryFn fn = new PartitionQueryFn(clock, ff, options);
    runFunction(fn);
    List<PartitionQueryPair> expected = newArrayList(new PartitionQueryPair(request1, response1));
    List<PartitionQueryPair> allValues = responses.getAllValues();
    assertEquals(expected, allValues);
}

Example 3 with PartitionQueryPair

use of org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.PartitionQueryPair in project beam by apache.

the class FirestoreV1FnPartitionQueryTest method endToEnd.

@Test
public void endToEnd() throws Exception {
    // 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"))).addPartitions(Cursor.newBuilder().addValues(Value.newBuilder().setReferenceValue("doc-400"))).build();
    when(callable.call(request1)).thenReturn(pagedResponse1);
    when(page1.getResponse()).thenReturn(response1);
    when(pagedResponse1.iteratePages()).thenReturn(ImmutableList.of(page1));
    when(stub.partitionQueryPagedCallable()).thenReturn(callable);
    when(ff.getFirestoreStub(any())).thenReturn(stub);
    RpcQosOptions options = RpcQosOptions.defaultOptions();
    when(ff.getRpcQos(any())).thenReturn(FirestoreStatefulComponentFactory.INSTANCE.getRpcQos(options));
    ArgumentCaptor<PartitionQueryPair> responses = ArgumentCaptor.forClass(PartitionQueryPair.class);
    doNothing().when(processContext).output(responses.capture());
    when(processContext.element()).thenReturn(request1);
    PartitionQueryFn fn = new PartitionQueryFn(clock, ff, options);
    runFunction(fn);
    List<PartitionQueryPair> expected = newArrayList(new PartitionQueryPair(request1, response1));
    List<PartitionQueryPair> allValues = responses.getAllValues();
    assertEquals(expected, allValues);
}

Example 4 with PartitionQueryPair

use of org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.PartitionQueryPair in project beam by apache.

the class PartitionQueryResponseToRunQueryRequestTest method ensureCursorPairingWorks.

@Test
public void ensureCursorPairingWorks() {
    StructuredQuery query = StructuredQuery.newBuilder().addFrom(CollectionSelector.newBuilder().setAllDescendants(true).setCollectionId("c1").build()).build();
    Cursor cursor1 = referenceValueCursor("projects/p1/databases/d1/documents/c1/doc1");
    Cursor cursor2 = referenceValueCursor("projects/p1/databases/d1/documents/c1/doc2");
    Cursor cursor3 = referenceValueCursor("projects/p1/databases/d1/documents/c1/doc2/c2/doc2");
    List<StructuredQuery> expectedQueries = newArrayList(newQueryWithCursors(query, null, cursor1), newQueryWithCursors(query, cursor1, cursor2), newQueryWithCursors(query, cursor2, cursor3), newQueryWithCursors(query, cursor3, null));
    PartitionQueryPair partitionQueryPair = new PartitionQueryPair(PartitionQueryRequest.newBuilder().setStructuredQuery(query).build(), PartitionQueryResponse.newBuilder().addPartitions(cursor3).addPartitions(cursor1).addPartitions(cursor2).build());
    ArgumentCaptor<RunQueryRequest> captor = ArgumentCaptor.forClass(RunQueryRequest.class);
    when(processContext.element()).thenReturn(partitionQueryPair);
    doNothing().when(processContext).output(captor.capture());
    PartitionQueryResponseToRunQueryRequest fn = new PartitionQueryResponseToRunQueryRequest();
    fn.processElement(processContext);
    List<StructuredQuery> actualQueries = captor.getAllValues().stream().map(RunQueryRequest::getStructuredQuery).collect(Collectors.toList());
    assertEquals(expectedQueries, actualQueries);
}

Example 5 with PartitionQueryPair

use of org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.PartitionQueryPair in project beam by apache.

the class PartitionQueryResponseToRunQueryRequestTest method ensureCursorPairingWorks_emptyCursorsInResponse.

@Test
public void ensureCursorPairingWorks_emptyCursorsInResponse() {
    StructuredQuery query = StructuredQuery.newBuilder().addFrom(CollectionSelector.newBuilder().setAllDescendants(true).setCollectionId("c1").build()).build();
    List<StructuredQuery> expectedQueries = newArrayList(query);
    PartitionQueryPair partitionQueryPair = new PartitionQueryPair(PartitionQueryRequest.newBuilder().setStructuredQuery(query).build(), PartitionQueryResponse.newBuilder().build());
    ArgumentCaptor<RunQueryRequest> captor = ArgumentCaptor.forClass(RunQueryRequest.class);
    when(processContext.element()).thenReturn(partitionQueryPair);
    doNothing().when(processContext).output(captor.capture());
    PartitionQueryResponseToRunQueryRequest fn = new PartitionQueryResponseToRunQueryRequest();
    fn.processElement(processContext);
    List<StructuredQuery> actualQueries = captor.getAllValues().stream().map(RunQueryRequest::getStructuredQuery).collect(Collectors.toList());
    assertEquals(expectedQueries, actualQueries);
}