Examples with TableDataList com.google.api.services.bigquery.model.TableDataList used on opensource projects

Example 1 with TableDataList

use of com.google.api.services.bigquery.model.TableDataList in project beam by apache.

the class BigQueryTableRowIterator method advance.

boolean advance() throws IOException, InterruptedException {
    while (true) {
        if (iteratorOverCurrentBatch != null && iteratorOverCurrentBatch.hasNext()) {
            // Embed schema information into the raw row, so that values have an
            // associated key.
            current = getTypedTableRow(schema.getFields(), iteratorOverCurrentBatch.next());
            return true;
        }
        if (lastPage) {
            return false;
        }
        Bigquery.Tabledata.List list = client.tabledata().list(ref.getProjectId(), ref.getDatasetId(), ref.getTableId());
        if (pageToken != null) {
            list.setPageToken(pageToken);
        }
        TableDataList result = executeWithBackOff(list, String.format("Error reading from BigQuery table %s of dataset %s.", ref.getTableId(), ref.getDatasetId()));
        pageToken = result.getPageToken();
        iteratorOverCurrentBatch = result.getRows() != null ? result.getRows().iterator() : Collections.<TableRow>emptyIterator();
        // The server may return a page token indefinitely on a zero-length table.
        if (pageToken == null || result.getTotalRows() != null && result.getTotalRows() == 0) {
            lastPage = true;
        }
    }
}

Example 2 with TableDataList

use of com.google.api.services.bigquery.model.TableDataList in project beam by apache.

the class BigQueryUtilTest method testReadWithTime.

@Test
public void testReadWithTime() throws IOException, InterruptedException {
    // The BigQuery JSON API returns timestamps in the following format: floating-point seconds
    // since epoch (UTC) with microsecond precision. Test that we faithfully preserve a set of
    // known values.
    TableDataList input = rawDataList(rawRow("1.430397296789E9"), rawRow("1.45206228E9"), rawRow("1.452062291E9"), rawRow("1.4520622911E9"), rawRow("1.45206229112E9"), rawRow("1.452062291123E9"), rawRow("1.4520622911234E9"), rawRow("1.45206229112345E9"), rawRow("1.452062291123456E9"));
    onTableGet(basicTableSchemaWithTime());
    onTableList(input);
    // Known results verified from BigQuery's export to JSON on GCS API.
    List<String> expected = ImmutableList.of("2015-04-30 12:34:56.789 UTC", "2016-01-06 06:38:00 UTC", "2016-01-06 06:38:11 UTC", "2016-01-06 06:38:11.1 UTC", "2016-01-06 06:38:11.12 UTC", "2016-01-06 06:38:11.123 UTC", "2016-01-06 06:38:11.1234 UTC", "2016-01-06 06:38:11.12345 UTC", "2016-01-06 06:38:11.123456 UTC");
    // Download the rows, verify the interactions.
    List<TableRow> rows = new ArrayList<>();
    try (BigQueryTableRowIterator iterator = BigQueryTableRowIterator.fromTable(BigQueryHelpers.parseTableSpec("project:dataset.table"), mockClient)) {
        iterator.open();
        while (iterator.advance()) {
            rows.add(iterator.getCurrent());
        }
    }
    verifyTableGet();
    verifyTabledataList();
    // Verify the timestamp converted as desired.
    assertEquals("Expected input and output rows to have the same size", expected.size(), rows.size());
    for (int i = 0; i < expected.size(); ++i) {
        assertEquals("i=" + i, expected.get(i), rows.get(i).get("time"));
    }
}

Example 3 with TableDataList

use of com.google.api.services.bigquery.model.TableDataList in project beam by apache.

the class BigQueryUtilTest method testReadEmpty.

@Test
public void testReadEmpty() throws IOException, InterruptedException {
    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.fromTable(BigQueryHelpers.parseTableSpec("project:dataset.table"), mockClient)) {
        iterator.open();
        Assert.assertFalse(iterator.advance());
        verifyTableGet();
        verifyTabledataList();
    }
}

Example 4 with TableDataList

use of com.google.api.services.bigquery.model.TableDataList in project beam by apache.

the class BigQueryUtilTest method testReadMultiPage.

@Test
public void testReadMultiPage() throws IOException, InterruptedException {
    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.fromTable(BigQueryHelpers.parseTableSpec("project:dataset.table"), mockClient)) {
        iterator.open();
        List<String> names = new LinkedList<>();
        while (iterator.advance()) {
            names.add((String) iterator.getCurrent().get("name"));
        }
        Assert.assertThat(names, Matchers.hasItems("Row1", "Row2"));
        verifyTableGet();
        verifyTabledataList();
        // The second call should have used a page token.
        verify(mockTabledataList).setPageToken("page2");
    }
}

Example 5 with TableDataList

use of com.google.api.services.bigquery.model.TableDataList in project beam by apache.

the class BigQueryUtilTest method testRead.

@Test
public void testRead() throws IOException, InterruptedException {
    onTableGet(basicTableSchema());
    TableDataList dataList = rawDataList(rawRow("Arthur", 42));
    onTableList(dataList);
    try (BigQueryTableRowIterator iterator = BigQueryTableRowIterator.fromTable(BigQueryHelpers.parseTableSpec("project:dataset.table"), mockClient)) {
        iterator.open();
        Assert.assertTrue(iterator.advance());
        TableRow row = iterator.getCurrent();
        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.advance());
        verifyTableGet();
        verifyTabledataList();
    }
}