Search in sources :

Example 1 with ListCollectionIdsRequest

use of com.google.firestore.v1.ListCollectionIdsRequest in project beam by apache.

the class FirestoreV1FnListCollectionIdsTest method endToEnd.

@Test
public void endToEnd() throws Exception {
    // First page of the response
    ListCollectionIdsRequest request1 = ListCollectionIdsRequest.newBuilder().setParent(String.format("projects/%s/databases/(default)/document", projectId)).build();
    ListCollectionIdsResponse response1 = ListCollectionIdsResponse.newBuilder().addCollectionIds("col_1-1").addCollectionIds("col_1-2").addCollectionIds("col_1-3").setNextPageToken("page2").build();
    when(page1.getNextPageToken()).thenReturn(response1.getNextPageToken());
    when(page1.getResponse()).thenReturn(response1);
    when(page1.hasNextPage()).thenReturn(true);
    // Second page of the response
    ListCollectionIdsResponse response2 = ListCollectionIdsResponse.newBuilder().addCollectionIds("col_2-1").build();
    when(page2.getResponse()).thenReturn(response2);
    when(page2.hasNextPage()).thenReturn(false);
    when(pagedResponse1.iteratePages()).thenReturn(ImmutableList.of(page1, page2));
    when(callable.call(request1)).thenReturn(pagedResponse1);
    when(stub.listCollectionIdsPagedCallable()).thenReturn(callable);
    when(ff.getFirestoreStub(any())).thenReturn(stub);
    RpcQosOptions options = RpcQosOptions.defaultOptions();
    when(ff.getRpcQos(any())).thenReturn(FirestoreStatefulComponentFactory.INSTANCE.getRpcQos(options));
    ArgumentCaptor<ListCollectionIdsResponse> responses = ArgumentCaptor.forClass(ListCollectionIdsResponse.class);
    doNothing().when(processContext).output(responses.capture());
    when(processContext.element()).thenReturn(request1);
    ListCollectionIdsFn fn = new ListCollectionIdsFn(clock, ff, options);
    runFunction(fn);
    List<ListCollectionIdsResponse> expected = newArrayList(response1, response2);
    List<ListCollectionIdsResponse> allValues = responses.getAllValues();
    assertEquals(expected, allValues);
}
Also used : ListCollectionIdsResponse(com.google.firestore.v1.ListCollectionIdsResponse) ListCollectionIdsRequest(com.google.firestore.v1.ListCollectionIdsRequest) ListCollectionIdsFn(org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.ListCollectionIdsFn) Test(org.junit.Test)

Example 2 with ListCollectionIdsRequest

use of com.google.firestore.v1.ListCollectionIdsRequest in project beam by apache.

the class FirestoreV1FnListCollectionIdsTest 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
    ListCollectionIdsRequest request1 = ListCollectionIdsRequest.newBuilder().setParent(String.format("projects/%s/databases/(default)/document", projectId)).build();
    ListCollectionIdsResponse response1 = ListCollectionIdsResponse.newBuilder().addCollectionIds("col_1-1").addCollectionIds("col_1-2").addCollectionIds("col_1-3").setNextPageToken("page2").build();
    when(page1.getNextPageToken()).thenReturn(response1.getNextPageToken());
    when(page1.getResponse()).thenReturn(response1);
    when(page1.hasNextPage()).thenReturn(true);
    when(callable.call(request1)).thenReturn(pagedResponse1);
    doNothing().when(attempt).checkCanRetry(any(), eq(RETRYABLE_ERROR));
    when(pagedResponse1.iteratePages()).thenAnswer(invocation -> new Iterable<ListCollectionIdsPage>() {

        @Override
        public Iterator<ListCollectionIdsPage> iterator() {
            return new AbstractIterator<ListCollectionIdsPage>() {

                private boolean first = true;

                @Override
                protected ListCollectionIdsPage computeNext() {
                    if (first) {
                        first = false;
                        return page1;
                    } else {
                        throw RETRYABLE_ERROR;
                    }
                }
            };
        }
    });
    // Second page of the response
    ListCollectionIdsRequest request2 = ListCollectionIdsRequest.newBuilder().setParent(String.format("projects/%s/databases/(default)/document", projectId)).setPageToken("page2").build();
    ListCollectionIdsResponse response2 = ListCollectionIdsResponse.newBuilder().addCollectionIds("col_2-1").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.listCollectionIdsPagedCallable()).thenReturn(callable);
    when(ff.getFirestoreStub(any())).thenReturn(stub);
    ArgumentCaptor<ListCollectionIdsResponse> responses = ArgumentCaptor.forClass(ListCollectionIdsResponse.class);
    doNothing().when(processContext).output(responses.capture());
    when(processContext.element()).thenReturn(request1);
    ListCollectionIdsFn fn = new ListCollectionIdsFn(clock, ff, rpcQosOptions);
    runFunction(fn);
    List<ListCollectionIdsResponse> expected = newArrayList(response1, response2);
    List<ListCollectionIdsResponse> allValues = responses.getAllValues();
    assertEquals(expected, allValues);
}
Also used : ListCollectionIdsResponse(com.google.firestore.v1.ListCollectionIdsResponse) ListCollectionIdsPage(com.google.cloud.firestore.v1.FirestoreClient.ListCollectionIdsPage) ListCollectionIdsRequest(com.google.firestore.v1.ListCollectionIdsRequest) Iterator(java.util.Iterator) AbstractIterator(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.AbstractIterator) ListCollectionIdsFn(org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.ListCollectionIdsFn)

Example 3 with ListCollectionIdsRequest

use of com.google.firestore.v1.ListCollectionIdsRequest in project beam by apache.

the class FirestoreTestingHelper method listCollectionIds.

