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);
}
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);
}
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);
}
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();
}
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;
}
Aggregations