|
32 | 32 | import com.google.cloud.bigquery.BigQueryOptions; |
33 | 33 | import com.google.cloud.bigquery.Field; |
34 | 34 | import com.google.cloud.bigquery.FieldList; |
| 35 | +import com.google.cloud.bigquery.FieldValueList; |
35 | 36 | import com.google.cloud.bigquery.Job; |
36 | 37 | import com.google.cloud.bigquery.JobId; |
37 | 38 | import com.google.cloud.bigquery.JobInfo; |
|
55 | 56 | import java.io.IOException; |
56 | 57 | import java.sql.ResultSet; |
57 | 58 | import java.sql.SQLException; |
| 59 | +import java.util.ArrayList; |
58 | 60 | import java.util.HashMap; |
59 | 61 | import java.util.List; |
60 | 62 | import java.util.Map; |
@@ -494,4 +496,98 @@ public void testGetStatementType(boolean isReadOnlyTokenUsed) throws Exception { |
494 | 496 | verify(bigquery, isReadOnlyTokenUsed ? Mockito.never() : Mockito.times(1)) |
495 | 497 | .create(any(JobInfo.class)); |
496 | 498 | } |
| 499 | + |
| 500 | + @Test |
| 501 | + public void testUseReadAPI_SafeguardSmallDataset() throws SQLException { |
| 502 | + // Setup: totalRows < MinTableSize, so it should not activate the Read API |
| 503 | + doReturn(true).when(bigQueryConnection).isEnableHighThroughputAPI(); |
| 504 | + doReturn(100).when(bigQueryConnection).getHighThroughputMinTableSize(); |
| 505 | + doReturn(2).when(bigQueryConnection).getHighThroughputActivationRatio(); |
| 506 | + doReturn(1000L).when(bigQueryConnection).getMaxResults(); |
| 507 | + |
| 508 | + BigQueryStatement statement = new BigQueryStatement(bigQueryConnection); |
| 509 | + TableResult tableResult = mock(TableResult.class); |
| 510 | + doReturn(50L).when(tableResult).getTotalRows(); |
| 511 | + |
| 512 | + // Standard java collection in values |
| 513 | + List<FieldValueList> valuesList = new ArrayList<>(); |
| 514 | + for (int i = 0; i < 50; i++) { |
| 515 | + valuesList.add(mock(FieldValueList.class)); |
| 516 | + } |
| 517 | + doReturn(valuesList).when(tableResult).getValues(); |
| 518 | + |
| 519 | + boolean useReadApi = statement.useReadAPI(tableResult); |
| 520 | + assertThat(useReadApi).isFalse(); |
| 521 | + } |
| 522 | + |
| 523 | + @Test |
| 524 | + public void testUseReadAPI_SafeguardNoNextPage() throws SQLException { |
| 525 | + // Setup: totalRows = 500 > MinTableSize (100), but hasNextPage() is false. |
| 526 | + // Safeguard should prevent double-fetching and not activate the Read API. |
| 527 | + doReturn(true).when(bigQueryConnection).isEnableHighThroughputAPI(); |
| 528 | + doReturn(100).when(bigQueryConnection).getHighThroughputMinTableSize(); |
| 529 | + doReturn(2).when(bigQueryConnection).getHighThroughputActivationRatio(); |
| 530 | + doReturn(1000L).when(bigQueryConnection).getMaxResults(); |
| 531 | + |
| 532 | + BigQueryStatement statement = new BigQueryStatement(bigQueryConnection); |
| 533 | + TableResult tableResult = mock(TableResult.class); |
| 534 | + doReturn(500L).when(tableResult).getTotalRows(); |
| 535 | + doReturn(false).when(tableResult).hasNextPage(); |
| 536 | + |
| 537 | + boolean useReadApi = statement.useReadAPI(tableResult); |
| 538 | + assertThat(useReadApi).isFalse(); |
| 539 | + } |
| 540 | + |
| 541 | + @Test |
| 542 | + public void testUseReadAPI_MeetsRatio() throws SQLException { |
| 543 | + // Setup: totalRows = 500, maxResultPerPage = 100, MinTableSize = 100, ActivationRatio = 2 |
| 544 | + // ratio = 5 > 2, should activate Read API |
| 545 | + doReturn(true).when(bigQueryConnection).isEnableHighThroughputAPI(); |
| 546 | + doReturn(100).when(bigQueryConnection).getHighThroughputMinTableSize(); |
| 547 | + doReturn(2).when(bigQueryConnection).getHighThroughputActivationRatio(); |
| 548 | + doReturn(100L).when(bigQueryConnection).getMaxResults(); |
| 549 | + |
| 550 | + BigQueryStatement statement = new BigQueryStatement(bigQueryConnection); |
| 551 | + TableResult tableResult = mock(TableResult.class); |
| 552 | + doReturn(500L).when(tableResult).getTotalRows(); |
| 553 | + doReturn(true).when(tableResult).hasNextPage(); |
| 554 | + |
| 555 | + boolean useReadApi = statement.useReadAPI(tableResult); |
| 556 | + assertThat(useReadApi).isTrue(); |
| 557 | + } |
| 558 | + |
| 559 | + @Test |
| 560 | + public void testUseReadAPI_FailsMinTableSize() throws SQLException { |
| 561 | + // Setup: totalRows = 80 < MinTableSize (100) |
| 562 | + doReturn(true).when(bigQueryConnection).isEnableHighThroughputAPI(); |
| 563 | + doReturn(100).when(bigQueryConnection).getHighThroughputMinTableSize(); |
| 564 | + doReturn(2).when(bigQueryConnection).getHighThroughputActivationRatio(); |
| 565 | + doReturn(1000L).when(bigQueryConnection).getMaxResults(); |
| 566 | + |
| 567 | + BigQueryStatement statement = new BigQueryStatement(bigQueryConnection); |
| 568 | + TableResult tableResult = mock(TableResult.class); |
| 569 | + doReturn(80L).when(tableResult).getTotalRows(); |
| 570 | + |
| 571 | + boolean useReadApi = statement.useReadAPI(tableResult); |
| 572 | + assertThat(useReadApi).isFalse(); |
| 573 | + } |
| 574 | + |
| 575 | + @Test |
| 576 | + public void testUseReadAPI_ZeroPageSizeDivisionByZeroSafeguard() throws SQLException { |
| 577 | + // Setup: totalRows = 500, MinTableSize = 100, ActivationRatio = 2, maxResultPerPage = 0 |
| 578 | + // Verify that the division by zero check safely guards and falls back to pageSize = 1 |
| 579 | + doReturn(true).when(bigQueryConnection).isEnableHighThroughputAPI(); |
| 580 | + doReturn(100).when(bigQueryConnection).getHighThroughputMinTableSize(); |
| 581 | + doReturn(2).when(bigQueryConnection).getHighThroughputActivationRatio(); |
| 582 | + doReturn(0L).when(bigQueryConnection).getMaxResults(); // maxResultPerPage = 0 |
| 583 | + |
| 584 | + BigQueryStatement statement = new BigQueryStatement(bigQueryConnection); |
| 585 | + TableResult tableResult = mock(TableResult.class); |
| 586 | + doReturn(500L).when(tableResult).getTotalRows(); |
| 587 | + doReturn(true).when(tableResult).hasNextPage(); |
| 588 | + |
| 589 | + // This should not throw ArithmeticException (/ by zero) and should evaluate safely |
| 590 | + boolean useReadApi = statement.useReadAPI(tableResult); |
| 591 | + assertThat(useReadApi).isTrue(); // ratio = 500 / 1 = 500 > 2 -> true |
| 592 | + } |
497 | 593 | } |
0 commit comments