fix(bigquery-jdbc): maxResults was not honored for child jobs

java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java

@@ -771,9 +771,17 @@ Job getNextJob() {
     return null;
   }
 
-  void handleQueryResult(String query, TableResult results, SqlType queryType)
-      throws SQLException, InterruptedException {
-    handleQueryResult(query, results, queryType, null);
+  SqlType handleQueryResult(String query, Job job) throws SQLException, InterruptedException {
+    JobStatistics statistics = job.getStatistics();
+    StatementType statementType =
+        (statistics instanceof QueryStatistics)
+            ? ((QueryStatistics) statistics).getStatementType()
+            : null;
+    SqlType sqlType = getQueryType(job.getConfiguration(), statementType);
+    TableResult results =
+        job.getQueryResults(QueryResultsOption.pageSize(querySettings.getMaxResultPerPage()));
+    handleQueryResult(query, results, sqlType, job);
+    return sqlType;
   }
 
   void handleQueryResult(String query, TableResult results, SqlType queryType, Job job)
@@ -819,10 +827,7 @@ void handleQueryResult(String query, TableResult results, SqlType queryType, Job
           if (currentJob == null) {
             return;
           }
-          StatementType statementType =
-              ((QueryStatistics) (currentJob.getStatistics())).getStatementType();
-          SqlType sqlType = getQueryType(currentJob.getConfiguration(), statementType);
-          handleQueryResult(query, currentJob.getQueryResults(), sqlType, currentJob);
+          handleQueryResult(query, currentJob);
         } catch (NullPointerException ex) {
           throw new BigQueryJdbcException(ex);
         }
@@ -1847,12 +1852,7 @@ private boolean getMoreResultsImpl(int current) throws SQLException {
 
       Job currentJob = getNextJob();
       if (currentJob != null) {
-        StatementType statementType =
-            ((QueryStatistics) (currentJob.getStatistics())).getStatementType();
-        SqlType sqlType = getQueryType(currentJob.getConfiguration(), statementType);
-        handleQueryResult(this.scriptQuery, currentJob.getQueryResults(), sqlType, currentJob);
-
-        return sqlType == SqlType.SELECT;
+        return handleQueryResult(this.scriptQuery, currentJob) == SqlType.SELECT;
       } else {
         resetStatementFields();
         return false;

java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java

@@ -32,13 +32,15 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.withSettings;
 
+import com.google.api.gax.paging.Page;
 import com.google.api.gax.rpc.ApiException;
 import com.google.api.gax.rpc.ServerStream;
 import com.google.api.gax.rpc.ServerStreamingCallable;
 import com.google.api.gax.rpc.StatusCode;
 import com.google.cloud.ServiceOptions;
 import com.google.cloud.Tuple;
 import com.google.cloud.bigquery.BigQuery;
+import com.google.cloud.bigquery.BigQuery.JobListOption;
 import com.google.cloud.bigquery.BigQuery.QueryResultsOption;
 import com.google.cloud.bigquery.BigQuery.TableDataListOption;
 import com.google.cloud.bigquery.BigQueryException;
@@ -55,6 +57,7 @@
 import com.google.cloud.bigquery.JobStatistics;
 import com.google.cloud.bigquery.JobStatistics.QueryStatistics;
 import com.google.cloud.bigquery.JobStatistics.QueryStatistics.StatementType;
+import com.google.cloud.bigquery.JobStatistics.ScriptStatistics;
 import com.google.cloud.bigquery.JobStatistics.SessionInfo;
 import com.google.cloud.bigquery.QueryJobConfiguration;
 import com.google.cloud.bigquery.QueryJobConfiguration.Priority;
@@ -1715,4 +1718,160 @@ public ReadRowsResponse next() {
       assertTrue(thrown.getCause().getCause() instanceof OutOfMemoryError);
     }
   }
+
+  @Test
+  public void testScriptAppliesMaxResultsPageSizeAndUsesReadApiAcrossChildJobs() throws Exception {
+    BigQueryStatement stmtSpy = createReadApiStatementSpy();
+
+    TableResult parentResult = mockTableResultWithJob("parent-script-job");
+    doReturn(StatementType.SCRIPT).when(parentResult).getStatementType();
+    doReturn(parentResult)
+        .when(bigquery)
+        .queryWithTimeout(any(QueryJobConfiguration.class), any(), any());
+
+    Job childJob1 =
+        createScriptChildJobMock("child-select-1", StatementType.SELECT, 5000L, 1000L, true, null);
+    Job childJob2 =
+        createScriptChildJobMock("child-dml-2", StatementType.INSERT, 0L, 0L, false, 4L);
+    Job childJob3 =
+        createScriptChildJobMock("child-select-3", StatementType.SELECT, 7000L, 1000L, true, null);
+
+    @SuppressWarnings("unchecked")
+    Page<Job> childJobsPage = mock(Page.class);
+    // BigQuery jobs.list returns child jobs in reverse chronological order; handleQueryResult
+    // reverses them into execution order.
+    doReturn(Arrays.asList(childJob3, childJob2, childJob1)).when(childJobsPage).iterateAll();
+    doReturn(childJobsPage)
+        .when(bigquery)
+        .listJobs(eq(JobListOption.parentJobId("parent-script-job")));
+
+    String scriptSql =
+        "SELECT * FROM UNNEST(GENERATE_ARRAY(1, 5000));"
+            + "INSERT INTO t VALUES (1),(2),(3),(4);"
+            + "SELECT * FROM UNNEST(GENERATE_ARRAY(1, 7000));";
+
+    assertTrue(stmtSpy.execute(scriptSql));
+    verify(childJob1).getQueryResults(QueryResultsOption.pageSize(1000L));
+    verify(childJob1, Mockito.never()).getQueryResults();
+    ResultSet rs1 = stmtSpy.getResultSet();
+    assertThat(rs1).isInstanceOf(BigQueryArrowResultSet.class);
+    assertEquals(-1, stmtSpy.getUpdateCount());
+
+    assertFalse(stmtSpy.getMoreResults());
+    verify(childJob2).getQueryResults(QueryResultsOption.pageSize(1000L));
+    verify(childJob2, Mockito.never()).getQueryResults();
+    assertNull(stmtSpy.getResultSet());
+    assertEquals(4, stmtSpy.getUpdateCount());
+
+    assertTrue(stmtSpy.getMoreResults());
+    verify(childJob3).getQueryResults(QueryResultsOption.pageSize(1000L));
+    verify(childJob3, Mockito.never()).getQueryResults();
+    ResultSet rs3 = stmtSpy.getResultSet();
+    assertThat(rs3).isInstanceOf(BigQueryArrowResultSet.class);
+    assertEquals(-1, stmtSpy.getUpdateCount());
+
+    assertFalse(stmtSpy.getMoreResults());
+    assertNull(stmtSpy.getResultSet());
+    assertEquals(-1, stmtSpy.getUpdateCount());
+    verify(stmtSpy, Mockito.times(2)).getReadSession(any(CreateReadSessionRequest.class));
+  }
+
+  @Test
+  public void testScriptUsesJsonForSmallSelectAndReadApiForLargeSelect() throws Exception {
+    BigQueryStatement stmtSpy = createReadApiStatementSpy();
+
+    TableResult parentResult = mockTableResultWithJob("parent-mixed-script-job");
+    doReturn(StatementType.SCRIPT).when(parentResult).getStatementType();
+    doReturn(parentResult)
+        .when(bigquery)
+        .queryWithTimeout(any(QueryJobConfiguration.class), any(), any());
+
+    Job smallChildJob =
+        createScriptChildJobMock("child-small-select", StatementType.SELECT, 10L, 10L, false, null);
+    Job largeChildJob =
+        createScriptChildJobMock(
+            "child-large-select", StatementType.SELECT, 5000L, 1000L, true, null);
+
+    @SuppressWarnings("unchecked")
+    Page<Job> childJobsPage = mock(Page.class);
+    doReturn(Arrays.asList(largeChildJob, smallChildJob)).when(childJobsPage).iterateAll();
+    doReturn(childJobsPage)
+        .when(bigquery)
+        .listJobs(eq(JobListOption.parentJobId("parent-mixed-script-job")));
+
+    assertTrue(stmtSpy.execute("SELECT 1; SELECT * FROM UNNEST(GENERATE_ARRAY(1, 5000));"));
+    verify(smallChildJob).getQueryResults(QueryResultsOption.pageSize(1000L));
+    assertThat(stmtSpy.getResultSet()).isInstanceOf(BigQueryJsonResultSet.class);
+
+    assertTrue(stmtSpy.getMoreResults());
+    verify(largeChildJob).getQueryResults(QueryResultsOption.pageSize(1000L));
+    assertThat(stmtSpy.getResultSet()).isInstanceOf(BigQueryArrowResultSet.class);
+
+    assertFalse(stmtSpy.getMoreResults());
+    verify(stmtSpy, Mockito.times(1)).getReadSession(any(CreateReadSessionRequest.class));
+  }
+
+  @Test
+  public void testHandleQueryResultFromJobCalculatesSqlTypeAndAppliesMaxResults() throws Exception {
+    doReturn(500L).when(bigQueryConnection).getMaxResults();
+    BigQueryStatement stmtSpy = createReadApiStatementSpy();
+
+    Job selectJob =
+        createScriptChildJobMock("job-select", StatementType.SELECT, 2000L, 500L, true, null);
+    assertEquals(
+        BigQueryStatement.SqlType.SELECT, stmtSpy.handleQueryResult("SELECT * FROM t", selectJob));
+    verify(selectJob).getQueryResults(QueryResultsOption.pageSize(500L));
+    assertThat(stmtSpy.getResultSet()).isInstanceOf(BigQueryArrowResultSet.class);
+
+    Job dmlJob = createScriptChildJobMock("job-insert", StatementType.INSERT, 0L, 0L, false, 7L);
+    assertEquals(
+        BigQueryStatement.SqlType.DML,
+        stmtSpy.handleQueryResult("INSERT INTO t VALUES (1)", dmlJob));
+    verify(dmlJob).getQueryResults(QueryResultsOption.pageSize(500L));
+    assertNull(stmtSpy.getResultSet());
+    assertEquals(7, stmtSpy.getUpdateCount());
+  }
+
+  private BigQueryStatement createReadApiStatementSpy() throws SQLException {
+    doReturn(true).when(bigQueryConnection).isEnableHighThroughputAPI();
+    doReturn(100).when(bigQueryConnection).getHighThroughputMinTableSize();
+    doReturn(2).when(bigQueryConnection).getHighThroughputActivationRatio();
+
+    BigQueryStatement stmtSpy = Mockito.spy(new BigQueryStatement(bigQueryConnection));
+    doReturn(storageReadClient).when(stmtSpy).getBigQueryReadClient();
+    doReturn(ReadSession.getDefaultInstance())
+        .when(stmtSpy)
+        .getReadSession(any(CreateReadSessionRequest.class));
+    doReturn(arrowSchema).when(stmtSpy).getArrowSchema(any(ReadSession.class));
+    doReturn(mock(Future.class))
+        .when(stmtSpy)
+        .populateArrowBufferedQueue(
+            any(ReadSession.class), any(BlockingQueue.class), any(BigQueryReadClient.class));
+    doReturn(TABLE_ID).when(stmtSpy).getDestinationTable(any());
+    return stmtSpy;
+  }
+
+  private Job createScriptChildJobMock(
+      String childId,
+      StatementType statementType,
+      long totalRows,
+      long rowsInPage,
+      boolean hasNextPage,
+      Long dmlAffectedRows)
+      throws InterruptedException {
+    TableResult childResult = mockTableResultWithJob(childId);
+    doReturn(Schema.of(fieldList)).when(childResult).getSchema();
+    doReturn(totalRows).when(childResult).getTotalRows();
+    doReturn(rowsInPage).when(childResult).getRowsInPage();
+    doReturn(hasNextPage).when(childResult).hasNextPage();
+    if (dmlAffectedRows != null) {
+      doReturn(dmlAffectedRows).when(childResult).getNumDmlAffectedRows();
+    }
+
+    Job childJob =
+        getJobMock(childResult, QueryJobConfiguration.newBuilder(query).build(), statementType);
+    QueryStatistics stats = childJob.getStatistics();
+    doReturn(mock(ScriptStatistics.class)).when(stats).getScriptStatistics();
+    return childJob;
+  }
 }

java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITCallableStatementTest.java

@@ -647,4 +647,43 @@ public void testCallableStatementScriptExecute() throws SQLException {
     assertEquals(-1, callableStatement.getUpdateCount());
     callableStatement.close();
   }
+
+  @Test
+  public void testCallableStatementScriptWithReadApi() throws SQLException {
+    String readApiUrl =
+        connection_uri
+            + "EnableHighThroughputAPI=1;"
+            + "MaxResults=1000;"
+            + "HighThroughputActivationRatio=2;"
+            + "HighThroughputMinTableSize=100;";
+    int randomNum = random.nextInt(99);
+    String callableStmtQuery =
+        String.format(
+                "DECLARE call_result STRING;CALL %s.%s(?,?,call_result);",
+                DATASET, CALLABLE_STMT_PROC_NAME)
+            + String.format(GENERATE_ARRAY_QUERY, 5000)
+            + String.format(GENERATE_ARRAY_QUERY, 7000);
+
+    try (Connection readApiConnection = DriverManager.getConnection(readApiUrl);
+        CallableStatement callableStatement = readApiConnection.prepareCall(callableStmtQuery)) {
+      callableStatement.setString(1, "callable-readapi-test");
+      callableStatement.setInt(2, randomNum);
+
+      assertFalse(callableStatement.execute());
+      assertEquals(1, callableStatement.getUpdateCount());
+
+      assertTrue(callableStatement.getMoreResults());
+      ResultSet resultSet1 = callableStatement.getResultSet();
+      assertArrowResultSet(resultSet1);
+      assertEquals(5000, resultSetRowCount(resultSet1));
+
+      assertTrue(callableStatement.getMoreResults());
+      ResultSet resultSet2 = callableStatement.getResultSet();
+      assertArrowResultSet(resultSet2);
+      assertEquals(7000, resultSetRowCount(resultSet2));
+
+      assertFalse(callableStatement.getMoreResults());
+      assertEquals(-1, callableStatement.getUpdateCount());
+    }
+  }
 }