Stream<String> listCollectionIds(String parent) {
    ListCollectionIdsRequest lcir = ListCollectionIdsRequest.newBuilder().setParent(parent).build();
    // LOGGER.debug("lcir = {}", lcir);
    ListCollectionIdsPagedResponse response = rpc.listCollectionIdsPagedCallable().call(lcir);
    return StreamSupport.stream(response.iteratePages().spliterator(), false).flatMap(page -> StreamSupport.stream(page.getValues().spliterator(), false)).map(colId -> String.format("%s/%s", parent, colId));
}
Also used : RunQueryRequest(com.google.firestore.v1.RunQueryRequest) BatchWriteRequest(com.google.firestore.v1.BatchWriteRequest) FieldReference(com.google.firestore.v1.StructuredQuery.FieldReference) Write(com.google.firestore.v1.Write) Direction(com.google.firestore.v1.StructuredQuery.Direction) LoggerFactory(org.slf4j.LoggerFactory) TimeoutException(java.util.concurrent.TimeoutException) ImmutableMap(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap) CollectionReference(com.google.cloud.firestore.CollectionReference) NonNull(edu.umd.cs.findbugs.annotations.NonNull) Iterables(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables) ZoneOffset(java.time.ZoneOffset) WriteResult(com.google.cloud.firestore.WriteResult) MoreExecutors(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.MoreExecutors) ApiFutures(com.google.api.core.ApiFutures) StructuredQuery(com.google.firestore.v1.StructuredQuery) Matchers.notNullValue(org.hamcrest.Matchers.notNullValue) ListDocumentsPagedResponse(com.google.cloud.firestore.v1.FirestoreClient.ListDocumentsPagedResponse) Description(org.junit.runner.Description) Instant(java.time.Instant) Collectors(java.util.stream.Collectors) ZoneId(java.time.ZoneId) ApiFuture(com.google.api.core.ApiFuture) ListDocumentsRequest(com.google.firestore.v1.ListDocumentsRequest) WriteBatch(com.google.cloud.firestore.WriteBatch) List(java.util.List) Stream(java.util.stream.Stream) Optional(java.util.Optional) Matchers.is(org.hamcrest.Matchers.is) ApiFunction(com.google.api.core.ApiFunction) IntStream(java.util.stream.IntStream) Statement(org.junit.runners.model.Statement) Firestore(com.google.cloud.firestore.Firestore) Assume.assumeThat(org.junit.Assume.assumeThat) TestRule(org.junit.rules.TestRule) Streams(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Streams) FirestoreOptions(com.google.cloud.firestore.FirestoreOptions) Retention(java.lang.annotation.Retention) ArrayList(java.util.ArrayList) ListCollectionIdsRequest(com.google.firestore.v1.ListCollectionIdsRequest) StreamSupport(java.util.stream.StreamSupport) ListCollectionIdsPagedResponse(com.google.cloud.firestore.v1.FirestoreClient.ListCollectionIdsPagedResponse) Logger(org.slf4j.Logger) Document(com.google.firestore.v1.Document) Target(java.lang.annotation.Target) ElementType(java.lang.annotation.ElementType) Projection(com.google.firestore.v1.StructuredQuery.Projection) ExecutionException(java.util.concurrent.ExecutionException) TimeUnit(java.util.concurrent.TimeUnit) FirestoreRpc(com.google.cloud.firestore.spi.v1.FirestoreRpc) DateTimeFormatter(java.time.format.DateTimeFormatter) Clock(java.time.Clock) RunQueryResponse(com.google.firestore.v1.RunQueryResponse) Order(com.google.firestore.v1.StructuredQuery.Order) DocumentReference(com.google.cloud.firestore.DocumentReference) CollectionSelector(com.google.firestore.v1.StructuredQuery.CollectionSelector) Collections(java.util.Collections) RetentionPolicy(java.lang.annotation.RetentionPolicy) ListCollectionIdsPagedResponse(com.google.cloud.firestore.v1.FirestoreClient.ListCollectionIdsPagedResponse) ListCollectionIdsRequest(com.google.firestore.v1.ListCollectionIdsRequest)

Aggregations

ListCollectionIdsRequest (com.google.firestore.v1.ListCollectionIdsRequest)3 ListCollectionIdsResponse (com.google.firestore.v1.ListCollectionIdsResponse)2 ListCollectionIdsFn (org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.ListCollectionIdsFn)2 ApiFunction (com.google.api.core.ApiFunction)1 ApiFuture (com.google.api.core.ApiFuture)1 ApiFutures (com.google.api.core.ApiFutures)1 CollectionReference (com.google.cloud.firestore.CollectionReference)1 DocumentReference (com.google.cloud.firestore.DocumentReference)1 Firestore (com.google.cloud.firestore.Firestore)1 FirestoreOptions (com.google.cloud.firestore.FirestoreOptions)1 WriteBatch (com.google.cloud.firestore.WriteBatch)1 WriteResult (com.google.cloud.firestore.WriteResult)1 FirestoreRpc (com.google.cloud.firestore.spi.v1.FirestoreRpc)1 ListCollectionIdsPage (com.google.cloud.firestore.v1.FirestoreClient.ListCollectionIdsPage)1 ListCollectionIdsPagedResponse (com.google.cloud.firestore.v1.FirestoreClient.ListCollectionIdsPagedResponse)1 ListDocumentsPagedResponse (com.google.cloud.firestore.v1.FirestoreClient.ListDocumentsPagedResponse)1 BatchWriteRequest (com.google.firestore.v1.BatchWriteRequest)1 Document (com.google.firestore.v1.Document)1 ListDocumentsRequest (com.google.firestore.v1.ListDocumentsRequest)1 RunQueryRequest (com.google.firestore.v1.RunQueryRequest)1