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

Example 1 with RunQueryFn

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

the class FirestoreV1FnRunQueryTest method endToEnd.

@Test
public void endToEnd() throws Exception {
    TestData testData = TestData.fieldEqualsBar().setProjectId(projectId).build();
    List<RunQueryResponse> responses = ImmutableList.of(testData.response1, testData.response2, testData.response3);
    when(responseStream1.iterator()).thenReturn(responses.iterator());
    when(callable.call(testData.request)).thenReturn(responseStream1);
    when(stub.runQueryCallable()).thenReturn(callable);
    when(ff.getFirestoreStub(any())).thenReturn(stub);
    RpcQosOptions options = RpcQosOptions.defaultOptions();
    when(ff.getRpcQos(any())).thenReturn(FirestoreStatefulComponentFactory.INSTANCE.getRpcQos(options));
    ArgumentCaptor<RunQueryResponse> responsesCaptor = ArgumentCaptor.forClass(RunQueryResponse.class);
    doNothing().when(processContext).output(responsesCaptor.capture());
    when(processContext.element()).thenReturn(testData.request);
    RunQueryFn fn = new RunQueryFn(clock, ff, options);
    runFunction(fn);
    List<RunQueryResponse> allValues = responsesCaptor.getAllValues();
    assertEquals(responses, allValues);
}

Example 2 with RunQueryFn

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

the class FirestoreV1FnRunQueryTest method resumeFromLastReadValue.

@Override
public void resumeFromLastReadValue() throws Exception {
    TestData testData = TestData.fieldEqualsBar().setProjectId(projectId).setOrderFunction(f -> Collections.singletonList(Order.newBuilder().setDirection(Direction.ASCENDING).setField(f).build())).build();
    RunQueryRequest request2 = RunQueryRequest.newBuilder().setParent(String.format("projects/%s/databases/(default)/document", projectId)).setStructuredQuery(testData.request.getStructuredQuery().toBuilder().setStartAt(Cursor.newBuilder().setBefore(false).addValues(Value.newBuilder().setStringValue("bar")))).build();
    List<RunQueryResponse> responses = ImmutableList.of(testData.response1, testData.response2, testData.response3);
    when(responseStream1.iterator()).thenReturn(new AbstractIterator<RunQueryResponse>() {

        private int invocationCount = 1;

        @Override
        protected RunQueryResponse computeNext() {
            int count = invocationCount++;
            if (count == 1) {
                return responses.get(0);
            } else if (count == 2) {
                return responses.get(1);
            } else {
                throw RETRYABLE_ERROR;
            }
        }
    });
    when(callable.call(testData.request)).thenReturn(responseStream1);
    doNothing().when(attempt).checkCanRetry(any(), eq(RETRYABLE_ERROR));
    when(responseStream2.iterator()).thenReturn(ImmutableList.of(responses.get(2)).iterator());
    when(callable.call(request2)).thenReturn(responseStream2);
    when(stub.runQueryCallable()).thenReturn(callable);
    when(ff.getFirestoreStub(any())).thenReturn(stub);
    when(ff.getRpcQos(any())).thenReturn(rpcQos);
    when(rpcQos.newReadAttempt(any())).thenReturn(attempt);
    when(attempt.awaitSafeToProceed(any())).thenReturn(true);
    ArgumentCaptor<RunQueryResponse> responsesCaptor = ArgumentCaptor.forClass(RunQueryResponse.class);
    doNothing().when(processContext).output(responsesCaptor.capture());
    when(processContext.element()).thenReturn(testData.request);
    RunQueryFn fn = new RunQueryFn(clock, ff, rpcQosOptions);
    runFunction(fn);
    List<RunQueryResponse> allValues = responsesCaptor.getAllValues();
    assertEquals(responses, allValues);
    verify(callable, times(1)).call(testData.request);
    verify(callable, times(1)).call(request2);
    verify(attempt, times(3)).recordStreamValue(any());
}

Example 3 with RunQueryFn

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

the class FirestoreV1FnRunQueryTest method resumeFromLastReadValue_withNoOrderBy.

@Test
public void resumeFromLastReadValue_withNoOrderBy() throws Exception {
    TestData testData = TestData.fieldEqualsBar().setProjectId(projectId).build();
    RunQueryRequest request2 = RunQueryRequest.newBuilder().setParent(String.format("projects/%s/databases/(default)/document", projectId)).setStructuredQuery(testData.request.getStructuredQuery().toBuilder().setStartAt(Cursor.newBuilder().setBefore(false).addValues(Value.newBuilder().setReferenceValue(testData.response2.getDocument().getName()))).addOrderBy(Order.newBuilder().setField(FieldReference.newBuilder().setFieldPath("__name__")).setDirection(Direction.ASCENDING))).build();
    List<RunQueryResponse> responses = ImmutableList.of(testData.response1, testData.response2, testData.response3);
    when(responseStream1.iterator()).thenReturn(new AbstractIterator<RunQueryResponse>() {

        private int invocationCount = 1;

        @Override
        protected RunQueryResponse computeNext() {
            int count = invocationCount++;
            if (count == 1) {
                return responses.get(0);
            } else if (count == 2) {
                return responses.get(1);
            } else {
                throw RETRYABLE_ERROR;
            }
        }
    });
    when(callable.call(testData.request)).thenReturn(responseStream1);
    doNothing().when(attempt).checkCanRetry(any(), eq(RETRYABLE_ERROR));
    when(responseStream2.iterator()).thenReturn(ImmutableList.of(testData.response3).iterator());
    when(callable.call(request2)).thenReturn(responseStream2);
    when(stub.runQueryCallable()).thenReturn(callable);
    when(ff.getFirestoreStub(any())).thenReturn(stub);
    when(ff.getRpcQos(any())).thenReturn(rpcQos);
    when(rpcQos.newReadAttempt(any())).thenReturn(attempt);
    when(attempt.awaitSafeToProceed(any())).thenReturn(true);
    ArgumentCaptor<RunQueryResponse> responsesCaptor = ArgumentCaptor.forClass(RunQueryResponse.class);
    doNothing().when(processContext).output(responsesCaptor.capture());
    when(processContext.element()).thenReturn(testData.request);
    RunQueryFn fn = new RunQueryFn(clock, ff, rpcQosOptions);
    runFunction(fn);
    List<RunQueryResponse> allValues = responsesCaptor.getAllValues();
    assertEquals(responses, allValues);
    verify(callable, times(1)).call(testData.request);
    verify(callable, times(1)).call(request2);
    verify(attempt, times(3)).recordStreamValue(any());
}