@Test
public void testRead() throws IOException {
onTableGet(basicTableSchema());
TableDataList dataList = rawDataList(rawRow("Arthur", 42));
onTableList(dataList);
try (BigQueryTableRowIterator iterator =
BigQueryTableRowIterator.of(
mockClient, BigQueryIO.parseTableSpec("project:dataset.table"))) {
Assert.assertTrue(iterator.hasNext());
TableRow row = iterator.next();
Assert.assertTrue(row.containsKey("name"));
Assert.assertTrue(row.containsKey("answer"));
Assert.assertEquals("Arthur", row.get("name"));
Assert.assertEquals(42, row.get("answer"));
Assert.assertFalse(iterator.hasNext());
verifyTableGet();
verifyTabledataList();
}
}
@Test
public void testReadMultiPage() throws IOException {
onTableGet(basicTableSchema());
TableDataList page1 = rawDataList(rawRow("Row1", 1)).setPageToken("page2");
TableDataList page2 = rawDataList(rawRow("Row2", 2)).setTotalRows(2L);
when(mockClient.tabledata()).thenReturn(mockTabledata);
when(mockTabledata.list(anyString(), anyString(), anyString())).thenReturn(mockTabledataList);
when(mockTabledataList.execute()).thenReturn(page1).thenReturn(page2);
try (BigQueryTableRowIterator iterator =
BigQueryTableRowIterator.of(
mockClient, BigQueryIO.parseTableSpec("project:dataset.table"))) {
List<String> names = new LinkedList<>();
Iterators.addAll(
names,
Iterators.transform(
iterator,
new Function<TableRow, String>() {
@Override
public String apply(TableRow input) {
return (String) input.get("name");
}
}));
Assert.assertThat(names, Matchers.hasItems("Row1", "Row2"));
verifyTableGet();
verifyTabledataList();
// The second call should have used a page token.
verify(mockTabledataList).setPageToken("page2");
}
}
@Test
public void testReadOpenFailure() throws IOException {
thrown.expect(RuntimeException.class);
when(mockClient.tables()).thenReturn(mockTables);
when(mockTables.get(anyString(), anyString(), anyString())).thenReturn(mockTablesGet);
when(mockTablesGet.execute()).thenThrow(new IOException("No such table"));
try (BigQueryTableRowIterator iterator =
BigQueryTableRowIterator.of(
mockClient, BigQueryIO.parseTableSpec("project:dataset.table"))) {
try {
Assert.assertFalse(iterator.hasNext()); // throws.
} finally {
verifyTableGet();
}
}
}
@Test
public void testReadEmpty() throws IOException {
onTableGet(basicTableSchema());
// BigQuery may respond with a page token for an empty table, ensure we
// handle it.
TableDataList dataList = new TableDataList().setPageToken("FEED==").setTotalRows(0L);
onTableList(dataList);
try (BigQueryTableRowIterator iterator =
BigQueryTableRowIterator.of(
mockClient, BigQueryIO.parseTableSpec("project:dataset.table"))) {
Assert.assertFalse(iterator.hasNext());
verifyTableGet();
verifyTabledataList();
}
}
