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

Example 6 with RunQueryRequest

use of com.google.datastore.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.datastore.v1.RunQueryRequest in project spring-cloud-gcp by spring-cloud.

the class FirestoreTemplateTests method findAllTest.

@Test
public void findAllTest() {
    mockRunQueryMethod();
    StepVerifier.create(this.firestoreTemplate.findAll(TestEntity.class)).expectNext(new TestEntity("e1", 100L), new TestEntity("e2", 200L)).verifyComplete();
    StructuredQuery structuredQuery = StructuredQuery.newBuilder().addFrom(StructuredQuery.CollectionSelector.newBuilder().setCollectionId("testEntities").build()).build();
    RunQueryRequest request = RunQueryRequest.newBuilder().setParent(this.parent).setStructuredQuery(structuredQuery).build();
    verify(this.firestoreStub, times(1)).runQuery(eq(request), any());
    verify(this.firestoreStub, times(1)).runQuery(any(), any());
}

Example 8 with RunQueryRequest

use of com.google.datastore.v1.RunQueryRequest in project spring-cloud-gcp by spring-cloud.

the class FirestoreTemplateTests method countTest.

@Test
public void countTest() {
    mockRunQueryMethod();
    StepVerifier.create(this.firestoreTemplate.count(TestEntity.class)).expectNext(2L).verifyComplete();
    StructuredQuery structuredQuery = StructuredQuery.newBuilder().addFrom(StructuredQuery.CollectionSelector.newBuilder().setCollectionId("testEntities").build()).setSelect(StructuredQuery.Projection.newBuilder().addFields(StructuredQuery.FieldReference.newBuilder().setFieldPath("__name__").build()).build()).build();
    RunQueryRequest request = RunQueryRequest.newBuilder().setParent(this.parent).setStructuredQuery(structuredQuery).build();
    verify(this.firestoreStub, times(1)).runQuery(eq(request), any());
    verify(this.firestoreStub, times(1)).runQuery(any(), any());
}

Example 9 with RunQueryRequest

use of com.google.datastore.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 10 with RunQueryRequest

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