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

Example 61 with TableRow

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

the class BigQueryTableRowIteratorTest method testReadFromQueryNoTables.

/**
   * Verifies that queries that reference no data can be read.
   */
@Test
public void testReadFromQueryNoTables() throws IOException, InterruptedException {
    // Mock job inserting.
    Job dryRunJob = new Job().setStatistics(new JobStatistics().setQuery(new JobStatistics2()));
    Job insertedJob = new Job().setJobReference(new JobReference());
    when(mockJobsInsert.execute()).thenReturn(dryRunJob, insertedJob);
    // Mock job polling.
    JobStatus status = new JobStatus().setState("DONE");
    JobConfigurationQuery resultQueryConfig = new JobConfigurationQuery().setDestinationTable(new TableReference().setProjectId("project").setDatasetId("tempdataset").setTableId("temptable"));
    Job getJob = new Job().setJobReference(new JobReference()).setStatus(status).setConfiguration(new JobConfiguration().setQuery(resultQueryConfig));
    when(mockJobsGet.execute()).thenReturn(getJob);
    // Mock table schema fetch.
    when(mockTablesGet.execute()).thenReturn(noTableQuerySchema());
    byte[] photoBytes = "photograph".getBytes();
    String photoBytesEncoded = BaseEncoding.base64().encode(photoBytes);
    // Mock table data fetch.
    when(mockTabledataList.execute()).thenReturn(rawDataList(rawRow("Arthur", 42, photoBytesEncoded)));
    // Run query and verify
    String query = String.format("SELECT \"Arthur\" as name, 42 as count, \"%s\" as photo", photoBytesEncoded);
    JobConfigurationQuery queryConfig = new JobConfigurationQuery().setQuery(query);
    try (BigQueryTableRowIterator iterator = BigQueryTableRowIterator.fromQuery(queryConfig, "project", mockClient)) {
        iterator.open();
        assertTrue(iterator.advance());
        TableRow row = iterator.getCurrent();
        assertTrue(row.containsKey("name"));
        assertTrue(row.containsKey("count"));
        assertTrue(row.containsKey("photo"));
        assertEquals("Arthur", row.get("name"));
        assertEquals(42, row.get("count"));
        assertEquals(photoBytesEncoded, row.get("photo"));
        assertFalse(iterator.advance());
    }
    // Temp dataset created and later deleted.
    verify(mockClient, times(2)).datasets();
    verify(mockDatasets).insert(anyString(), any(Dataset.class));
    verify(mockDatasetsInsert).execute();
    verify(mockDatasets).delete(anyString(), anyString());
    verify(mockDatasetsDelete).execute();
    // Job inserted to run the query, polled once.
    verify(mockClient, times(3)).jobs();
    verify(mockJobs, times(2)).insert(anyString(), any(Job.class));
    verify(mockJobsInsert, times(2)).execute();
    verify(mockJobs).get(anyString(), anyString());
    verify(mockJobsGet).execute();
    // Temp table get after query finish, deleted after reading.
    verify(mockClient, times(2)).tables();
    verify(mockTables, times(1)).get(anyString(), anyString(), anyString());
    verify(mockTablesGet, times(1)).execute();
    verify(mockTables).delete(anyString(), anyString(), anyString());
    verify(mockTablesDelete).execute();
    // Table data read.
    verify(mockClient).tabledata();
    verify(mockTabledata).list("project", "tempdataset", "temptable");
    verify(mockTabledataList).execute();
}

Example 62 with TableRow

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

the class FakeJobService method runQueryJob.

private JobStatus runQueryJob(JobConfigurationQuery query) throws IOException, InterruptedException {
    List<TableRow> rows = FakeBigQueryServices.rowsFromEncodedQuery(query.getQuery());
    datasetService.createTable(new Table().setTableReference(query.getDestinationTable()));
    datasetService.insertAll(query.getDestinationTable(), rows, null);
    return new JobStatus().setState("DONE");
}

Example 63 with TableRow

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

the class BigQueryTornadoesTest method testExtractTornadoes.

@Test
public void testExtractTornadoes() throws Exception {
    TableRow row = new TableRow().set("month", "6").set("tornado", true);
    DoFnTester<TableRow, Integer> extractWordsFn = DoFnTester.of(new ExtractTornadoesFn());
    Assert.assertThat(extractWordsFn.processBundle(row), CoreMatchers.hasItems(6));
}

Example 64 with TableRow

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

the class CombinePerKeyExamplesTest method testExtractLargeWordsFn.

@Test
public void testExtractLargeWordsFn() throws Exception {
    DoFnTester<TableRow, KV<String, String>> extractLargeWordsFn = DoFnTester.of(new ExtractLargeWordsFn());
    List<KV<String, String>> results = extractLargeWordsFn.processBundle(ROWS_ARRAY);
    Assert.assertThat(results, CoreMatchers.hasItem(tuple1));
    Assert.assertThat(results, CoreMatchers.hasItem(tuple2));
    Assert.assertThat(results, CoreMatchers.hasItem(tuple3));
}

Example 65 with TableRow

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

the class TriggerExample method main.

public static void main(String[] args) throws Exception {
    TrafficFlowOptions options = PipelineOptionsFactory.fromArgs(args).withValidation().as(TrafficFlowOptions.class);
    options.setStreaming(true);
    options.setBigQuerySchema(getSchema());
    ExampleUtils exampleUtils = new ExampleUtils(options);
    exampleUtils.setup();
    Pipeline pipeline = Pipeline.create(options);
    TableReference tableRef = getTableReference(options.getProject(), options.getBigQueryDataset(), options.getBigQueryTable());
    PCollectionList<TableRow> resultList = pipeline.apply("ReadMyFile", TextIO.read().from(options.getInput())).apply("InsertRandomDelays", ParDo.of(new InsertDelays())).apply(ParDo.of(new ExtractFlowInfo())).apply(new CalculateTotalFlow(options.getWindowDuration()));
    for (int i = 0; i < resultList.size(); i++) {
        resultList.get(i).apply(BigQueryIO.writeTableRows().to(tableRef).withSchema(getSchema()));
    }
    PipelineResult result = pipeline.run();
    // ExampleUtils will try to cancel the pipeline and the injector before the program exits.
    exampleUtils.waitToFinish(result);
}