Search in sources :

Example 1 with Builder

use of com.google.cloud.teleport.spanner.ExportProtos.Export.Builder in project DataflowTemplates by GoogleCloudPlatform.

the class ExportTransformTest method buildDatabaseManifestFile.

@Test
public void buildDatabaseManifestFile() throws InvalidProtocolBufferException {
    Map<String, String> tablesAndManifests = ImmutableMap.of("table1", "table1 manifest", "table2", "table2 manifest", "changeStream", "changeStream manifest");
    PCollection<List<Export.Table>> metadataTables = pipeline.apply("Initialize table manifests", Create.of(tablesAndManifests)).apply("Combine table manifests", Combine.globally(new CombineTableMetadata()));
    ImmutableList<Export.DatabaseOption> databaseOptions = ImmutableList.of(Export.DatabaseOption.newBuilder().setOptionName("version_retention_period").setOptionValue("5d").build());
    Ddl.Builder ddlBuilder = Ddl.builder();
    ddlBuilder.mergeDatabaseOptions(databaseOptions);
    ddlBuilder.createChangeStream("changeStream").endChangeStream();
    Ddl ddl = ddlBuilder.build();
    PCollectionView<Ddl> ddlView = pipeline.apply(Create.of(ddl)).apply(View.asSingleton());
    PCollectionView<Dialect> dialectView = pipeline.apply("CreateSingleton", Create.of(Dialect.GOOGLE_STANDARD_SQL)).apply("As PCollectionView", View.asSingleton());
    PCollection<String> databaseManifest = metadataTables.apply("Test adding database option to manifest", ParDo.of(new CreateDatabaseManifest(ddlView, dialectView)).withSideInputs(ddlView, dialectView));
    // The output JSON may contain the tables in any order, so a string comparison is not
    // sufficient. Have to convert the manifest string to a protobuf. Also for the checker function
    // to be serializable, it has to be written as a lambda.
    PAssert.thatSingleton(databaseManifest).satisfies((SerializableFunction<String, Void>) input -> {
        Builder builder1 = Export.newBuilder();
        try {
            JsonFormat.parser().merge(input, builder1);
        } catch (InvalidProtocolBufferException e) {
            throw new RuntimeException(e);
        }
        Export manifestProto = builder1.build();
        assertThat(manifestProto.getTablesCount(), is(2));
        assertThat(manifestProto.getDialect(), is(ProtoDialect.GOOGLE_STANDARD_SQL));
        String table1Name = manifestProto.getTables(0).getName();
        assertThat(table1Name, startsWith("table"));
        assertThat(manifestProto.getTables(0).getManifestFile(), is(table1Name + "-manifest.json"));
        Export.DatabaseOption dbOptions = manifestProto.getDatabaseOptions(0);
        String optionName = dbOptions.getOptionName();
        String optionValue = dbOptions.getOptionValue();
        assertThat(optionName, is("version_retention_period"));
        assertThat(optionValue, is("5d"));
        assertThat(manifestProto.getChangeStreamsCount(), is(1));
        assertThat(manifestProto.getChangeStreams(0).getName(), is("changeStream"));
        assertThat(manifestProto.getChangeStreams(0).getManifestFile(), is("changeStream-manifest.json"));
        return null;
    });
    pipeline.run();
}
Also used : Assertions.assertThrows(org.junit.jupiter.api.Assertions.assertThrows) KV(org.apache.beam.sdk.values.KV) CombineTableMetadata(com.google.cloud.teleport.spanner.ExportTransform.CombineTableMetadata) Dialect(com.google.cloud.spanner.Dialect) Combine(org.apache.beam.sdk.transforms.Combine) SerializableFunction(org.apache.beam.sdk.transforms.SerializableFunction) View(org.apache.beam.sdk.transforms.View) Timestamp(com.google.cloud.Timestamp) BuildTableManifests(com.google.cloud.teleport.spanner.ExportTransform.BuildTableManifests) Assert.assertThat(org.junit.Assert.assertThat) Builder(com.google.cloud.teleport.spanner.ExportProtos.Export.Builder) ImmutableList(com.google.common.collect.ImmutableList) Create(org.apache.beam.sdk.transforms.Create) Map(java.util.Map) TestPipeline(org.apache.beam.sdk.testing.TestPipeline) TableManifest(com.google.cloud.teleport.spanner.ExportProtos.TableManifest) Ddl(com.google.cloud.teleport.spanner.ddl.Ddl) Path(java.nio.file.Path) InvalidProtocolBufferException(com.google.protobuf.InvalidProtocolBufferException) ImmutableMap(com.google.common.collect.ImmutableMap) Files(java.nio.file.Files) PAssert(org.apache.beam.sdk.testing.PAssert) ProtoDialect(com.google.cloud.teleport.spanner.ExportProtos.ProtoDialect) Test(org.junit.Test) TimestampBound(com.google.cloud.spanner.TimestampBound) PCollection(org.apache.beam.sdk.values.PCollection) Matchers.startsWith(org.hamcrest.Matchers.startsWith) Export(com.google.cloud.teleport.spanner.ExportProtos.Export) List(java.util.List) Rule(org.junit.Rule) CreateDatabaseManifest(com.google.cloud.teleport.spanner.ExportTransform.CreateDatabaseManifest) JsonFormat(com.google.protobuf.util.JsonFormat) ParDo(org.apache.beam.sdk.transforms.ParDo) PCollectionView(org.apache.beam.sdk.values.PCollectionView) Matchers.is(org.hamcrest.Matchers.is) Assert.assertEquals(org.junit.Assert.assertEquals) CombineTableMetadata(com.google.cloud.teleport.spanner.ExportTransform.CombineTableMetadata) Builder(com.google.cloud.teleport.spanner.ExportProtos.Export.Builder) InvalidProtocolBufferException(com.google.protobuf.InvalidProtocolBufferException) Ddl(com.google.cloud.teleport.spanner.ddl.Ddl) CreateDatabaseManifest(com.google.cloud.teleport.spanner.ExportTransform.CreateDatabaseManifest) Dialect(com.google.cloud.spanner.Dialect) ProtoDialect(com.google.cloud.teleport.spanner.ExportProtos.ProtoDialect) Export(com.google.cloud.teleport.spanner.ExportProtos.Export) ImmutableList(com.google.common.collect.ImmutableList) List(java.util.List) Test(org.junit.Test)

