Skip to content
Permalink
Browse files
fix: fix start index with page size for list rows (#793)
  • Loading branch information
suraj-qlogic committed Oct 13, 2020
1 parent d786194 commit 71e7d4804e1facfb921892dc702adaec35e47cd0
@@ -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());
@@ -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()
@@ -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 {

0 comments on commit 71e7d48

Please sign in to comment.