Search in sources :

Example 6 with Query

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

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

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

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

the class DatastoreTest method buildResponsesForQueryPagination.

private List<RunQueryResponse> buildResponsesForQueryPagination() {
    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);
    List<RunQueryResponse> responses = new ArrayList<>();
    Query<Key> query = Query.newKeyQueryBuilder().build();
    RunQueryRequest.Builder requestPb = RunQueryRequest.newBuilder();
    query.populatePb(requestPb);
    QueryResultBatch queryResultBatchPb = RunQueryResponse.newBuilder().mergeFrom(((DatastoreImpl) datastore).runQuery(requestPb.build())).getBatch();
    QueryResultBatch queryResultBatchPb1 = QueryResultBatch.newBuilder().mergeFrom(queryResultBatchPb).setMoreResults(QueryResultBatch.MoreResultsType.NOT_FINISHED).clearEntityResults().addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(0, 1)).setEndCursor(queryResultBatchPb.getEntityResultsList().get(0).getCursor()).build();
    responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb1).build());
    QueryResultBatch queryResultBatchPb2 = QueryResultBatch.newBuilder().mergeFrom(queryResultBatchPb).setMoreResults(QueryResultBatch.MoreResultsType.NOT_FINISHED).clearEntityResults().addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(1, 3)).setEndCursor(queryResultBatchPb.getEntityResultsList().get(2).getCursor()).build();
    responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb2).build());
    QueryResultBatch queryResultBatchPb3 = QueryResultBatch.newBuilder().mergeFrom(queryResultBatchPb).setMoreResults(QueryResultBatch.MoreResultsType.NO_MORE_RESULTS).clearEntityResults().addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(3, 5)).setEndCursor(queryResultBatchPb.getEntityResultsList().get(4).getCursor()).build();
    responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb3).build());
    return responses;
}
Also used : QueryResultBatch(com.google.datastore.v1.QueryResultBatch) RunQueryResponse(com.google.datastore.v1.RunQueryResponse) RunQueryRequest(com.google.datastore.v1.RunQueryRequest) ArrayList(java.util.ArrayList)

Example 10 with Query

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

the class DatastoreTest method testEventualConsistencyQuery.

@Test
public void testEventualConsistencyQuery() {
    ReadOptions readOption = ReadOptions.newBuilder().setReadConsistencyValue(ReadConsistency.EVENTUAL_VALUE).build();
    com.google.datastore.v1.GqlQuery query = com.google.datastore.v1.GqlQuery.newBuilder().setQueryString("FROM * SELECT *").build();
    RunQueryRequest.Builder expectedRequest = RunQueryRequest.newBuilder().setReadOptions(readOption).setGqlQuery(query).setPartitionId(PartitionId.newBuilder().setProjectId(PROJECT_ID).build());
    EasyMock.expect(rpcMock.runQuery(expectedRequest.build())).andReturn(RunQueryResponse.newBuilder().build());
    EasyMock.replay(rpcFactoryMock, rpcMock);
    Datastore datastore = rpcMockOptions.getService();
    datastore.run(Query.newGqlQueryBuilder("FROM * SELECT *").build(), ReadOption.eventualConsistency());
    EasyMock.verify(rpcFactoryMock, rpcMock);
}
Also used : ReadOptions(com.google.datastore.v1.ReadOptions) RunQueryRequest(com.google.datastore.v1.RunQueryRequest) Test(org.junit.Test)

Aggregations

Test (org.junit.Test)11 Query (com.google.datastore.v1.Query)10 RunQueryRequest (com.google.datastore.v1.RunQueryRequest)9 GqlQuery (com.google.datastore.v1.GqlQuery)8 RunQueryResponse (com.google.datastore.v1.RunQueryResponse)6 Datastore (com.google.datastore.v1.client.Datastore)4 SplitQueryFn (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.Read.SplitQueryFn)4 KV (org.apache.beam.sdk.values.KV)4 Entity (com.google.datastore.v1.Entity)2 PartitionId (com.google.datastore.v1.PartitionId)2 QueryResultBatch (com.google.datastore.v1.QueryResultBatch)2 ArrayList (java.util.ArrayList)2 DatastoreRpc (com.google.cloud.datastore.spi.v1.DatastoreRpc)1 ReadOptions (com.google.datastore.v1.ReadOptions)1 DatastoreException (com.google.datastore.v1.client.DatastoreException)1 ByteString (com.google.protobuf.ByteString)1 Pipeline (org.apache.beam.sdk.Pipeline)1 DeleteEntity (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DeleteEntity)1 ReadFn (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.Read.ReadFn)1 TestPipeline (org.apache.beam.sdk.testing.TestPipeline)1