fix: fix start index with page size for list rows (#793)

googleapis · Oct 13, 2020 · 71e7d48 · 71e7d48
1 parent d786194
commit 71e7d48
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 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
@@ -1089,9 +1089,11 @@ public TableDataList call() {
               EXCEPTION_HANDLER,
               serviceOptions.getClock());
       String cursor = result.getPageToken();
+      Map<BigQueryRpc.Option, ?> 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
@@ -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<BigQueryRpc.Option, ?> 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
@@ -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<String> 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<FieldValueList> 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 {