From 71e7d4804e1facfb921892dc702adaec35e47cd0 Mon Sep 17 00:00:00 2001 From: Suraj Dhamecha <48670070+suraj-qlogic@users.noreply.github.com> Date: Wed, 14 Oct 2020 01:33:12 +0530 Subject: [PATCH] fix: fix start index with page size for list rows (#793) --- .../google/cloud/bigquery/BigQueryImpl.java | 4 ++- .../cloud/bigquery/BigQueryImplTest.java | 35 +++++++++++++++++++ .../cloud/bigquery/it/ITBigQueryTest.java | 23 ++++++++++++ 3 files changed, 61 insertions(+), 1 deletion(-) diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java index d53f45909..76067be0d 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java @@ -1089,9 +1089,11 @@ public TableDataList call() { EXCEPTION_HANDLER, serviceOptions.getClock()); String cursor = result.getPageToken(); + Map pageOptionMap = + Strings.isNullOrEmpty(cursor) ? optionsMap : optionMap(TableDataListOption.startIndex(0)); return Tuple.of( new PageImpl<>( - new TableDataPageFetcher(tableId, schema, serviceOptions, cursor, optionsMap), + new TableDataPageFetcher(tableId, schema, serviceOptions, cursor, pageOptionMap), cursor, transformTableData(result.getRows(), schema)), result.getTotalRows()); diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java index 9ed0edbae..a9b2293fa 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java @@ -26,6 +26,7 @@ import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; import static org.mockito.Mockito.any; +import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; @@ -1470,6 +1471,40 @@ public void testListTableDataWithOptions() { verify(bigqueryRpcMock).listTableData(PROJECT, DATASET, TABLE, TABLE_DATA_LIST_OPTIONS); } + @Test + public void testListTableDataWithNextPage() { + doReturn(TABLE_DATA_PB) + .when(bigqueryRpcMock) + .listTableData(PROJECT, DATASET, TABLE, TABLE_DATA_LIST_OPTIONS); + bigquery = options.getService(); + TableResult page = + bigquery.listTableData( + DATASET, + TABLE, + TABLE_DATA_LIST_PAGE_SIZE, + TABLE_DATA_LIST_PAGE_TOKEN, + TABLE_DATA_LIST_START_INDEX); + assertEquals(CURSOR, page.getNextPageToken()); + verify(bigqueryRpcMock).listTableData(PROJECT, DATASET, TABLE, TABLE_DATA_LIST_OPTIONS); + assertArrayEquals(TABLE_DATA.toArray(), Iterables.toArray(page.getValues(), List.class)); + Map SECOND_TABLE_DATA_LIST_OPTIONS = + ImmutableMap.of(BigQueryRpc.Option.PAGE_TOKEN, CURSOR, BigQueryRpc.Option.START_INDEX, 0L); + doReturn( + new TableDataList() + .setPageToken(null) + .setTotalRows(1L) + .setRows( + ImmutableList.of( + new TableRow().setF(ImmutableList.of(new TableCell().setV("Value3"))), + new TableRow().setF(ImmutableList.of(new TableCell().setV("Value4")))))) + .when(bigqueryRpcMock) + .listTableData(PROJECT, DATASET, TABLE, SECOND_TABLE_DATA_LIST_OPTIONS); + assertTrue(page.hasNextPage()); + page = page.getNextPage(); + assertNull(page.getNextPageToken()); + verify(bigqueryRpcMock).listTableData(PROJECT, DATASET, TABLE, SECOND_TABLE_DATA_LIST_OPTIONS); + } + // The "minimally initialized" Job that lets Job.fromPb run without throwing. private static com.google.api.services.bigquery.model.Job newJobPb() { return new com.google.api.services.bigquery.model.Job() diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java index 44125ce77..05332ba6a 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java @@ -112,6 +112,7 @@ import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; import java.nio.file.FileSystems; +import java.util.Collection; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; @@ -365,6 +366,9 @@ public class ITBigQueryTest { private static final Set PUBLIC_DATASETS = ImmutableSet.of("github_repos", "hacker_news", "noaa_gsod", "samples", "usa_names"); + private static final String PUBLIC_PROJECT = "bigquery-public-data"; + private static final String PUBLIC_DATASET = "census_bureau_international"; + private static BigQuery bigquery; private static Storage storage; @@ -1342,6 +1346,25 @@ public void testListAllTableData() { assertEquals(2, rowCount); } + @Test + public void testListPageWithStartIndex() { + String tableName = "midyear_population_agespecific"; + TableId tableId = TableId.of(PUBLIC_PROJECT, PUBLIC_DATASET, tableName); + Table table = bigquery.getTable(tableId); + long numRows = table.getNumRows().longValue(); + Page tableResult = + bigquery.listTableData( + tableId, + BigQuery.TableDataListOption.startIndex(numRows - 300_000L), + BigQuery.TableDataListOption.pageSize(600_000L)); + assertNotNull(tableResult.getNextPageToken()); + long totalRows = ((Collection) tableResult.getValues()).size(); + tableResult = tableResult.getNextPage(); + totalRows = totalRows + ((Collection) tableResult.getValues()).size(); + assertNull(tableResult.getNextPageToken()); + assertEquals(300_000L, totalRows); + } + @Test public void testModelLifecycle() throws InterruptedException {