diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java index e036030e6d6f..05fb99f1577f 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java +++ b/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; diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java index 4ad8358f3894..a750163c5e42 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java @@ -32,6 +32,7 @@ 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; @@ -39,6 +40,7 @@ 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 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 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; + } } diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITCallableStatementTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITCallableStatementTest.java index 725167cc5660..fbee7a45f349 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITCallableStatementTest.java +++ b/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()); + } + } } diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITStatementTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITStatementTest.java index 3bdd48774e03..3fb7a816cf5e 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITStatementTest.java +++ b/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