Search in sources :

Example 1 with RunQueryResponse

use of com.google.datastore.v1.RunQueryResponse in project google-cloud-java by GoogleCloudPlatform.

the class DatastoreTest method testQueryPaginationWithLimit.

@Test
public void testQueryPaginationWithLimit() throws DatastoreException {
    List<RunQueryResponse> responses = buildResponsesForQueryPaginationWithLimit();
    List<ByteString> endCursors = Lists.newArrayListWithCapacity(responses.size());
    for (RunQueryResponse response : responses) {
        EasyMock.expect(rpcMock.runQuery(EasyMock.anyObject(RunQueryRequest.class))).andReturn(response);
        if (response.getBatch().getMoreResults() != QueryResultBatch.MoreResultsType.NOT_FINISHED) {
            endCursors.add(response.getBatch().getEndCursor());
        }
    }
    EasyMock.replay(rpcFactoryMock, rpcMock);
    Datastore datastore = rpcMockOptions.getService();
    int limit = 2;
    int totalCount = 0;
    Iterator<ByteString> cursorIter = endCursors.iterator();
    StructuredQuery<Entity> query = Query.newEntityQueryBuilder().setLimit(limit).build();
    while (true) {
        QueryResults<Entity> results = datastore.run(query);
        int resultCount = 0;
        while (results.hasNext()) {
            results.next();
            resultCount++;
            totalCount++;
        }
        assertTrue(cursorIter.hasNext());
        Cursor expectedEndCursor = Cursor.copyFrom(cursorIter.next().toByteArray());
        assertEquals(expectedEndCursor, results.getCursorAfter());
        if (resultCount < limit) {
            break;
        }
        query = query.toBuilder().setStartCursor(results.getCursorAfter()).build();
    }
    assertEquals(5, totalCount);
    EasyMock.verify(rpcFactoryMock, rpcMock);
}
Also used : RunQueryResponse(com.google.datastore.v1.RunQueryResponse) ByteString(com.google.protobuf.ByteString) RunQueryRequest(com.google.datastore.v1.RunQueryRequest) Test(org.junit.Test)

Example 2 with RunQueryResponse

use of com.google.datastore.v1.RunQueryResponse in project google-cloud-java by GoogleCloudPlatform.

the class DatastoreTest method buildResponsesForQueryPaginationWithLimit.

private List<RunQueryResponse> buildResponsesForQueryPaginationWithLimit() {
    Entity entity4 = Entity.newBuilder(KEY4).set("value", StringValue.of("value")).build();
    Entity entity5 = Entity.newBuilder(KEY5).set("value", "value").build();
    datastore.add(ENTITY3, entity4, entity5);
    DatastoreRpc datastoreRpc = datastore.getOptions().getDatastoreRpcV1();
    List<RunQueryResponse> responses = new ArrayList<>();
    Query<Entity> query = Query.newEntityQueryBuilder().build();
    RunQueryRequest.Builder requestPb = RunQueryRequest.newBuilder();
    query.populatePb(requestPb);
    QueryResultBatch queryResultBatchPb = RunQueryResponse.newBuilder().mergeFrom(datastoreRpc.runQuery(requestPb.build())).getBatch();
    QueryResultBatch queryResultBatchPb1 = QueryResultBatch.newBuilder().mergeFrom(queryResultBatchPb).setMoreResults(QueryResultBatch.MoreResultsType.NOT_FINISHED).clearEntityResults().addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(0, 1)).setEndCursor(ByteString.copyFromUtf8("a")).build();
    responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb1).build());
    QueryResultBatch queryResultBatchPb2 = QueryResultBatch.newBuilder().mergeFrom(queryResultBatchPb).setMoreResults(QueryResultBatch.MoreResultsType.MORE_RESULTS_AFTER_LIMIT).clearEntityResults().addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(1, 2)).setEndCursor(// test invalid UTF-8 string
    ByteString.copyFrom(new byte[] { (byte) 0x80 })).build();
    responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb2).build());
    QueryResultBatch queryResultBatchPb3 = QueryResultBatch.newBuilder().mergeFrom(queryResultBatchPb).setMoreResults(QueryResultBatch.MoreResultsType.MORE_RESULTS_AFTER_LIMIT).clearEntityResults().addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(2, 4)).setEndCursor(ByteString.copyFromUtf8("b")).build();
    responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb3).build());
    QueryResultBatch queryResultBatchPb4 = QueryResultBatch.newBuilder().mergeFrom(queryResultBatchPb).setMoreResults(QueryResultBatch.MoreResultsType.NO_MORE_RESULTS).clearEntityResults().addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(4, 5)).setEndCursor(ByteString.copyFromUtf8("c")).build();
    responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb4).build());
    return responses;
}
Also used : QueryResultBatch(com.google.datastore.v1.QueryResultBatch) RunQueryResponse(com.google.datastore.v1.RunQueryResponse) RunQueryRequest(com.google.datastore.v1.RunQueryRequest) DatastoreRpc(com.google.cloud.datastore.spi.v1.DatastoreRpc) ArrayList(java.util.ArrayList)

Example 3 with RunQueryResponse

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

Example 4 with RunQueryResponse

use of com.google.datastore.v1.RunQueryResponse 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 5 with RunQueryResponse

use of com.google.datastore.v1.RunQueryResponse 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)

Aggregations

RunQueryResponse (com.google.datastore.v1.RunQueryResponse)6 Test (org.junit.Test)6 RunQueryRequest (com.google.datastore.v1.RunQueryRequest)5 RunQueryResponse (com.google.firestore.v1.RunQueryResponse)4 QueryResultBatch (com.google.datastore.v1.QueryResultBatch)3 RunQueryRequest (com.google.firestore.v1.RunQueryRequest)3 RunQueryFn (org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.RunQueryFn)3 Document (com.google.firestore.v1.Document)2 FieldReference (com.google.firestore.v1.StructuredQuery.FieldReference)2 ArrayList (java.util.ArrayList)2 ServerStream (com.google.api.gax.rpc.ServerStream)1 ServerStreamingCallable (com.google.api.gax.rpc.ServerStreamingCallable)1 DatastoreRpc (com.google.cloud.datastore.spi.v1.DatastoreRpc)1 FirestoreStub (com.google.cloud.firestore.v1.stub.FirestoreStub)1 Entity (com.google.datastore.v1.Entity)1 EntityResult (com.google.datastore.v1.EntityResult)1 GqlQuery (com.google.datastore.v1.GqlQuery)1 PartitionId (com.google.datastore.v1.PartitionId)1 Query (com.google.datastore.v1.Query)1 Datastore (com.google.datastore.v1.client.Datastore)1