From 35c0c90c002c50ab37ce55a74ddbca9301c59bcd Mon Sep 17 00:00:00 2001 From: kirl Date: Wed, 11 Nov 2026 18:47:31 +0000 Subject: [PATCH 1/2] fix(bigquery-jdbc): pass MaxResults page size to SCRIPT child jobs for Read API --- .../bigquery/jdbc/BigQueryStatement.java | 26 +-- .../bigquery/jdbc/BigQueryStatementTest.java | 216 ++++++++++++++++++ .../jdbc/it/ITCallableStatementTest.java | 40 ++++ .../bigquery/jdbc/it/ITStatementTest.java | 17 +- 4 files changed, 285 insertions(+), 14 deletions(-) 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..b3845cfe8194 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,217 @@ public ReadRowsResponse next() { assertTrue(thrown.getCause().getCause() instanceof OutOfMemoryError); } } + + @Test + public void testScriptAppliesMaxResultsPageSizeAndUsesReadApiAcrossChildJobs() throws Exception { + doReturn(true).when(bigQueryConnection).isEnableHighThroughputAPI(); + doReturn(100).when(bigQueryConnection).getHighThroughputMinTableSize(); + doReturn(2).when(bigQueryConnection).getHighThroughputActivationRatio(); + doReturn(1000L).when(bigQueryConnection).getMaxResults(); + + BigQueryStatement stmt = new BigQueryStatement(bigQueryConnection); + BigQueryStatement stmtSpy = Mockito.spy(stmt); + BigQueryReadClient readClientMock = mock(BigQueryReadClient.class); + ReadSession readSession = ReadSession.getDefaultInstance(); + Future mockWorker = mock(Future.class); + doReturn(readClientMock).when(stmtSpy).getBigQueryReadClient(); + doReturn(readSession).when(stmtSpy).getReadSession(any(CreateReadSessionRequest.class)); + doReturn(arrowSchema).when(stmtSpy).getArrowSchema(any(ReadSession.class)); + doReturn(mockWorker) + .when(stmtSpy) + .populateArrowBufferedQueue( + any(ReadSession.class), any(BlockingQueue.class), any(BigQueryReadClient.class)); + doReturn(TableId.of(PROJECT, DEFAULT_TEST_DATASET, "dest_table")) + .when(stmtSpy) + .getDestinationTable(any(JobId.class)); + + JobId parentId = JobId.of(PROJECT, "parent-script-job"); + TableResult parentResult = mock(TableResult.class); + doReturn(parentId).when(parentResult).getJobId(); + 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 { + doReturn(true).when(bigQueryConnection).isEnableHighThroughputAPI(); + doReturn(100).when(bigQueryConnection).getHighThroughputMinTableSize(); + doReturn(2).when(bigQueryConnection).getHighThroughputActivationRatio(); + doReturn(1000L).when(bigQueryConnection).getMaxResults(); + + BigQueryStatement stmt = new BigQueryStatement(bigQueryConnection); + BigQueryStatement stmtSpy = Mockito.spy(stmt); + BigQueryReadClient readClientMock = mock(BigQueryReadClient.class); + ReadSession readSession = ReadSession.getDefaultInstance(); + Future mockWorker = mock(Future.class); + doReturn(readClientMock).when(stmtSpy).getBigQueryReadClient(); + doReturn(readSession).when(stmtSpy).getReadSession(any(CreateReadSessionRequest.class)); + doReturn(arrowSchema).when(stmtSpy).getArrowSchema(any(ReadSession.class)); + doReturn(mockWorker) + .when(stmtSpy) + .populateArrowBufferedQueue( + any(ReadSession.class), any(BlockingQueue.class), any(BigQueryReadClient.class)); + doReturn(TableId.of(PROJECT, DEFAULT_TEST_DATASET, "dest_table")) + .when(stmtSpy) + .getDestinationTable(any(JobId.class)); + + JobId parentId = JobId.of(PROJECT, "parent-mixed-script-job"); + TableResult parentResult = mock(TableResult.class); + doReturn(parentId).when(parentResult).getJobId(); + 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(true).when(bigQueryConnection).isEnableHighThroughputAPI(); + doReturn(100).when(bigQueryConnection).getHighThroughputMinTableSize(); + doReturn(2).when(bigQueryConnection).getHighThroughputActivationRatio(); + doReturn(500L).when(bigQueryConnection).getMaxResults(); + + BigQueryStatement stmt = new BigQueryStatement(bigQueryConnection); + BigQueryStatement stmtSpy = Mockito.spy(stmt); + BigQueryReadClient readClientMock = mock(BigQueryReadClient.class); + ReadSession readSession = ReadSession.getDefaultInstance(); + Future mockWorker = mock(Future.class); + doReturn(readClientMock).when(stmtSpy).getBigQueryReadClient(); + doReturn(readSession).when(stmtSpy).getReadSession(any(CreateReadSessionRequest.class)); + doReturn(arrowSchema).when(stmtSpy).getArrowSchema(any(ReadSession.class)); + doReturn(mockWorker) + .when(stmtSpy) + .populateArrowBufferedQueue( + any(ReadSession.class), any(BlockingQueue.class), any(BigQueryReadClient.class)); + doReturn(TableId.of(PROJECT, DEFAULT_TEST_DATASET, "dest_table")) + .when(stmtSpy) + .getDestinationTable(any(JobId.class)); + + 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 Job createScriptChildJobMock( + String childId, + StatementType statementType, + long totalRows, + long rowsInPage, + boolean hasNextPage, + Long dmlAffectedRows) + throws InterruptedException { + Job childJob = mock(Job.class); + JobId childJobId = JobId.of(PROJECT, childId); + QueryStatistics stats = mock(QueryStatistics.class); + ScriptStatistics scriptStats = mock(ScriptStatistics.class); + QueryJobConfiguration config = + QueryJobConfiguration.newBuilder("SELECT 1") + .setDestinationTable(TableId.of(PROJECT, DEFAULT_TEST_DATASET, "dest_" + childId)) + .build(); + TableResult childResult = mock(TableResult.class); + Schema schema = Schema.of(fieldList); + + doReturn(childJobId).when(childJob).getJobId(); + doReturn(config).when(childJob).getConfiguration(); + doReturn(stats).when(childJob).getStatistics(); + doReturn(statementType).when(stats).getStatementType(); + doReturn(scriptStats).when(stats).getScriptStatistics(); + doReturn("STATEMENT").when(scriptStats).getEvaluationKind(); + + doReturn(childJobId).when(childResult).getJobId(); + doReturn(schema).when(childResult).getSchema(); + doReturn(totalRows).when(childResult).getTotalRows(); + doReturn(rowsInPage).when(childResult).getRowsInPage(); + doReturn(hasNextPage).when(childResult).hasNextPage(); + doReturn(Collections.emptyList()).when(childResult).getValues(); + if (dmlAffectedRows != null) { + doReturn(dmlAffectedRows).when(childResult).getNumDmlAffectedRows(); + } + doReturn(childResult).when(childJob).getQueryResults(any(QueryResultsOption.class)); + doReturn(childJob).when(bigquery).getJob(childJobId); + 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..656e083431e1 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,44 @@ 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);" + + "SELECT * FROM UNNEST(GENERATE_ARRAY(1, 5000));" + + "SELECT * FROM UNNEST(GENERATE_ARRAY(1, 7000));", + DATASET, CALLABLE_STMT_PROC_NAME); + + 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(); + assertTrue(resultSet1.getClass().getName().contains("BigQueryArrowResultSet")); + assertEquals(5000, resultSetRowCount(resultSet1)); + + assertTrue(callableStatement.getMoreResults()); + ResultSet resultSet2 = callableStatement.getResultSet(); + assertTrue(resultSet2.getClass().getName().contains("BigQueryArrowResultSet")); + 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..0c05d0e69c55 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 @@ -237,10 +237,25 @@ public void testScript() throws SQLException { 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); + assertTrue(statement.execute(longQuery + longerQuery + longerQuery2)); + ResultSet arrowResultSet = statement.getResultSet(); + assertTrue(arrowResultSet.getClass().getName().contains("BigQueryArrowResultSet")); assertEquals(500000, resultSetRowCount(arrowResultSet)); + + assertTrue(statement.getMoreResults()); + arrowResultSet = statement.getResultSet(); + assertTrue(arrowResultSet.getClass().getName().contains("BigQueryArrowResultSet")); + assertEquals(700000, resultSetRowCount(arrowResultSet)); + + assertTrue(statement.getMoreResults()); + arrowResultSet = statement.getResultSet(); + assertTrue(arrowResultSet.getClass().getName().contains("BigQueryArrowResultSet")); + assertEquals(900000, resultSetRowCount(arrowResultSet)); + + assertFalse(statement.getMoreResults()); arrowResultSet.close(); + statement.close(); connection.close(); } From cd138426499fea63adf929676d1a67f832680263 Mon Sep 17 00:00:00 2001 From: kirl Date: Fri, 13 Nov 2026 06:15:10 +0000 Subject: [PATCH 2/2] test(bigquery-jdbc): refactor script Read API tests to reuse shared helpers --- .../bigquery/jdbc/BigQueryStatementTest.java | 119 +++++------------- .../google/cloud/bigquery/jdbc/it/ITBase.java | 10 ++ .../jdbc/it/ITCallableStatementTest.java | 13 +- .../bigquery/jdbc/it/ITStatementTest.java | 52 ++++---- 4 files changed, 71 insertions(+), 123 deletions(-) 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 b3845cfe8194..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 @@ -1721,30 +1721,9 @@ public ReadRowsResponse next() { @Test public void testScriptAppliesMaxResultsPageSizeAndUsesReadApiAcrossChildJobs() throws Exception { - doReturn(true).when(bigQueryConnection).isEnableHighThroughputAPI(); - doReturn(100).when(bigQueryConnection).getHighThroughputMinTableSize(); - doReturn(2).when(bigQueryConnection).getHighThroughputActivationRatio(); - doReturn(1000L).when(bigQueryConnection).getMaxResults(); + BigQueryStatement stmtSpy = createReadApiStatementSpy(); - BigQueryStatement stmt = new BigQueryStatement(bigQueryConnection); - BigQueryStatement stmtSpy = Mockito.spy(stmt); - BigQueryReadClient readClientMock = mock(BigQueryReadClient.class); - ReadSession readSession = ReadSession.getDefaultInstance(); - Future mockWorker = mock(Future.class); - doReturn(readClientMock).when(stmtSpy).getBigQueryReadClient(); - doReturn(readSession).when(stmtSpy).getReadSession(any(CreateReadSessionRequest.class)); - doReturn(arrowSchema).when(stmtSpy).getArrowSchema(any(ReadSession.class)); - doReturn(mockWorker) - .when(stmtSpy) - .populateArrowBufferedQueue( - any(ReadSession.class), any(BlockingQueue.class), any(BigQueryReadClient.class)); - doReturn(TableId.of(PROJECT, DEFAULT_TEST_DATASET, "dest_table")) - .when(stmtSpy) - .getDestinationTable(any(JobId.class)); - - JobId parentId = JobId.of(PROJECT, "parent-script-job"); - TableResult parentResult = mock(TableResult.class); - doReturn(parentId).when(parentResult).getJobId(); + TableResult parentResult = mockTableResultWithJob("parent-script-job"); doReturn(StatementType.SCRIPT).when(parentResult).getStatementType(); doReturn(parentResult) .when(bigquery) @@ -1799,30 +1778,9 @@ public void testScriptAppliesMaxResultsPageSizeAndUsesReadApiAcrossChildJobs() t @Test public void testScriptUsesJsonForSmallSelectAndReadApiForLargeSelect() throws Exception { - doReturn(true).when(bigQueryConnection).isEnableHighThroughputAPI(); - doReturn(100).when(bigQueryConnection).getHighThroughputMinTableSize(); - doReturn(2).when(bigQueryConnection).getHighThroughputActivationRatio(); - doReturn(1000L).when(bigQueryConnection).getMaxResults(); - - BigQueryStatement stmt = new BigQueryStatement(bigQueryConnection); - BigQueryStatement stmtSpy = Mockito.spy(stmt); - BigQueryReadClient readClientMock = mock(BigQueryReadClient.class); - ReadSession readSession = ReadSession.getDefaultInstance(); - Future mockWorker = mock(Future.class); - doReturn(readClientMock).when(stmtSpy).getBigQueryReadClient(); - doReturn(readSession).when(stmtSpy).getReadSession(any(CreateReadSessionRequest.class)); - doReturn(arrowSchema).when(stmtSpy).getArrowSchema(any(ReadSession.class)); - doReturn(mockWorker) - .when(stmtSpy) - .populateArrowBufferedQueue( - any(ReadSession.class), any(BlockingQueue.class), any(BigQueryReadClient.class)); - doReturn(TableId.of(PROJECT, DEFAULT_TEST_DATASET, "dest_table")) - .when(stmtSpy) - .getDestinationTable(any(JobId.class)); + BigQueryStatement stmtSpy = createReadApiStatementSpy(); - JobId parentId = JobId.of(PROJECT, "parent-mixed-script-job"); - TableResult parentResult = mock(TableResult.class); - doReturn(parentId).when(parentResult).getJobId(); + TableResult parentResult = mockTableResultWithJob("parent-mixed-script-job"); doReturn(StatementType.SCRIPT).when(parentResult).getStatementType(); doReturn(parentResult) .when(bigquery) @@ -1855,26 +1813,8 @@ public void testScriptUsesJsonForSmallSelectAndReadApiForLargeSelect() throws Ex @Test public void testHandleQueryResultFromJobCalculatesSqlTypeAndAppliesMaxResults() throws Exception { - doReturn(true).when(bigQueryConnection).isEnableHighThroughputAPI(); - doReturn(100).when(bigQueryConnection).getHighThroughputMinTableSize(); - doReturn(2).when(bigQueryConnection).getHighThroughputActivationRatio(); doReturn(500L).when(bigQueryConnection).getMaxResults(); - - BigQueryStatement stmt = new BigQueryStatement(bigQueryConnection); - BigQueryStatement stmtSpy = Mockito.spy(stmt); - BigQueryReadClient readClientMock = mock(BigQueryReadClient.class); - ReadSession readSession = ReadSession.getDefaultInstance(); - Future mockWorker = mock(Future.class); - doReturn(readClientMock).when(stmtSpy).getBigQueryReadClient(); - doReturn(readSession).when(stmtSpy).getReadSession(any(CreateReadSessionRequest.class)); - doReturn(arrowSchema).when(stmtSpy).getArrowSchema(any(ReadSession.class)); - doReturn(mockWorker) - .when(stmtSpy) - .populateArrowBufferedQueue( - any(ReadSession.class), any(BlockingQueue.class), any(BigQueryReadClient.class)); - doReturn(TableId.of(PROJECT, DEFAULT_TEST_DATASET, "dest_table")) - .when(stmtSpy) - .getDestinationTable(any(JobId.class)); + BigQueryStatement stmtSpy = createReadApiStatementSpy(); Job selectJob = createScriptChildJobMock("job-select", StatementType.SELECT, 2000L, 500L, true, null); @@ -1892,6 +1832,25 @@ public void testHandleQueryResultFromJobCalculatesSqlTypeAndAppliesMaxResults() 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, @@ -1900,35 +1859,19 @@ private Job createScriptChildJobMock( boolean hasNextPage, Long dmlAffectedRows) throws InterruptedException { - Job childJob = mock(Job.class); - JobId childJobId = JobId.of(PROJECT, childId); - QueryStatistics stats = mock(QueryStatistics.class); - ScriptStatistics scriptStats = mock(ScriptStatistics.class); - QueryJobConfiguration config = - QueryJobConfiguration.newBuilder("SELECT 1") - .setDestinationTable(TableId.of(PROJECT, DEFAULT_TEST_DATASET, "dest_" + childId)) - .build(); - TableResult childResult = mock(TableResult.class); - Schema schema = Schema.of(fieldList); - - doReturn(childJobId).when(childJob).getJobId(); - doReturn(config).when(childJob).getConfiguration(); - doReturn(stats).when(childJob).getStatistics(); - doReturn(statementType).when(stats).getStatementType(); - doReturn(scriptStats).when(stats).getScriptStatistics(); - doReturn("STATEMENT").when(scriptStats).getEvaluationKind(); - - doReturn(childJobId).when(childResult).getJobId(); - doReturn(schema).when(childResult).getSchema(); + 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(); - doReturn(Collections.emptyList()).when(childResult).getValues(); if (dmlAffectedRows != null) { doReturn(dmlAffectedRows).when(childResult).getNumDmlAffectedRows(); } - doReturn(childResult).when(childJob).getQueryResults(any(QueryResultsOption.class)); - doReturn(childJob).when(bigquery).getJob(childJobId); + + 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/ITBase.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBase.java index 98d7a7c014ba..4ad678312fbb 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBase.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBase.java @@ -284,6 +284,11 @@ public static String getBaseConnectionUrl() { getBaseConnectionUrl() + "ProjectId=" + DEFAULT_CATALOG + ";OAuthType=3;Timeout=3600;"; public static final BigQuery bigQuery = BigQueryJdbcBaseTest.getBigQuery(connectionUrl); + /** + * Query template returning rows 1..N; use with {@code String.format(GENERATE_ARRAY_QUERY, n)}. + */ + public static final String GENERATE_ARRAY_QUERY = "SELECT * FROM UNNEST(GENERATE_ARRAY(1, %s));"; + public static final String FORCE_READ_API_PROPERTIES = ";EnableHighThroughputAPI=1;HighThroughputActivationRatio=0;HighThroughputMinTableSize=0;"; @@ -438,6 +443,11 @@ protected static GoogleCredentials getCredentials() throws IOException { new ByteArrayInputStream(authJson.toString().getBytes(StandardCharsets.UTF_8))); } + /** Asserts that the result set was served by the Storage Read (High-Throughput) API. */ + protected static void assertArrowResultSet(ResultSet resultSet) { + assertTrue(resultSet.getClass().getName().contains("BigQueryArrowResultSet")); + } + protected int resultSetRowCount(ResultSet resultSet) throws SQLException { int rowCount = 0; while (resultSet.next()) { 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 656e083431e1..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 @@ -659,11 +659,10 @@ public void testCallableStatementScriptWithReadApi() throws SQLException { int randomNum = random.nextInt(99); String callableStmtQuery = String.format( - "DECLARE call_result STRING;" - + "CALL %s.%s(?,?,call_result);" - + "SELECT * FROM UNNEST(GENERATE_ARRAY(1, 5000));" - + "SELECT * FROM UNNEST(GENERATE_ARRAY(1, 7000));", - DATASET, CALLABLE_STMT_PROC_NAME); + "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)) { @@ -675,12 +674,12 @@ public void testCallableStatementScriptWithReadApi() throws SQLException { assertTrue(callableStatement.getMoreResults()); ResultSet resultSet1 = callableStatement.getResultSet(); - assertTrue(resultSet1.getClass().getName().contains("BigQueryArrowResultSet")); + assertArrowResultSet(resultSet1); assertEquals(5000, resultSetRowCount(resultSet1)); assertTrue(callableStatement.getMoreResults()); ResultSet resultSet2 = callableStatement.getResultSet(); - assertTrue(resultSet2.getClass().getName().contains("BigQueryArrowResultSet")); + assertArrowResultSet(resultSet2); assertEquals(7000, resultSetRowCount(resultSet2)); assertFalse(callableStatement.getMoreResults()); 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 0c05d0e69c55..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,36 +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); - assertTrue(statement.execute(longQuery + longerQuery + longerQuery2)); - - ResultSet arrowResultSet = statement.getResultSet(); - assertTrue(arrowResultSet.getClass().getName().contains("BigQueryArrowResultSet")); - assertEquals(500000, resultSetRowCount(arrowResultSet)); - - assertTrue(statement.getMoreResults()); - arrowResultSet = statement.getResultSet(); - assertTrue(arrowResultSet.getClass().getName().contains("BigQueryArrowResultSet")); - assertEquals(700000, resultSetRowCount(arrowResultSet)); - - assertTrue(statement.getMoreResults()); - arrowResultSet = statement.getResultSet(); - assertTrue(arrowResultSet.getClass().getName().contains("BigQueryArrowResultSet")); - assertEquals(900000, resultSetRowCount(arrowResultSet)); - - assertFalse(statement.getMoreResults()); - arrowResultSet.close(); - statement.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