use of com.google.firestore.v1.RunQueryRequest 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());
}
use of com.google.firestore.v1.RunQueryRequest 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());
}
use of com.google.firestore.v1.RunQueryRequest 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.RunQueryRequest 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);
}
use of com.google.firestore.v1.RunQueryRequest 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);
}
Aggregations