use of com.google.firestore.v1.RunQueryResponse 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);
}
use of com.google.firestore.v1.RunQueryResponse 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);
}
use of com.google.firestore.v1.RunQueryResponse 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;
}
use of com.google.firestore.v1.RunQueryResponse 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.firestore.v1.RunQueryResponse in project beam by apache.
the class DatastoreV1Test method makeLatestTimestampResponse.
/**
* Builds a response of the given timestamp.
*/
private static RunQueryResponse makeLatestTimestampResponse(long timestamp) {
RunQueryResponse.Builder timestampResponse = RunQueryResponse.newBuilder();
Entity.Builder entity = Entity.newBuilder();
entity.setKey(makeKey("dummyKind", "dummyId"));
entity.putProperties("timestamp", makeValue(new Date(timestamp * 1000)).build());
EntityResult.Builder entityResult = EntityResult.newBuilder();
entityResult.setEntity(entity);
QueryResultBatch.Builder batch = QueryResultBatch.newBuilder();
batch.addEntityResults(entityResult);
timestampResponse.setBatch(batch);
return timestampResponse.build();
}
Aggregations