com.google.api.services.datastore.DatastoreV1.Query

Here are the examples of the java api class com.google.api.services.datastore.DatastoreV1.Query taken from open source projects.

1. DatastoreWordCount#readDataFromDatastore()

View license
/**
   * An example that creates a pipeline to do DatastoreIO.Read from Datastore.
   */
public static void readDataFromDatastore(Options options) {
    Query query = makeAncestorKindQuery(options);
    // For Datastore sources, the read namespace can be set on the entire query.
    DatastoreIO.Source source = DatastoreIO.source().withDataset(options.getDataset()).withQuery(query).withNamespace(options.getNamespace());
    Pipeline p = Pipeline.create(options);
    p.apply("ReadShakespeareFromDatastore", Read.from(source)).apply("StringifyEntity", ParDo.of(new GetContentFn())).apply("CountWords", new WordCount.CountWords()).apply("PrintWordCount", MapElements.via(new WordCount.FormatAsTextFn())).apply("WriteLines", TextIO.Write.to(options.getOutput()).withNumShards(options.getNumShards()));
    p.run();
}

2. DatastoreIOTest#testQueryDoesNotSplitWithLimitSet()

View license
/**
   * Tests that a query with a user-provided limit field does not split, and does not even
   * interact with a query splitter.
   */
@Test
public void testQueryDoesNotSplitWithLimitSet() throws Exception {
    // Minimal query with a limit
    Query query = Query.newBuilder().setLimit(5).build();
    // Mock query splitter, should not be invoked.
    QuerySplitter splitter = mock(QuerySplitter.class);
    when(splitter.getSplits(any(Query.class), any(PartitionId.class), eq(2), any(Datastore.class))).thenThrow(new AssertionError("Splitter should not be invoked"));
    List<DatastoreIO.Source> bundles = initialSource.withQuery(query).withMockSplitter(splitter).splitIntoBundles(1024, testPipelineOptions(null));
    assertEquals(1, bundles.size());
    assertEquals(query, bundles.get(0).getQuery());
    verifyNoMoreInteractions(splitter);
}

3. DatastoreIOTest#testQuerySplitWithZeroSize()

View license
@Test
public void testQuerySplitWithZeroSize() throws Exception {
    KindExpression mykind = KindExpression.newBuilder().setName("mykind").build();
    Query query = Query.newBuilder().addKind(mykind).build();
    List<Query> mockSplits = Lists.newArrayList(Query.newBuilder().addKind(mykind).build());
    QuerySplitter splitter = mock(QuerySplitter.class);
    when(splitter.getSplits(any(Query.class), any(PartitionId.class), eq(1), any(Datastore.class))).thenReturn(mockSplits);
    DatastoreIO.Source io = initialSource.withQuery(query).withMockSplitter(splitter).withMockEstimateSizeBytes(0L);
    List<DatastoreIO.Source> bundles = io.splitIntoBundles(1024, testPipelineOptions(null));
    assertEquals(1, bundles.size());
    verify(splitter, never()).getSplits(any(Query.class), any(PartitionId.class), eq(1), any(Datastore.class));
    DatastoreIO.Source bundle = bundles.get(0);
    Query bundleQuery = bundle.getQuery();
    assertEquals("mykind", bundleQuery.getKind(0).getName());
    assertFalse(bundleQuery.hasFilter());
}

4. DatastoreIOTest#testQuerySplitNoWorkers()

View license
@Test
public void testQuerySplitNoWorkers() throws Exception {
    KindExpression mykind = KindExpression.newBuilder().setName("mykind").build();
    Query query = Query.newBuilder().addKind(mykind).build();
    List<Query> mockSplits = Lists.newArrayList(Query.newBuilder().addKind(mykind).build());
    QuerySplitter splitter = mock(QuerySplitter.class);
    when(splitter.getSplits(any(Query.class), any(PartitionId.class), eq(12), any(Datastore.class))).thenReturn(mockSplits);
    DatastoreIO.Source io = initialSource.withQuery(query).withMockSplitter(splitter).withMockEstimateSizeBytes(8 * 1024L);
    DatastoreIO.Source spiedIo = spy(io);
    when(spiedIo.getEstimatedSizeBytes(any(PipelineOptions.class))).thenThrow(new NoSuchElementException());
    List<DatastoreIO.Source> bundles = spiedIo.splitIntoBundles(1024, testPipelineOptions(0));
    assertEquals(1, bundles.size());
    verify(splitter, never()).getSplits(any(Query.class), any(PartitionId.class), eq(1), any(Datastore.class));
    DatastoreIO.Source bundle = bundles.get(0);
    Query bundleQuery = bundle.getQuery();
    assertEquals("mykind", bundleQuery.getKind(0).getName());
    assertFalse(bundleQuery.hasFilter());
}

5. DatastoreIOTest#testSourceValidationFailsQueryLimitZero()

View license
@Test
public void testSourceValidationFailsQueryLimitZero() throws Exception {
    Query invalidLimit = Query.newBuilder().setLimit(0).build();
    thrown.expect(IllegalArgumentException.class);
    thrown.expectMessage("Invalid query limit 0");
    DatastoreIO.source().withQuery(invalidLimit);
}

6. DatastoreIOTest#testSourceValidationFailsQueryLimitNegative()

View license
@Test
public void testSourceValidationFailsQueryLimitNegative() throws Exception {
    Query invalidLimit = Query.newBuilder().setLimit(-5).build();
    thrown.expect(IllegalArgumentException.class);
    thrown.expectMessage("Invalid query limit -5");
    DatastoreIO.source().withQuery(invalidLimit);
}

7. DatastoreIOTest#testQuerySplitBasic()