Aggregations

Timestamp (com.google.cloud.Timestamp)1 Dialect (com.google.cloud.spanner.Dialect)1 TimestampBound (com.google.cloud.spanner.TimestampBound)1 Export (com.google.cloud.teleport.spanner.ExportProtos.Export)1 Builder (com.google.cloud.teleport.spanner.ExportProtos.Export.Builder)1 ProtoDialect (com.google.cloud.teleport.spanner.ExportProtos.ProtoDialect)1 TableManifest (com.google.cloud.teleport.spanner.ExportProtos.TableManifest)1 BuildTableManifests (com.google.cloud.teleport.spanner.ExportTransform.BuildTableManifests)1 CombineTableMetadata (com.google.cloud.teleport.spanner.ExportTransform.CombineTableMetadata)1 CreateDatabaseManifest (com.google.cloud.teleport.spanner.ExportTransform.CreateDatabaseManifest)1 Ddl (com.google.cloud.teleport.spanner.ddl.Ddl)1 ImmutableList (com.google.common.collect.ImmutableList)1 ImmutableMap (com.google.common.collect.ImmutableMap)1 InvalidProtocolBufferException (com.google.protobuf.InvalidProtocolBufferException)1 JsonFormat (com.google.protobuf.util.JsonFormat)1 Files (java.nio.file.Files)1 Path (java.nio.file.Path)1 List (java.util.List)1 Map (java.util.Map)1 PAssert (org.apache.beam.sdk.testing.PAssert)1