Skip to content

Instantly share code, notes, and snippets.

@polleyg
Created May 31, 2019 14:17
Show Gist options
  • Save polleyg/240f65a718deebd8b7fb410fb1e303f6 to your computer and use it in GitHub Desktop.
Save polleyg/240f65a718deebd8b7fb410fb1e303f6 to your computer and use it in GitHub Desktop.
A trivial Cloud Dataflow pipeline that reads from Spanner and writes to BigQuery
package org.polleyg;
import com.google.api.services.bigquery.model.TableFieldSchema;
import com.google.api.services.bigquery.model.TableRow;
import com.google.api.services.bigquery.model.TableSchema;
import com.google.cloud.spanner.Struct;
import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO;
import org.apache.beam.sdk.io.gcp.spanner.SpannerIO;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.values.PCollection;
import java.util.ArrayList;
import java.util.List;
import static org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED;
import static org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.WriteDisposition.WRITE_TRUNCATE;
/**
* Do some randomness
*/
public class TemplatePipeline {
public static void main(String[] args) {
PipelineOptionsFactory.register(DataflowPipelineOptions.class);
DataflowPipelineOptions options = PipelineOptionsFactory.fromArgs(args).withValidation().as(DataflowPipelineOptions.class);
Pipeline pipeline = Pipeline.create(options);
PCollection<Struct> records = pipeline.apply("read_from_spanner",
SpannerIO.read()
.withInstanceId("spanner-to-dataflow-to-bq")
.withDatabaseId("the-dude")
.withQuery("SELECT * FROM Singers"));
records.apply("convert-2-bq-row", ParDo.of(new DoFn<Struct, TableRow>() {
@ProcessElement
public void processElement(ProcessContext c) throws Exception {
TableRow row = new TableRow();
row.set("id", c.element().getLong("SingerId"));
row.set("first", c.element().getString("FirstName"));
row.set("last", c.element().getString("LastName"));
c.output(row);
}
})).apply("write-to-bq", BigQueryIO.writeTableRows()
.to(String.format("%s:spanner_to_bigquery.singers", options.getProject()))
.withCreateDisposition(CREATE_IF_NEEDED)
.withWriteDisposition(WRITE_TRUNCATE)
.withSchema(getTableSchema()));
pipeline.run();
}
private static TableSchema getTableSchema() {
List<TableFieldSchema> fields = new ArrayList<>();
fields.add(new TableFieldSchema().setName("id").setType("INTEGER"));
fields.add(new TableFieldSchema().setName("first").setType("STRING"));
fields.add(new TableFieldSchema().setName("last").setType("STRING"));
return new TableSchema().setFields(fields);
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment