fixups

Author: dhalperi

sdk/src/main/java/com/google/cloud/dataflow/sdk/util/BigQueryTableRowIterator.java

@@ -445,6 +445,7 @@ private TableReference executeQueryAndWaitForCompletion()
     destinationTable.setDatasetId(temporaryDatasetId);
     destinationTable.setTableId(temporaryTableId);
     queryConfig.setDestinationTable(destinationTable);
+    queryConfig.setAllowLargeResults(Boolean.TRUE);
 
     Job queryJob = executeWithBackOff(
         client.jobs().insert(projectId, job),

sdk/src/test/java/com/google/cloud/dataflow/sdk/util/BigQueryTableRowIteratorTest.java

@@ -258,14 +258,18 @@ public void testReadFromQueryNoTables() throws IOException, InterruptedException
 
     // Mock job polling.
     JobStatus status = new JobStatus().setState("DONE");
-    TableReference tableRef =
-        new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
-    JobConfigurationQuery queryConfig = new JobConfigurationQuery().setDestinationTable(tableRef);
+    JobConfigurationQuery resultQueryConfig =
+        new JobConfigurationQuery().setDestinationTable(
+            new TableReference()
+                .setProjectId("project")
+                .setDatasetId("tempdataset")
+                .setTableId("temptable")
+        );
     Job getJob =
         new Job()
             .setJobReference(new JobReference())
             .setStatus(status)
-            .setConfiguration(new JobConfiguration().setQuery(queryConfig));
+            .setConfiguration(new JobConfiguration().setQuery(resultQueryConfig));
     when(mockJobsGet.execute()).thenReturn(getJob);
 
     // Mock table schema fetch.
@@ -281,8 +285,9 @@ public void testReadFromQueryNoTables() throws IOException, InterruptedException
     String query = String.format(
         "SELECT \"Arthur\" as name, 42 as count, \"%s\" as photo",
         photoBytesEncoded);
+    JobConfigurationQuery queryConfig = new JobConfigurationQuery().setQuery(query);
     try (BigQueryTableRowIterator iterator =
-        BigQueryTableRowIterator.fromQuery(query, "project", mockClient, null, null)) {
+        BigQueryTableRowIterator.fromQuery(queryConfig, "project", mockClient)) {
       iterator.open();
       assertTrue(iterator.advance());
       TableRow row = iterator.getCurrent();
@@ -317,7 +322,7 @@ public void testReadFromQueryNoTables() throws IOException, InterruptedException
     verify(mockTablesDelete).execute();
     // Table data read.
     verify(mockClient).tabledata();
-    verify(mockTabledata).list("project", "dataset", "table");
+    verify(mockTabledata).list("project", "tempdataset", "temptable");
     verify(mockTabledataList).execute();
   }
 
@@ -334,18 +339,16 @@ public void testQueryFailed() throws IOException {
     when(mockJobsInsert.execute()).thenThrow(exception, exception, exception, exception);
 
     String query = "NOT A QUERY";
+    JobConfigurationQuery queryConfig = new JobConfigurationQuery().setQuery(query);
     try (BigQueryTableRowIterator iterator =
-            BigQueryTableRowIterator.fromQuery(query, "project", mockClient, null, null)) {
-
-      try {
-        iterator.open();
-        fail();
-      } catch (Exception expected) {
-        // Verify message explains cause and reports the query.
-        assertThat(expected.getMessage(), containsString("Error"));
-        assertThat(expected.getMessage(), containsString(query));
-        assertThat(expected.getCause().getMessage(), containsString(errorReason));
-      }
+            BigQueryTableRowIterator.fromQuery(queryConfig, "project", mockClient)) {
+      iterator.open();
+      fail();
+    } catch (Exception expected) {
+      // Verify message explains cause and reports the query.
+      assertThat(expected.getMessage(), containsString("Error"));
+      assertThat(expected.getMessage(), containsString(query));
+      assertThat(expected.getCause().getMessage(), containsString(errorReason));
     }
 
     // Job inserted to run the query, then polled once.