java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITStatementTest.java

@@ -227,21 +227,32 @@ public void testExecuteUpdate() throws SQLException {
 
   @Test
   public void testScript() throws SQLException {
-    String connection_uri = ITBase.connectionUrl;
     Properties withReadApi = new Properties();
     withReadApi.setProperty("EnableHighThroughputAPI", "1");
-    Connection connection = DriverManager.getConnection(connection_uri, withReadApi);
-    Statement statement = connection.createStatement();
-    String BASE_QUERY = "SELECT * FROM UNNEST(GENERATE_ARRAY(1, %s));";
-    int expectedCnt = 500000;
-    String longQuery = String.format(BASE_QUERY, expectedCnt);
-    String longerQuery = String.format(BASE_QUERY, 700000);
-    String longerQuery2 = String.format(BASE_QUERY, 900000);
-    statement.execute(longQuery + longerQuery + longerQuery2);
-    ResultSet arrowResultSet = statement.getResultSet();
-    assertEquals(500000, resultSetRowCount(arrowResultSet));
-    arrowResultSet.close();
-    connection.close();
+    String longQuery = String.format(GENERATE_ARRAY_QUERY, 500000);
+    String longerQuery = String.format(GENERATE_ARRAY_QUERY, 700000);
+    String longerQuery2 = String.format(GENERATE_ARRAY_QUERY, 900000);
+
+    try (Connection connection = DriverManager.getConnection(ITBase.connectionUrl, withReadApi);
+        Statement statement = connection.createStatement()) {
+      assertTrue(statement.execute(longQuery + longerQuery + longerQuery2));
+
+      ResultSet arrowResultSet = statement.getResultSet();
+      assertArrowResultSet(arrowResultSet);
+      assertEquals(500000, resultSetRowCount(arrowResultSet));
+
+      assertTrue(statement.getMoreResults());
+      arrowResultSet = statement.getResultSet();
+      assertArrowResultSet(arrowResultSet);
+      assertEquals(700000, resultSetRowCount(arrowResultSet));
+
+      assertTrue(statement.getMoreResults());
+      arrowResultSet = statement.getResultSet();
+      assertArrowResultSet(arrowResultSet);
+      assertEquals(900000, resultSetRowCount(arrowResultSet));
+
+      assertFalse(statement.getMoreResults());
+    }
   }
 
   @Test