View license
@Test
public void testQuerySplitBasic() throws Exception {
    KindExpression mykind = KindExpression.newBuilder().setName("mykind").build();
    Query query = Query.newBuilder().addKind(mykind).build();
    List<Query> mockSplits = new ArrayList<>();
    for (int i = 0; i < 8; ++i) {
        mockSplits.add(Query.newBuilder().addKind(mykind).setFilter(DatastoreHelper.makeFilter("foo", PropertyFilter.Operator.EQUAL, Value.newBuilder().setIntegerValue(i).build())).build());
    }
    QuerySplitter splitter = mock(QuerySplitter.class);
    /* No namespace */
    PartitionId partition = PartitionId.newBuilder().build();
    when(splitter.getSplits(any(Query.class), eq(partition), eq(8), any(Datastore.class))).thenReturn(mockSplits);
    DatastoreIO.Source io = initialSource.withNamespace(null).withQuery(query).withMockSplitter(splitter).withMockEstimateSizeBytes(8 * 1024L);
    List<DatastoreIO.Source> bundles = io.splitIntoBundles(1024, testPipelineOptions(null));
    assertEquals(8, bundles.size());
    for (int i = 0; i < 8; ++i) {
        DatastoreIO.Source bundle = bundles.get(i);
        Query bundleQuery = bundle.getQuery();
        assertEquals("mykind", bundleQuery.getKind(0).getName());
        assertEquals(i, bundleQuery.getFilter().getPropertyFilter().getValue().getIntegerValue());
    }
}

8. DatastoreIOTest#testQuerySplitterThrows()

View license
/**
   * Tests that when {@link QuerySplitter} cannot split a query, {@link DatastoreIO} falls back to
   * a single split.
   */
@Test
public void testQuerySplitterThrows() throws Exception {
    // Mock query splitter that throws IllegalArgumentException
    IllegalArgumentException exception = new IllegalArgumentException("query not supported by splitter");
    QuerySplitter splitter = mock(QuerySplitter.class);
    when(splitter.getSplits(any(Query.class), any(PartitionId.class), any(Integer.class), any(Datastore.class))).thenThrow(exception);
    Query query = Query.newBuilder().addKind(KindExpression.newBuilder().setName("myKind")).build();
    List<DatastoreIO.Source> bundles = initialSource.withQuery(query).withMockSplitter(splitter).withMockEstimateSizeBytes(10240L).splitIntoBundles(1024, testPipelineOptions(null));
    assertEquals(1, bundles.size());
    assertEquals(query, bundles.get(0).getQuery());
    verify(splitter, times(1)).getSplits(any(Query.class), any(PartitionId.class), any(Integer.class), any(Datastore.class));
    logged.verifyWarn("Unable to parallelize the given query", exception);
}

9. DatastoreIOTest#testQuerySplitSizeUnavailable()

View license
@Test
public void testQuerySplitSizeUnavailable() throws Exception {
    KindExpression mykind = KindExpression.newBuilder().setName("mykind").build();
    Query query = Query.newBuilder().addKind(mykind).build();
    List<Query> mockSplits = new ArrayList<>();
    for (int i = 0; i < 2; i++) {
        mockSplits.add(Query.newBuilder().addKind(mykind).setFilter(DatastoreHelper.makeFilter("foo", PropertyFilter.Operator.EQUAL, Value.newBuilder().setIntegerValue(i).build())).build());
    }
    QuerySplitter splitter = mock(QuerySplitter.class);
    when(splitter.getSplits(any(Query.class), any(PartitionId.class), eq(2), any(Datastore.class))).thenReturn(mockSplits);
    DatastoreIO.Source io = initialSource.withQuery(query).withMockSplitter(splitter).withMockEstimateSizeBytes(8 * 1024L);
    DatastoreIO.Source spiedIo = spy(io);
    when(spiedIo.getEstimatedSizeBytes(any(PipelineOptions.class))).thenThrow(new IOException());
    List<DatastoreIO.Source> bundles = spiedIo.splitIntoBundles(1024, testPipelineOptions(2));
    assertEquals(2, bundles.size());
    for (int i = 0; i < 2; ++i) {
        DatastoreIO.Source bundle = bundles.get(i);
        Query bundleQuery = bundle.getQuery();
        assertEquals("mykind", bundleQuery.getKind(0).getName());
        assertEquals(i, bundleQuery.getFilter().getPropertyFilter().getValue().getIntegerValue());
    }
}

10. DatastoreIOTest#runQueryLimitReadTest()

View license
/** Helper function to run a test reading from a limited-result query. */
private void runQueryLimitReadTest(int numEntities) throws Exception {
    // An empty query to read entities.
    Query query = Query.newBuilder().setLimit(numEntities).build();
    DatastoreIO.Source source = DatastoreIO.source().withQuery(query).withDataset("mockDataset");
    // Use mockResponseForQuery to generate results.
    when(mockDatastore.runQuery(any(RunQueryRequest.class))).thenAnswer(new Answer<RunQueryResponse>() {

        @Override
        public RunQueryResponse answer(InvocationOnMock invocation) throws Throwable {
            Query q = ((RunQueryRequest) invocation.getArguments()[0]).getQuery();
            return mockResponseForQuery(q);
        }
    });
    // Actually instantiate the reader.
    DatastoreReader reader = new DatastoreReader(source, mockDatastore);
    // Simply count the number of results returned by the reader.
    assertTrue(reader.start());
    int resultCount = 1;
    while (reader.advance()) {
        resultCount++;
    }
    reader.close();
    // Validate the number of results.
    assertEquals(numEntities, resultCount);
}