Examples with RunQueryRequest com.google.firestore.v1.RunQueryRequest used on opensource projects

Example 6 with RunQueryRequest

use of com.google.firestore.v1.RunQueryRequest in project beam by apache.

the class DatastoreV1Test method testReadFnRetriesErrors.

/**
 * Tests that {@link ReadFn} retries after an error.
 */
@Test
public void testReadFnRetriesErrors() throws Exception {
    // An empty query to read entities.
    Query query = Query.newBuilder().setLimit(Int32Value.newBuilder().setValue(1)).build();
    // Use mockResponseForQuery to generate results.
    when(mockDatastore.runQuery(any(RunQueryRequest.class))).thenThrow(new DatastoreException("RunQuery", Code.DEADLINE_EXCEEDED, "", null)).thenAnswer(invocationOnMock -> {
        Query q = ((RunQueryRequest) invocationOnMock.getArguments()[0]).getQuery();
        return mockResponseForQuery(q);
    });
    ReadFn readFn = new ReadFn(V_1_OPTIONS, mockDatastoreFactory);
    DoFnTester<Query, Entity> doFnTester = DoFnTester.of(readFn);
    doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
    doFnTester.processBundle(query);
    verifyMetricWasSet("BatchDatastoreRead", "ok", NAMESPACE, 1);
    verifyMetricWasSet("BatchDatastoreRead", "unknown", NAMESPACE, 1);
}

Example 7 with RunQueryRequest

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

Example 8 with RunQueryRequest

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

Example 9 with RunQueryRequest

use of com.google.firestore.v1.RunQueryRequest in project beam by apache.

the class DatastoreV1Test method readFnTest.

/**
 * Helper function to run a test reading from a {@link ReadFn}.
 */
private void readFnTest(int numEntities) throws Exception {
    // An empty query to read entities.
    Query query = Query.newBuilder().setLimit(Int32Value.newBuilder().setValue(numEntities)).build();
    // Use mockResponseForQuery to generate results.
    when(mockDatastore.runQuery(any(RunQueryRequest.class))).thenAnswer(invocationOnMock -> {
        Query q = ((RunQueryRequest) invocationOnMock.getArguments()[0]).getQuery();
        return mockResponseForQuery(q);
    });
    ReadFn readFn = new ReadFn(V_1_OPTIONS, mockDatastoreFactory);
    DoFnTester<Query, Entity> doFnTester = DoFnTester.of(readFn);
    /**
     * Although Datastore client is marked transient in {@link ReadFn}, when injected through mock
     * factory using a when clause for unit testing purposes, it is not serializable because it
     * doesn't have a no-arg constructor. Thus disabling the cloning to prevent the test object from
     * being serialized.
     */
    doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
    List<Entity> entities = doFnTester.processBundle(query);
    int expectedNumCallsToRunQuery = (int) Math.ceil((double) numEntities / QUERY_BATCH_LIMIT);
    verify(mockDatastore, times(expectedNumCallsToRunQuery)).runQuery(any(RunQueryRequest.class));
    // Validate the number of results.
    assertEquals(numEntities, entities.size());
}

Example 10 with RunQueryRequest

use of com.google.firestore.v1.RunQueryRequest in project beam by apache.

the class PartitionQueryResponseToRunQueryRequestTest method ensureCursorPairingWorks.

@Test
public void ensureCursorPairingWorks() {
    StructuredQuery query = StructuredQuery.newBuilder().addFrom(CollectionSelector.newBuilder().setAllDescendants(true).setCollectionId("c1").build()).build();
    Cursor cursor1 = referenceValueCursor("projects/p1/databases/d1/documents/c1/doc1");
    Cursor cursor2 = referenceValueCursor("projects/p1/databases/d1/documents/c1/doc2");
    Cursor cursor3 = referenceValueCursor("projects/p1/databases/d1/documents/c1/doc2/c2/doc2");
    List<StructuredQuery> expectedQueries = newArrayList(newQueryWithCursors(query, null, cursor1), newQueryWithCursors(query, cursor1, cursor2), newQueryWithCursors(query, cursor2, cursor3), newQueryWithCursors(query, cursor3, null));
    PartitionQueryPair partitionQueryPair = new PartitionQueryPair(PartitionQueryRequest.newBuilder().setStructuredQuery(query).build(), PartitionQueryResponse.newBuilder().addPartitions(cursor3).addPartitions(cursor1).addPartitions(cursor2).build());
    ArgumentCaptor<RunQueryRequest> captor = ArgumentCaptor.forClass(RunQueryRequest.class);
    when(processContext.element()).thenReturn(partitionQueryPair);
    doNothing().when(processContext).output(captor.capture());
    PartitionQueryResponseToRunQueryRequest fn = new PartitionQueryResponseToRunQueryRequest();
    fn.processElement(processContext);
    List<StructuredQuery> actualQueries = captor.getAllValues().stream().map(RunQueryRequest::getStructuredQuery).collect(Collectors.toList());
    assertEquals(expectedQueries, actualQueries);
}