import static com.google.common.base.Preconditions.checkState;

import com.google.cloud.bigquery.BigQuery;
import com.google.cloud.bigquery.BigQuery.TableDataListOption;
import com.google.cloud.bigquery.BigQueryOptions;
import com.google.cloud.bigquery.FieldValueList;
import com.google.cloud.bigquery.Table;
import com.google.cloud.bigquery.TableId;
import com.google.cloud.bigquery.TableResult;

public class BigQueryListDataTest {

    public static void main(String[] args) {
        BigQuery bigquery = BigQueryOptions.getDefaultInstance().getService();

        TableId tableId = TableId.of(
                "bigquery-public-data",
                "census_bureau_international",
                "midyear_population_agespecific");

        Table table = bigquery.getTable(tableId);
        long numRows = table.getNumRows().longValue();
        
        System.out.println("Table records # = " + numRows);

        checkRecordsConsumed(bigquery, tableId, numRows - 100_000L, 100_000L);
        checkRecordsConsumed(bigquery, tableId, numRows - 200_000L, 200_000L);
        // this will loop forever, although it should get to end of the table after 300k rows
        checkRecordsConsumed(bigquery, tableId, numRows - 300_000L, 300_000L);
    }

    private static void checkRecordsConsumed(BigQuery bigquery, TableId tableId, long startRow, long expectedRows) {
        TableResult tableResult = bigquery.listTableData(
                tableId,
                TableDataListOption.startIndex(startRow),
                TableDataListOption.pageSize(600_000L));

        System.out.println("Reading records starting row = " + startRow);
        int rowsConsumed = 0;
        for (FieldValueList row : tableResult.iterateAll()) {
            rowsConsumed++;
            if (rowsConsumed % 50_000 == 0) {
                System.out.println("% rows consumed: " + rowsConsumed);
            }
        }
        
        System.out.println("Consumed record count = " + rowsConsumed);

        checkState(rowsConsumed == expectedRows, "Expected %s records but iterated over %s", expectedRows, rowsConsumed);
    }

}