Examples with Write org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write used on opensource projects

Example 6 with Write

use of org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write in project beam by apache.

the class BigQueryIOWriteTest method testWriteValidateFailsBothFormatFunctions.

@Test
public void testWriteValidateFailsBothFormatFunctions() {
    if (useStorageApi) {
        return;
    }
    p.enableAbandonedNodeEnforcement(false);
    thrown.expect(IllegalArgumentException.class);
    thrown.expectMessage("Only one of withFormatFunction or withAvroFormatFunction/withAvroWriter maybe set, not both.");
    p.apply(Create.empty(INPUT_RECORD_CODER)).apply(BigQueryIO.<InputRecord>write().to("dataset.table").withSchema(new TableSchema()).withFormatFunction(r -> new TableRow()).withAvroFormatFunction(r -> new GenericData.Record(r.getSchema())).withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED));
}

Example 7 with Write

use of org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write in project beam by apache.

the class BigQuerySchemaUpdateOptionsIT method runWriteTest.

/**
 * Runs a write test against a BigQuery table to check that SchemaUpdateOption sets are taking
 * effect.
 *
 * <p>Attempt write a row via BigQueryIO.writeTables with the given params, then run the given
 * query, and finaly check the results of the query.
 *
 * @param schemaUpdateOptions The SchemaUpdateOption set to use
 * @param tableName The table to write to
 * @param schema The schema to use for the table
 * @param rowToInsert The row to insert
 * @param testQuery A testing SQL query to run after writing the row
 * @param expectedResult The expected result of the query as a nested list of column values (one
 *     list per result row)
 */
private void runWriteTest(Set<SchemaUpdateOption> schemaUpdateOptions, String tableName, TableSchema schema, TableRow rowToInsert, String testQuery, List<List<String>> expectedResult) throws Exception {
    Options options = TestPipeline.testingPipelineOptions().as(Options.class);
    options.setTempLocation(options.getTempRoot() + "/bq_it_temp");
    Pipeline p = Pipeline.create(options);
    Create.Values<TableRow> input = Create.<TableRow>of(rowToInsert);
    Write<TableRow> writer = BigQueryIO.writeTableRows().to(String.format("%s:%s.%s", options.getProject(), BIG_QUERY_DATASET_ID, tableName)).withSchema(schema).withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED).withWriteDisposition(BigQueryIO.Write.WriteDisposition.WRITE_APPEND).withSchemaUpdateOptions(schemaUpdateOptions);
    p.apply(input).apply(writer);
    p.run().waitUntilFinish();
    QueryResponse response = BQ_CLIENT.queryWithRetries(testQuery, project);
    List<List<String>> result = response.getRows().stream().map(row -> row.getF().stream().map(cell -> cell.getV().toString()).collect(Collectors.toList())).collect(Collectors.toList());
    assertEquals(expectedResult, result);
}