use of com.google.cloud.teleport.spanner.ddl.ChangeStream in project DataflowTemplates by GoogleCloudPlatform.
the class ExportTransform method expand.
/**
* Read the Cloud Spanner schema and all the rows in all the tables of the database. Create and
* write the exported Avro files to GCS.
*/
@Override
public WriteFilesResult<String> expand(PBegin begin) {
Pipeline p = begin.getPipeline();
/*
* Allow users to specify read timestamp.
* CreateTransaction and CreateTransactionFn classes in SpannerIO
* only take a timestamp object for exact staleness which works when
* parameters are provided during template compile time. They do not work with
* a Timestamp valueProvider which can take parameters at runtime. Hence a new
* ParDo class CreateTransactionFnWithTimestamp had to be created for this
* purpose.
*/
PCollectionView<Transaction> tx = p.apply("CreateTransaction", Create.of(1)).apply("Create transaction", ParDo.of(new CreateTransactionFnWithTimestamp(spannerConfig, snapshotTime))).apply("Tx As PCollectionView", View.asSingleton());
PCollectionView<Dialect> dialectView = p.apply("Read Dialect", new ReadDialect(spannerConfig)).apply("Dialect As PCollectionView", View.asSingleton());
PCollection<Ddl> ddl = p.apply("Read Information Schema", new ReadInformationSchema(spannerConfig, tx, dialectView));
PCollection<Ddl> exportState = ddl.apply("Check export conditions", ParDo.of(new DoFn<Ddl, Ddl>() {
@ProcessElement
public void processElement(ProcessContext c) throws Exception {
Ddl ddl = c.element();
List<String> tablesList = Collections.emptyList();
// a list of export tables, throw an exception.
if (tableNames.get().trim().isEmpty() && exportRelatedTables.get()) {
throw new Exception("Invalid usage of --tableNames and --shouldExportRelatedTables. Set" + " --shouldExportRelatedTables=true only if --tableNames is given" + " selected tables for export.");
}
// If the user provides a comma-separated list of strings, parse it into a List
if (!tableNames.get().trim().isEmpty()) {
tablesList = Arrays.asList(tableNames.get().split(",\\s*"));
}
// If the user provided any invalid table names, throw an exception.
List<String> allSpannerTables = ddl.allTables().stream().map(t -> t.name()).collect(Collectors.toList());
List<String> invalidTables = tablesList.stream().distinct().filter(t -> !allSpannerTables.contains(t)).collect(Collectors.toList());
if (invalidTables.size() != 0) {
throw new Exception("INVALID_ARGUMENT: Table(s) not found: " + String.join(", ", invalidTables) + ".");
}
List<String> filteredTables = getFilteredTables(ddl, tablesList).stream().map(t -> t.name()).collect(Collectors.toList());
// Save any missing necessary export table names; save a copy of the original
// table list to bypass 'final or effectively final' condition of the lambda
// expression below.
List<String> usersTables = tablesList.stream().collect(Collectors.toList());
List<String> missingTables = filteredTables.stream().distinct().filter(t -> !usersTables.contains(t)).collect(Collectors.toList());
Collections.sort(missingTables);
// throw an exception.
if (tablesList.size() != 0 && !(tablesList.equals(filteredTables)) && !exportRelatedTables.get()) {
throw new Exception("Attempted to export table(s) requiring parent and/or foreign keys tables" + " without setting the shouldExportRelatedTables parameter. Set" + " --shouldExportRelatedTables=true to export all necessary" + " tables, or add " + String.join(", ", missingTables) + " to --tableNames.");
}
c.output(ddl);
}
}));
PCollection<ReadOperation> tables = ddl.apply("Build table read operations", new BuildReadFromTableOperations(tableNames));
PCollection<KV<String, Void>> allTableAndViewNames = ddl.apply("List all table and view names", ParDo.of(new DoFn<Ddl, KV<String, Void>>() {
@ProcessElement
public void processElement(ProcessContext c) {
Ddl ddl = c.element();
for (Table t : ddl.allTables()) {
c.output(KV.of(t.name(), null));
}
// we need to add the names of all views separately here.
for (com.google.cloud.teleport.spanner.ddl.View v : ddl.views()) {
c.output(KV.of(v.name(), null));
}
}
}));
PCollection<String> allChangeStreamNames = ddl.apply("List all change stream names", ParDo.of(new DoFn<Ddl, String>() {
@ProcessElement
public void processElement(ProcessContext c) {
Ddl ddl = c.element();
for (ChangeStream changeStream : ddl.changeStreams()) {
c.output(changeStream.name());
}
}
}));
// Generate a unique output directory name.
final PCollectionView<String> outputDirectoryName = p.apply(Create.of(1)).apply("Create Avro output folder", ParDo.of(new DoFn<Integer, String>() {
@ProcessElement
public void processElement(ProcessContext c) {
String instanceId = spannerConfig.getInstanceId().get();
String dbId = spannerConfig.getDatabaseId().get();
// For direct runner or tests we need a deterministic jobId.
String testJobId = ExportTransform.this.testJobId.get();
if (!Strings.isNullOrEmpty(testJobId)) {
c.output(testJobId);
return;
}
try {
DataflowWorkerHarnessOptions workerHarnessOptions = c.getPipelineOptions().as(DataflowWorkerHarnessOptions.class);
String jobId = workerHarnessOptions.getJobId();
c.output(instanceId + "-" + dbId + "-" + jobId);
} catch (Exception e) {
throw new IllegalStateException("Please specify --testJobId to run with non-dataflow runner");
}
}
})).apply(View.asSingleton());
final PCollectionView<Map<String, SerializableSchemaSupplier>> avroSchemas = ddl.apply("Build Avro schemas from DDL", ParDo.of(new DoFn<Ddl, KV<String, SerializableSchemaSupplier>>() {
@ProcessElement
public void processElement(ProcessContext c) {
Collection<Schema> avroSchemas = new DdlToAvroSchemaConverter("spannerexport", "1.0.0", shouldExportTimestampAsLogicalType.get()).convert(c.element());
for (Schema schema : avroSchemas) {
c.output(KV.of(schema.getName(), new SerializableSchemaSupplier(schema)));
}
}
})).apply("As view", View.asMap());
PCollection<Struct> rows = tables.apply("Read all rows from Spanner", SpannerIO.readAll().withTransaction(tx).withSpannerConfig(spannerConfig));
ValueProvider<ResourceId> resource = ValueProvider.NestedValueProvider.of(outputDir, (SerializableFunction<String, ResourceId>) s -> FileSystems.matchNewResource(s, true));
ValueProvider<ResourceId> tempResource = ValueProvider.NestedValueProvider.of(eitherOrValueProvider(avroTempDirectory, outputDir), (SerializableFunction<String, ResourceId>) s -> FileSystems.matchNewResource(s, true));
WriteFilesResult<String> fileWriteResults = rows.apply("Store Avro files", AvroIO.<Struct>writeCustomTypeToGenericRecords().to(new SchemaBasedDynamicDestinations(avroSchemas, outputDirectoryName, dialectView, resource)).withTempDirectory(tempResource));
// Generate the manifest file.
PCollection<KV<String, Iterable<String>>> tableFiles = fileWriteResults.getPerDestinationOutputFilenames().apply(GroupByKey.create());
final TupleTag<Void> allTables = new TupleTag<>();
final TupleTag<Iterable<String>> nonEmptyTables = new TupleTag<>();
PCollection<KV<String, CoGbkResult>> groupedTables = KeyedPCollectionTuple.of(allTables, allTableAndViewNames).and(nonEmptyTables, tableFiles).apply("Group with all tables", CoGroupByKey.create());
// The following is to export empty tables and views from the database. Empty tables and views
// are handled together because we do not export any rows for views, only their metadata,
// including the queries defining them.
PCollection<KV<String, Iterable<String>>> emptyTablesAndViews = groupedTables.apply("Export empty tables and views", ParDo.of(new DoFn<KV<String, CoGbkResult>, KV<String, Iterable<String>>>() {
@ProcessElement
public void processElement(ProcessContext c) {
KV<String, CoGbkResult> kv = c.element();
String table = kv.getKey();
CoGbkResult coGbkResult = kv.getValue();
Iterable<String> only = coGbkResult.getOnly(nonEmptyTables, null);
if (only == null) {
LOG.info("Exporting empty table or view: " + table);
// This file will contain the schema definition: column definitions for empty
// tables or defining queries for views.
c.output(KV.of(table, Collections.singleton(table + ".avro-00000-of-00001")));
}
}
}));
PCollection<KV<String, Iterable<String>>> changeStreams = allChangeStreamNames.apply("Export change streams", ParDo.of(new DoFn<String, KV<String, Iterable<String>>>() {
@ProcessElement
public void processElement(ProcessContext c) {
String changeStreamName = c.element();
LOG.info("Exporting change stream: " + changeStreamName);
// This file will contain the schema definition for the change stream.
c.output(KV.of(changeStreamName, Collections.singleton(changeStreamName + ".avro-00000-of-00001")));
}
}));
// Empty tables, views and change streams are handled together, because we export them as empty
// Avro files that only contain the Avro schemas.
PCollection<KV<String, Iterable<String>>> emptySchemaFiles = PCollectionList.of(emptyTablesAndViews).and(changeStreams).apply("Combine all empty schema files", Flatten.pCollections());
emptySchemaFiles = emptySchemaFiles.apply("Save empty schema files", ParDo.of(new DoFn<KV<String, Iterable<String>>, KV<String, Iterable<String>>>() {
@ProcessElement
public void processElement(ProcessContext c) {
Map<String, SerializableSchemaSupplier> schemaMap = c.sideInput(avroSchemas);
KV<String, Iterable<String>> kv = c.element();
String objectName = kv.getKey();
String fileName = kv.getValue().iterator().next();
Schema schema = schemaMap.get(objectName).get();
DatumWriter<GenericRecord> datumWriter = new GenericDatumWriter<>(schema);
Path fullPath = createOutputPath(outputDir.get(), c.sideInput(outputDirectoryName), fileName);
try (DataFileWriter<GenericRecord> dataFileWriter = new DataFileWriter<>(datumWriter)) {
dataFileWriter.create(schema, createOutputStream(fullPath, c));
} catch (IOException e) {
throw new RuntimeException(e);
}
c.output(KV.of(objectName, Collections.singleton(fullPath.toString())));
}
/**
* Resolves the complete path name for Avro files for both GCS and local FS
* (for testing).
*
* @param outputDirectoryPath Initial directory path for the file.
* @param outputDirectoryName Terminal directory for the file.
* @param fileName Name of the Avro file
* @return The full {@link Path} of the output Avro file.
*/
private Path createOutputPath(String outputDirectoryPath, String outputDirectoryName, String fileName) {
if (GcsPath.GCS_URI.matcher(outputDirectoryPath).matches()) {
// Avro file path in GCS.
return GcsPath.fromUri(outputDirectoryPath).resolve(outputDirectoryName).resolve(fileName);
} else {
// Avro file path in local filesystem
return Paths.get(outputDirectoryPath, outputDirectoryName, fileName);
}
}
/**
* Creates the {@link OutputStream} for the output file either on GCS or on
* local FS (for testing).
*
* @param outputPath The full path of the output file.
* @param c The {@link org.apache.beam.sdk.transforms.DoFn.ProcessContext}
* @return An {@link OutputStream} for the opened output file.
* @throws IOException if the output file cannot be opened.
*/
private OutputStream createOutputStream(Path outputPath, ProcessContext c) throws IOException {
if (GcsPath.GCS_URI.matcher(outputPath.toString()).matches()) {
// Writing the Avro file to GCS.
org.apache.beam.sdk.extensions.gcp.util.GcsUtil gcsUtil = c.getPipelineOptions().as(GcsOptions.class).getGcsUtil();
String gcsType = "application/octet-stream";
WritableByteChannel gcsChannel = gcsUtil.create((GcsPath) outputPath, gcsType);
return Channels.newOutputStream(gcsChannel);
} else {
// Avro file is created on local filesystem (for testing).
Files.createDirectories(outputPath.getParent());
return Files.newOutputStream(outputPath);
}
}
}).withSideInputs(avroSchemas, outputDirectoryName));
PCollection<KV<String, Iterable<String>>> allFiles = PCollectionList.of(tableFiles).and(emptySchemaFiles).apply("Combine all files", Flatten.pCollections());
PCollection<KV<String, String>> tableManifests = allFiles.apply("Build table manifests", ParDo.of(new BuildTableManifests()));
Contextful.Fn<String, FileIO.Write.FileNaming> tableManifestNaming = (element, c) -> (window, pane, numShards, shardIndex, compression) -> GcsUtil.joinPath(outputDir.get(), c.sideInput(outputDirectoryName), tableManifestFileName(element));
tableManifests.apply("Store table manifests", FileIO.<String, KV<String, String>>writeDynamic().by(KV::getKey).withDestinationCoder(StringUtf8Coder.of()).withNaming(Contextful.of(tableManifestNaming, Requirements.requiresSideInputs(outputDirectoryName))).via(Contextful.fn(KV::getValue), TextIO.sink()).withTempDirectory(eitherOrValueProvider(avroTempDirectory, outputDir)));
PCollection<List<Export.Table>> metadataTables = tableManifests.apply("Combine table metadata", Combine.globally(new CombineTableMetadata()));
PCollectionView<Ddl> ddlView = ddl.apply("Cloud Spanner DDL as view", View.asSingleton());
PCollection<String> metadataContent = metadataTables.apply("Create database manifest", ParDo.of(new CreateDatabaseManifest(ddlView, dialectView)).withSideInputs(ddlView, dialectView));
Contextful.Fn<String, FileIO.Write.FileNaming> manifestNaming = (element, c) -> (window, pane, numShards, shardIndex, compression) -> GcsUtil.joinPath(outputDir.get(), c.sideInput(outputDirectoryName), "spanner-export.json");
metadataContent.apply("Store the database manifest", FileIO.<String, String>writeDynamic().by(SerializableFunctions.constant("")).withDestinationCoder(StringUtf8Coder.of()).via(TextIO.sink()).withNaming(Contextful.of(manifestNaming, Requirements.requiresSideInputs(outputDirectoryName))).withTempDirectory(eitherOrValueProvider(avroTempDirectory, outputDir)));
return fileWriteResults;
}
use of com.google.cloud.teleport.spanner.ddl.ChangeStream in project DataflowTemplates by GoogleCloudPlatform.
the class DdlToAvroSchemaConverter method convert.
public Collection<Schema> convert(Ddl ddl) {
Collection<Schema> schemas = new ArrayList<>();
for (Table table : ddl.allTables()) {
SchemaBuilder.RecordBuilder<Schema> recordBuilder = SchemaBuilder.record(table.name()).namespace(this.namespace);
recordBuilder.prop("googleFormatVersion", version);
recordBuilder.prop("googleStorage", "CloudSpanner");
if (table.interleaveInParent() != null) {
recordBuilder.prop("spannerParent", table.interleaveInParent());
recordBuilder.prop("spannerOnDeleteAction", table.onDeleteCascade() ? "cascade" : "no action");
}
if (table.dialect() == Dialect.GOOGLE_STANDARD_SQL) {
if (table.primaryKeys() != null) {
String encodedPk = table.primaryKeys().stream().map(IndexColumn::prettyPrint).collect(Collectors.joining(","));
recordBuilder.prop("spannerPrimaryKey", encodedPk);
}
for (int i = 0; i < table.primaryKeys().size(); i++) {
recordBuilder.prop("spannerPrimaryKey_" + i, table.primaryKeys().get(i).prettyPrint());
}
} else if (table.dialect() == Dialect.POSTGRESQL) {
if (table.primaryKeys() != null) {
String encodedPk = table.primaryKeys().stream().map(c -> "\"" + c.name() + "\"").collect(Collectors.joining(", "));
recordBuilder.prop("spannerPrimaryKey", encodedPk);
}
for (int i = 0; i < table.primaryKeys().size(); i++) {
IndexColumn pk = table.primaryKeys().get(i);
recordBuilder.prop("spannerPrimaryKey_" + i, "\"" + pk.name() + "\" ASC");
}
}
for (int i = 0; i < table.indexes().size(); i++) {
recordBuilder.prop("spannerIndex_" + i, table.indexes().get(i));
}
for (int i = 0; i < table.foreignKeys().size(); i++) {
recordBuilder.prop("spannerForeignKey_" + i, table.foreignKeys().get(i));
}
for (int i = 0; i < table.checkConstraints().size(); i++) {
recordBuilder.prop("spannerCheckConstraint_" + i, table.checkConstraints().get(i));
}
SchemaBuilder.FieldAssembler<Schema> fieldsAssembler = recordBuilder.fields();
for (Column cm : table.columns()) {
SchemaBuilder.FieldBuilder<Schema> fieldBuilder = fieldsAssembler.name(cm.name());
fieldBuilder.prop("sqlType", cm.typeString());
for (int i = 0; i < cm.columnOptions().size(); i++) {
fieldBuilder.prop("spannerOption_" + i, cm.columnOptions().get(i));
}
if (cm.isGenerated()) {
fieldBuilder.prop("notNull", Boolean.toString(cm.notNull()));
fieldBuilder.prop("generationExpression", cm.generationExpression());
fieldBuilder.prop("stored", Boolean.toString(cm.isStored()));
// Make the type null to allow us not export the generated column values,
// which are semantically logical entities.
fieldBuilder.type(SchemaBuilder.builder().nullType()).withDefault(null);
} else {
if (cm.defaultExpression() != null) {
fieldBuilder.prop("defaultExpression", cm.defaultExpression());
}
Schema avroType = avroType(cm.type());
if (!cm.notNull()) {
avroType = wrapAsNullable(avroType);
}
fieldBuilder.type(avroType).noDefault();
}
}
Schema schema = fieldsAssembler.endRecord();
schemas.add(schema);
}
for (View view : ddl.views()) {
SchemaBuilder.RecordBuilder<Schema> recordBuilder = SchemaBuilder.record(view.name()).namespace(this.namespace);
recordBuilder.prop("googleFormatVersion", version);
recordBuilder.prop("googleStorage", "CloudSpanner");
recordBuilder.prop("spannerViewQuery", view.query());
if (view.security() != null) {
recordBuilder.prop("spannerViewSecurity", view.security().toString());
}
schemas.add(recordBuilder.fields().endRecord());
}
for (ChangeStream changeStream : ddl.changeStreams()) {
SchemaBuilder.RecordBuilder<Schema> recordBuilder = SchemaBuilder.record(changeStream.name()).namespace(this.namespace);
recordBuilder.prop("googleFormatVersion", version);
recordBuilder.prop("googleStorage", "CloudSpanner");
recordBuilder.prop(AvroUtil.CHANGE_STREAM_FOR_CLAUSE, changeStream.forClause() == null ? "" : changeStream.forClause());
if (changeStream.options() != null) {
for (int i = 0; i < changeStream.options().size(); i++) {
recordBuilder.prop("spannerOption_" + i, changeStream.options().get(i));
}
}
schemas.add(recordBuilder.fields().endRecord());
}
return schemas;
}
Aggregations