Search in sources :

Example 1 with RunQueryRequest

use of com.google.firestore.v1.RunQueryRequest 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());
}
Also used : ArgumentMatchers.any(org.mockito.ArgumentMatchers.any) RunQueryRequest(com.google.firestore.v1.RunQueryRequest) FieldReference(com.google.firestore.v1.StructuredQuery.FieldReference) ArgumentMatchers.eq(org.mockito.ArgumentMatchers.eq) Mock(org.mockito.Mock) Direction(com.google.firestore.v1.StructuredQuery.Direction) FirestoreStub(com.google.cloud.firestore.v1.stub.FirestoreStub) Operator(com.google.firestore.v1.StructuredQuery.FieldFilter.Operator) ImmutableMap(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap) Function(java.util.function.Function) AbstractIterator(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.AbstractIterator) ArgumentCaptor(org.mockito.ArgumentCaptor) Value(com.google.firestore.v1.Value) Mockito.verifyNoMoreInteractions(org.mockito.Mockito.verifyNoMoreInteractions) Objects.requireNonNull(java.util.Objects.requireNonNull) ServerStream(com.google.api.gax.rpc.ServerStream) ServerStreamingCallable(com.google.api.gax.rpc.ServerStreamingCallable) StructuredQuery(com.google.firestore.v1.StructuredQuery) Filter(com.google.firestore.v1.StructuredQuery.Filter) Document(com.google.firestore.v1.Document) Mockito.times(org.mockito.Mockito.times) Test(org.junit.Test) Mockito.doNothing(org.mockito.Mockito.doNothing) Mockito.when(org.mockito.Mockito.when) FieldFilter(com.google.firestore.v1.StructuredQuery.FieldFilter) Mockito.verify(org.mockito.Mockito.verify) RunQueryFn(org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.RunQueryFn) List(java.util.List) Cursor(com.google.firestore.v1.Cursor) ImmutableList(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList) RunQueryResponse(com.google.firestore.v1.RunQueryResponse) Order(com.google.firestore.v1.StructuredQuery.Order) CollectionSelector(com.google.firestore.v1.StructuredQuery.CollectionSelector) Collections(java.util.Collections) Assert.assertEquals(org.junit.Assert.assertEquals) RunQueryResponse(com.google.firestore.v1.RunQueryResponse) RunQueryRequest(com.google.firestore.v1.RunQueryRequest) RunQueryFn(org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.RunQueryFn)

Example 2 with RunQueryRequest

use of com.google.firestore.v1.RunQueryRequest 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());
}
Also used : RunQueryResponse(com.google.firestore.v1.RunQueryResponse) RunQueryRequest(com.google.firestore.v1.RunQueryRequest) RunQueryFn(org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.RunQueryFn) Test(org.junit.Test)

Example 3 with RunQueryRequest

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

the class FirestoreTestingHelper method listDocumentsViaQuery.

Stream<String> listDocumentsViaQuery(String collectionPath) {
    int index = collectionPath.lastIndexOf('/');
    String parent = collectionPath.substring(0, index);
    String collectionId = collectionPath.substring(index + 1);
    FieldReference nameField = FieldReference.newBuilder().setFieldPath("__name__").build();
    RunQueryRequest rqr = RunQueryRequest.newBuilder().setParent(parent).setStructuredQuery(StructuredQuery.newBuilder().addFrom(CollectionSelector.newBuilder().setCollectionId(collectionId)).addOrderBy(Order.newBuilder().setField(nameField).setDirection(Direction.ASCENDING).build()).setSelect(Projection.newBuilder().addFields(nameField).build())).build();
    return StreamSupport.stream(rpc.runQueryCallable().call(rqr).spliterator(), false).filter(RunQueryResponse::hasDocument).map(RunQueryResponse::getDocument).map(Document::getName);
}
Also used : FieldReference(com.google.firestore.v1.StructuredQuery.FieldReference) RunQueryResponse(com.google.firestore.v1.RunQueryResponse) RunQueryRequest(com.google.firestore.v1.RunQueryRequest) Document(com.google.firestore.v1.Document)

Example 4 with RunQueryRequest

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

the class DatastoreV1Test method testTranslateGqlQueryWithNoLimit.

