Search in sources :

Example 11 with Query

use of com.google.datastore.v1.Query in project beam by apache.

the class DatastoreV1Test method testTranslateGqlQueryWithLimit.

@Test
public void testTranslateGqlQueryWithLimit() throws Exception {
    String gql = "SELECT * from DummyKind LIMIT 10";
    String gqlWithZeroLimit = gql + " LIMIT 0";
    GqlQuery gqlQuery = GqlQuery.newBuilder().setQueryString(gql).setAllowLiterals(true).build();
    GqlQuery gqlQueryWithZeroLimit = GqlQuery.newBuilder().setQueryString(gqlWithZeroLimit).setAllowLiterals(true).build();
    RunQueryRequest gqlRequest = makeRequest(gqlQuery, V_1_OPTIONS.getNamespace());
    RunQueryRequest gqlRequestWithZeroLimit = makeRequest(gqlQueryWithZeroLimit, V_1_OPTIONS.getNamespace());
    when(mockDatastore.runQuery(gqlRequestWithZeroLimit)).thenThrow(new DatastoreException("runQuery", Code.INVALID_ARGUMENT, "invalid query", // dummy
    new RuntimeException()));
    when(mockDatastore.runQuery(gqlRequest)).thenReturn(RunQueryResponse.newBuilder().setQuery(QUERY).build());
    assertEquals(translateGqlQueryWithLimitCheck(gql, mockDatastore, V_1_OPTIONS.getNamespace()), QUERY);
    verify(mockDatastore, times(1)).runQuery(gqlRequest);
    verify(mockDatastore, times(1)).runQuery(gqlRequestWithZeroLimit);
}
Also used : RunQueryRequest(com.google.datastore.v1.RunQueryRequest) DatastoreException(com.google.datastore.v1.client.DatastoreException) GqlQuery(com.google.datastore.v1.GqlQuery) Test(org.junit.Test)

Example 12 with Query

use of com.google.datastore.v1.Query in project beam by apache.

the class DatastoreV1Test method testReadValidationFailsQueryLimitNegative.

@Test
public void testReadValidationFailsQueryLimitNegative() throws Exception {
    Query invalidLimit = Query.newBuilder().setLimit(Int32Value.newBuilder().setValue(-5)).build();
    thrown.expect(IllegalArgumentException.class);
    thrown.expectMessage("Invalid query limit -5: must be positive");
    DatastoreIO.v1().read().withQuery(invalidLimit);
}
Also used : GqlQuery(com.google.datastore.v1.GqlQuery) Query(com.google.datastore.v1.Query) Test(org.junit.Test)

Example 13 with Query

use of com.google.datastore.v1.Query in project beam by apache.

the class DatastoreV1Test method testEstimatedSizeBytes.

/**
   * Tests {@link DatastoreV1.Read#getEstimatedSizeBytes} to fetch and return estimated size for a
   * query.
   */
@Test
public void testEstimatedSizeBytes() throws Exception {
    long entityBytes = 100L;
    // 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);
    assertEquals(entityBytes, getEstimatedSizeBytes(mockDatastore, QUERY, NAMESPACE));
    verify(mockDatastore, times(1)).runQuery(latestTimestampRequest);
    verify(mockDatastore, times(1)).runQuery(statRequest);
}
Also used : RunQueryResponse(com.google.datastore.v1.RunQueryResponse) RunQueryRequest(com.google.datastore.v1.RunQueryRequest) Test(org.junit.Test)

Example 14 with Query

use of com.google.datastore.v1.Query in project beam by apache.

the class V1ReadIT method testE2EV1Read.

/**
   * An end-to-end test for {@link DatastoreV1.Read#withQuery(Query)}
   *
   * <p>Write some test entities to datastore and then run a pipeline that
   * reads and counts the total number of entities. Verify that the count matches the
   * number of entities written.
   */
@Test
public void testE2EV1Read() throws Exception {
    // Read from datastore
    Query query = V1TestUtil.makeAncestorKindQuery(options.getKind(), options.getNamespace(), ancestor);
    DatastoreV1.Read read = DatastoreIO.v1().read().withProjectId(project).withQuery(query).withNamespace(options.getNamespace());
    // Count the total number of entities
    Pipeline p = Pipeline.create(options);
    PCollection<Long> count = p.apply(read).apply(Count.<Entity>globally());
    PAssert.thatSingleton(count).isEqualTo(numEntities);
    p.run();
}
Also used : Query(com.google.datastore.v1.Query) TestPipeline(org.apache.beam.sdk.testing.TestPipeline) Pipeline(org.apache.beam.sdk.Pipeline) Test(org.junit.Test)

Example 15 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)

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