Examples with Query org.molgenis.emx2.Query used on opensource projects

Example 21 with Query

use of org.molgenis.emx2.Query 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 22 with Query

use of org.molgenis.emx2.Query in project beam by apache.

the class DatastoreV1Test method testReadValidationFailsQueryLimitZero.

@Test
public void testReadValidationFailsQueryLimitZero() throws Exception {
    Query invalidLimit = Query.newBuilder().setLimit(Int32Value.newBuilder().setValue(0)).build();
    thrown.expect(IllegalArgumentException.class);
    thrown.expectMessage("Invalid query limit 0: must be positive");
    DatastoreIO.v1().read().withQuery(invalidLimit);
}

Example 23 with Query

use of org.molgenis.emx2.Query 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);
}

Example 24 with Query

use of org.molgenis.emx2.Query in project beam by apache.

the class DatastoreV1Test method testSplitQueryFnWithNumSplits.

/**
 * Tests {@link SplitQueryFn} when number of query splits is specified.
 */
@Test
public void testSplitQueryFnWithNumSplits() throws Exception {
    int numSplits = 100;
    when(mockQuerySplitter.getSplits(eq(QUERY), any(PartitionId.class), eq(numSplits), any(Datastore.class))).thenReturn(splitQuery(QUERY, numSplits));
    SplitQueryFn splitQueryFn = new SplitQueryFn(V_1_OPTIONS, numSplits, mockDatastoreFactory);
    DoFnTester<Query, Query> doFnTester = DoFnTester.of(splitQueryFn);
    /**
     * Although Datastore client is marked transient in {@link SplitQueryFn}, 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 doFn from being
     * serialized.
     */
    doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
    List<Query> queries = doFnTester.processBundle(QUERY);
    assertEquals(queries.size(), numSplits);
    // Confirms that sub-queries are not equal to original when there is more than one split.
    for (Query subQuery : queries) {
        assertNotEquals(subQuery, QUERY);
    }
    verify(mockQuerySplitter, times(1)).getSplits(eq(QUERY), any(PartitionId.class), eq(numSplits), any(Datastore.class));
    verifyZeroInteractions(mockDatastore);
}

Example 25 with Query

use of org.molgenis.emx2.Query in project beam by apache.

the class SplitQueryFnIT method testSplitQueryFn.

/**
 * A helper method to test {@link SplitQueryFn} to generate the expected number of splits.
 */
private void testSplitQueryFn(String projectId, String kind, @Nullable String namespace, int expectedNumSplits) throws Exception {
    Query.Builder query = Query.newBuilder();
    query.addKindBuilder().setName(kind);
    SplitQueryFn splitQueryFn = new SplitQueryFn(V1Options.from(projectId, namespace, null), 0);
    DoFnTester<Query, Query> doFnTester = DoFnTester.of(splitQueryFn);
    List<Query> queries = doFnTester.processBundle(query.build());
    assertEquals(expectedNumSplits, queries.size());
}