@Test
public void testTranslateGqlQueryWithNoLimit() throws Exception {
    String gql = "SELECT * from DummyKind";
    String gqlWithZeroLimit = gql + " LIMIT 0";
    GqlQuery gqlQueryWithZeroLimit = GqlQuery.newBuilder().setQueryString(gqlWithZeroLimit).setAllowLiterals(true).build();
    RunQueryRequest gqlRequestWithZeroLimit = makeRequest(gqlQueryWithZeroLimit, V_1_OPTIONS.getNamespace());
    when(mockDatastore.runQuery(gqlRequestWithZeroLimit)).thenReturn(RunQueryResponse.newBuilder().setQuery(QUERY).build());
    assertEquals(translateGqlQueryWithLimitCheck(gql, mockDatastore, V_1_OPTIONS.getNamespace()), QUERY);
    verify(mockDatastore, times(1)).runQuery(gqlRequestWithZeroLimit);
}
Also used : RunQueryRequest(com.google.datastore.v1.RunQueryRequest) GqlQuery(com.google.datastore.v1.GqlQuery) Test(org.junit.Test)

Example 5 with RunQueryRequest

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

the class DatastoreV1Test method testSplitQueryFnWithoutNumSplits.

/**
 * Tests {@link SplitQueryFn} when no query splits is specified.
 */
@Test
public void testSplitQueryFnWithoutNumSplits() throws Exception {
    // Force SplitQueryFn to compute the number of query splits
    int numSplits = 0;
    int expectedNumSplits = 20;
    long entityBytes = expectedNumSplits * DEFAULT_BUNDLE_SIZE_BYTES;
    // In seconds
    long timestamp = 1234L;
    RunQueryRequest latestTimestampRequest = makeRequest(makeLatestTimestampQuery(NAMESPACE), NAMESPACE);
    RunQueryResponse latestTimestampResponse = makeLatestTimestampResponse(timestamp);
    // Per Kind statistics request and response
    RunQueryRequest statRequest = makeRequest(makeStatKindQuery(NAMESPACE, timestamp), NAMESPACE);
    RunQueryResponse statResponse = makeStatKindResponse(entityBytes);
    when(mockDatastore.runQuery(latestTimestampRequest)).thenReturn(latestTimestampResponse);
    when(mockDatastore.runQuery(statRequest)).thenReturn(statResponse);
    when(mockQuerySplitter.getSplits(eq(QUERY), any(PartitionId.class), eq(expectedNumSplits), any(Datastore.class))).thenReturn(splitQuery(QUERY, expectedNumSplits));
    SplitQueryFn splitQueryFn = new SplitQueryFn(V_1_OPTIONS, numSplits, mockDatastoreFactory);
    DoFnTester<Query, Query> doFnTester = DoFnTester.of(splitQueryFn);
    doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
    List<Query> queries = doFnTester.processBundle(QUERY);
    assertEquals(expectedNumSplits, queries.size());
    verify(mockQuerySplitter, times(1)).getSplits(eq(QUERY), any(PartitionId.class), eq(expectedNumSplits), any(Datastore.class));
    verify(mockDatastore, times(1)).runQuery(latestTimestampRequest);
    verify(mockDatastore, times(1)).runQuery(statRequest);
}
Also used : Datastore(com.google.datastore.v1.client.Datastore) GqlQuery(com.google.datastore.v1.GqlQuery) Query(com.google.datastore.v1.Query) RunQueryResponse(com.google.datastore.v1.RunQueryResponse) RunQueryRequest(com.google.datastore.v1.RunQueryRequest) SplitQueryFn(org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.Read.SplitQueryFn) PartitionId(com.google.datastore.v1.PartitionId) Test(org.junit.Test)

Aggregations

Test (org.junit.Test)9 RunQueryRequest (com.google.datastore.v1.RunQueryRequest)6 GqlQuery (com.google.datastore.v1.GqlQuery)5 RunQueryRequest (com.google.firestore.v1.RunQueryRequest)5 Query (com.google.datastore.v1.Query)3 RunQueryResponse (com.google.firestore.v1.RunQueryResponse)3 StructuredQuery (com.google.firestore.v1.StructuredQuery)3 Entity (com.google.datastore.v1.Entity)2 RunQueryResponse (com.google.datastore.v1.RunQueryResponse)2 DatastoreException (com.google.datastore.v1.client.DatastoreException)2 Cursor (com.google.firestore.v1.Cursor)2 Document (com.google.firestore.v1.Document)2 FieldReference (com.google.firestore.v1.StructuredQuery.FieldReference)2 DeleteEntity (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DeleteEntity)2 ReadFn (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.Read.ReadFn)2 PartitionQueryResponseToRunQueryRequest (org.apache.beam.sdk.io.gcp.firestore.FirestoreV1.PartitionQuery.PartitionQueryResponseToRunQueryRequest)2 PartitionQueryPair (org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.PartitionQueryPair)2 RunQueryFn (org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.RunQueryFn)2 ServerStream (com.google.api.gax.rpc.ServerStream)1 ServerStreamingCallable (com.google.api.gax.rpc.ServerStreamingCallable)1