Examples with BigqueryClient org.apache.beam.sdk.io.gcp.testing.BigqueryClient used on opensource projects

Example 1 with BigqueryClient

use of org.apache.beam.sdk.io.gcp.testing.BigqueryClient in project beam by apache.

the class TrafficMaxLaneFlowIT method setupTestEnvironment.

@Before
public void setupTestEnvironment() throws Exception {
    PipelineOptionsFactory.register(TrafficMaxLaneFlowOptions.class);
    this.options = TestPipeline.testingPipelineOptions().as(TrafficMaxLaneFlowOptions.class);
    this.projectId = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject();
    this.bqClient = new BigqueryClient("TrafficMaxLaneFlowIT");
    this.bqClient.createNewDataset(this.projectId, this.outputDatasetId);
}

Example 2 with BigqueryClient

use of org.apache.beam.sdk.io.gcp.testing.BigqueryClient in project beam by apache.

the class BigQueryNestedRecordsIT method runPipeline.

private static void runPipeline(Options options) throws Exception {
    // Create flattened and unflattened collections via Dataflow, via normal and side input
    // paths.
    Pipeline p = Pipeline.create(options);
    BigQueryOptions bigQueryOptions = options.as(BigQueryOptions.class);
    PCollection<TableRow> flattenedCollection = p.apply("ReadFlattened", BigQueryIO.readTableRows().fromQuery(options.getInput()));
    PCollection<TableRow> nonFlattenedCollection = p.apply("ReadNonFlattened", BigQueryIO.readTableRows().fromQuery(options.getInput()).withoutResultFlattening());
    PCollection<TableRow> unflattenableCollection = p.apply("ReadUnflattenable", BigQueryIO.readTableRows().fromQuery(options.getUnflattenableInput()).withoutResultFlattening());
    // Also query BigQuery directly.
    BigqueryClient bigQueryClient = new BigqueryClient(bigQueryOptions.getAppName());
    TableRow queryFlattenedTyped = bigQueryClient.queryWithRetries(options.getInput(), bigQueryOptions.getProject(), true).getRows().get(0);
    TableRow queryUnflattened = bigQueryClient.queryUnflattened(options.getInput(), bigQueryOptions.getProject(), true).get(0);
    TableRow queryUnflattenable = bigQueryClient.queryUnflattened(options.getUnflattenableInput(), bigQueryOptions.getProject(), true).get(0);
    // Verify that the results are the same.
    PAssert.thatSingleton(flattenedCollection).isEqualTo(queryFlattenedTyped);
    PAssert.thatSingleton(nonFlattenedCollection).isEqualTo(queryUnflattened);
    PAssert.thatSingleton(unflattenableCollection).isEqualTo(queryUnflattenable);
    PAssert.thatSingleton(flattenedCollection).notEqualTo(queryUnflattened);
    p.run().waitUntilFinish();
}

Example 3 with BigqueryClient

use of org.apache.beam.sdk.io.gcp.testing.BigqueryClient in project beam by apache.

the class TrafficRoutesIT method setupTestEnvironment.

@Before
public void setupTestEnvironment() throws Exception {
    PipelineOptionsFactory.register(TrafficRoutesOptions.class);
    this.options = TestPipeline.testingPipelineOptions().as(TrafficRoutesOptions.class);
    this.projectId = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject();
    this.bqClient = new BigqueryClient("TrafficRoutesIT");
    this.bqClient.createNewDataset(this.projectId, this.outputDatasetId);
}