Examples with Query org.skife.jdbi.v2.Query used on opensource projects

Example 16 with Query

use of org.skife.jdbi.v2.Query in project beam by apache.

the class V1ReadIT method testE2EV1Read.

/**
   * An end-to-end test for {@link DatastoreV1.Read#withQuery(Query)}
   *
   * <p>Write some test entities to datastore and then run a pipeline that
   * reads and counts the total number of entities. Verify that the count matches the
   * number of entities written.
   */
@Test
public void testE2EV1Read() throws Exception {
    // Read from datastore
    Query query = V1TestUtil.makeAncestorKindQuery(options.getKind(), options.getNamespace(), ancestor);
    DatastoreV1.Read read = DatastoreIO.v1().read().withProjectId(project).withQuery(query).withNamespace(options.getNamespace());
    // Count the total number of entities
    Pipeline p = Pipeline.create(options);
    PCollection<Long> count = p.apply(read).apply(Count.<Entity>globally());
    PAssert.thatSingleton(count).isEqualTo(numEntities);
    p.run();
}

Example 17 with Query

use of org.skife.jdbi.v2.Query 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(new Answer<RunQueryResponse>() {

        @Override
        public RunQueryResponse answer(InvocationOnMock invocationOnMock) throws Throwable {
            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());
}