DatastoreIO: do not split when QuerySplitter fails

As opposed to throwing an exception. Fixes GoogleCloudPlatform#101.

----Release Notes----
[]
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=112999725

Author: dhalperi

sdk/src/main/java/com/google/cloud/dataflow/sdk/io/DatastoreIO.java

@@ -326,8 +326,16 @@ public List<Source> splitIntoBundles(long desiredBundleSizeBytes, PipelineOption
         return ImmutableList.of(this);
       }
 
+      List<Query> datastoreSplits;
+      try {
+        datastoreSplits = getSplitQueries(Ints.checkedCast(numSplits), options);
+      } catch (IllegalArgumentException | DatastoreException e) {
+        LOG.warn("Unable to parallelize the given query: {}", query, e);
+        return ImmutableList.of(this);
+      }
+
       ImmutableList.Builder<Source> splits = ImmutableList.builder();
-      for (Query splitQuery : getSplitQueries(Ints.checkedCast(numSplits), options)) {
+      for (Query splitQuery : datastoreSplits) {
         splits.add(new Source(host, datasetId, splitQuery, namespace));
       }
       return splits.build();

sdk/src/test/java/com/google/cloud/dataflow/sdk/io/DatastoreIOTest.java

@@ -29,6 +29,7 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
@@ -52,6 +53,7 @@
 import com.google.cloud.dataflow.sdk.options.DataflowPipelineOptions;
 import com.google.cloud.dataflow.sdk.options.PipelineOptions;
 import com.google.cloud.dataflow.sdk.options.PipelineOptionsFactory;
+import com.google.cloud.dataflow.sdk.testing.ExpectedLogs;
 import com.google.cloud.dataflow.sdk.util.TestCredential;
 import com.google.common.collect.Lists;
 
@@ -97,6 +99,8 @@ public class DatastoreIOTest {
   @Rule
   public final ExpectedException thrown = ExpectedException.none();
 
+  @Rule public final ExpectedLogs logged = ExpectedLogs.none(DatastoreIO.Source.class);
+
   @Before
   public void setUp() {
     MockitoAnnotations.initMocks(this);
@@ -330,6 +334,37 @@ public void testQueryDoesNotSplitWithLimitSet() throws Exception {
     verifyNoMoreInteractions(splitter);
   }
 
+  /**
+   * 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);
+  }
+
   @Test
   public void testQuerySplitSizeUnavailable() throws Exception {
     KindExpression mykind = KindExpression.newBuilder().setName("mykind").build();