ProtoOverrides { public static Pipeline updateTransform( String urn, Pipeline originalPipeline, TransformReplacement compositeBuilder) { Components.Builder resultComponents = originalPipeline.getComponents().toBuilder(); for (Map.Entry<String, PTransform> pt : originalPipeline.getComponents().getTransformsMap().entrySet()) { if (pt.getValue().getSpec() != null && urn.equals(pt.getValue().getSpec().getUrn())) { MessageWithComponents updated = compositeBuilder.getReplacement(pt.getKey(), originalPipeline.getComponents()); if (updated == null) { continue; } checkArgument( updated.getPtransform().getOutputsMap().equals(pt.getValue().getOutputsMap()), "A %s must produce all of the outputs of the original %s", TransformReplacement.class.getSimpleName(), PTransform.class.getSimpleName()); removeSubtransforms(pt.getValue(), resultComponents); resultComponents .mergeFrom(updated.getComponents()) .putTransforms(pt.getKey(), updated.getPtransform()); } } return originalPipeline.toBuilder().setComponents(resultComponents).build(); }  static Pipeline updateTransform(
      String urn, Pipeline originalPipeline, TransformReplacement compositeBuilder);  }
@Test public void replacesOnlyMatching() { RunnerApi.Pipeline p = Pipeline.newBuilder() .addAllRootTransformIds(ImmutableList.of("first", "second")) .setComponents( Components.newBuilder() .putTransforms( "first", PTransform.newBuilder() .setSpec(FunctionSpec.newBuilder().setUrn("beam:first")) .build()) .putTransforms( "second", PTransform.newBuilder() .setSpec(FunctionSpec.newBuilder().setUrn("beam:second")) .build()) .putPcollections( "intermediatePc", PCollection.newBuilder().setUniqueName("intermediate").build()) .putCoders( "coder", Coder.newBuilder().setSpec(FunctionSpec.getDefaultInstance()).build())) .build(); PTransform secondReplacement = PTransform.newBuilder() .addSubtransforms("second_sub") .setSpec( FunctionSpec.newBuilder() .setUrn("beam:second:replacement") .setPayload( ByteString.copyFrom("foo-bar-baz".getBytes(StandardCharsets.UTF_8)))) .build(); WindowingStrategy introducedWS = WindowingStrategy.newBuilder() .setAccumulationMode(AccumulationMode.Enum.ACCUMULATING) .build(); RunnerApi.Components extraComponents = Components.newBuilder() .putPcollections( "intermediatePc", PCollection.newBuilder().setUniqueName("intermediate_replacement").build()) .putWindowingStrategies("new_ws", introducedWS) .putTransforms("second_sub", PTransform.getDefaultInstance()) .build(); Pipeline updated = ProtoOverrides.updateTransform( "beam:second", p, new TestReplacer(secondReplacement, extraComponents)); PTransform updatedSecond = updated.getComponents().getTransformsOrThrow("second"); assertThat(updatedSecond, equalTo(secondReplacement)); assertThat( updated.getComponents().getWindowingStrategiesOrThrow("new_ws"), equalTo(introducedWS)); assertThat( updated.getComponents().getTransformsOrThrow("second_sub"), equalTo(PTransform.getDefaultInstance())); assertThat( updated.getComponents().getPcollectionsOrThrow("intermediatePc").getUniqueName(), equalTo("intermediate_replacement")); assertThat( updated.getComponents().getTransformsOrThrow("first"), equalTo(p.getComponents().getTransformsOrThrow("first"))); assertThat( updated.getComponents().getCodersOrThrow("coder"), equalTo(p.getComponents().getCodersOrThrow("coder"))); assertThat(updated.getRootTransformIdsList(), equalTo(p.getRootTransformIdsList())); }
@Test public void replacesMultiple() { RunnerApi.Pipeline p = Pipeline.newBuilder() .addAllRootTransformIds(ImmutableList.of("first", "second")) .setComponents( Components.newBuilder() .putTransforms( "first", PTransform.newBuilder() .setSpec(FunctionSpec.newBuilder().setUrn("beam:first")) .build()) .putTransforms( "second", PTransform.newBuilder() .setSpec(FunctionSpec.newBuilder().setUrn("beam:repeated")) .build()) .putTransforms( "third", PTransform.newBuilder() .setSpec(FunctionSpec.newBuilder().setUrn("beam:repeated")) .build()) .putPcollections( "intermediatePc", PCollection.newBuilder().setUniqueName("intermediate").build()) .putCoders( "coder", Coder.newBuilder().setSpec(FunctionSpec.getDefaultInstance()).build())) .build(); ByteString newPayload = ByteString.copyFrom("foo-bar-baz".getBytes(StandardCharsets.UTF_8)); Pipeline updated = ProtoOverrides.updateTransform( "beam:repeated", p, (transformId, existingComponents) -> { String subtransform = String.format("%s_sub", transformId); return MessageWithComponents.newBuilder() .setPtransform( PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn("beam:repeated:replacement") .setPayload(newPayload)) .addSubtransforms(subtransform)) .setComponents( Components.newBuilder() .putTransforms( subtransform, PTransform.newBuilder().setUniqueName(subtransform).build())) .build(); }); PTransform updatedSecond = updated.getComponents().getTransformsOrThrow("second"); PTransform updatedThird = updated.getComponents().getTransformsOrThrow("third"); assertThat(updatedSecond, not(equalTo(p.getComponents().getTransformsOrThrow("second")))); assertThat(updatedThird, not(equalTo(p.getComponents().getTransformsOrThrow("third")))); assertThat(updatedSecond.getSubtransformsList(), contains("second_sub")); assertThat(updatedSecond.getSpec().getPayload(), equalTo(newPayload)); assertThat(updatedThird.getSubtransformsList(), contains("third_sub")); assertThat(updatedThird.getSpec().getPayload(), equalTo(newPayload)); assertThat(updated.getComponents().getTransformsMap(), hasKey("second_sub")); assertThat(updated.getComponents().getTransformsMap(), hasKey("third_sub")); assertThat( updated.getComponents().getTransformsOrThrow("second_sub").getUniqueName(), equalTo("second_sub")); assertThat( updated.getComponents().getTransformsOrThrow("third_sub").getUniqueName(), equalTo("third_sub")); }
@Test public void replaceExistingCompositeSucceeds() { Pipeline p = Pipeline.newBuilder() .addRootTransformIds("root") .setComponents( Components.newBuilder() .putTransforms( "root", PTransform.newBuilder() .addSubtransforms("sub_first") .setSpec(FunctionSpec.newBuilder().setUrn("beam:composite")) .build()) .putTransforms( "sub_first", PTransform.newBuilder() .setSpec(FunctionSpec.newBuilder().setUrn("beam:inner")) .build())) .build(); Pipeline pipeline = ProtoOverrides.updateTransform( "beam:composite", p, new TestReplacer( PTransform.newBuilder() .addSubtransforms("foo") .addSubtransforms("bar") .setSpec( FunctionSpec.getDefaultInstance() .newBuilderForType() .setUrn("beam:composite")) .build(), Components.getDefaultInstance())); assertThat( pipeline.getComponents().getTransformsOrThrow("root").getSpec().getUrn(), equalTo("beam:composite")); assertThat( pipeline.getComponents().getTransformsOrThrow("root").getSubtransformsList(), contains("foo", "bar")); }
@Test public void replacesOnlyMatching() { RunnerApi.Pipeline p = Pipeline.newBuilder() .addAllRootTransformIds(ImmutableList.of("first", "second")) .setComponents( Components.newBuilder() .putTransforms( "first", PTransform.newBuilder() .setSpec(FunctionSpec.newBuilder().setUrn("beam:first")) .build()) .putTransforms( "second", PTransform.newBuilder() .setSpec(FunctionSpec.newBuilder().setUrn("beam:second")) .build()) .putPcollections( "intermediatePc", PCollection.newBuilder().setUniqueName("intermediate").build()) .putCoders( "coder", Coder.newBuilder().setSpec(FunctionSpec.getDefaultInstance()).build())) .build(); PTransform secondReplacement = PTransform.newBuilder() .addSubtransforms("second_sub") .setSpec( FunctionSpec.newBuilder() .setUrn("beam:second:replacement") .setPayload( ByteString.copyFrom("foo-bar-baz".getBytes(StandardCharsets.UTF_8)))) .build(); WindowingStrategy introducedWS = WindowingStrategy.newBuilder().setAccumulationMode(Enum.ACCUMULATING).build(); RunnerApi.Components extraComponents = Components.newBuilder() .putPcollections( "intermediatePc", PCollection.newBuilder().setUniqueName("intermediate_replacement").build()) .putWindowingStrategies("new_ws", introducedWS) .putTransforms("second_sub", PTransform.getDefaultInstance()) .build(); Pipeline updated = ProtoOverrides.updateTransform( "beam:second", p, new TestReplacer(secondReplacement, extraComponents)); PTransform updatedSecond = updated.getComponents().getTransformsOrThrow("second"); assertThat(updatedSecond, equalTo(secondReplacement)); assertThat( updated.getComponents().getWindowingStrategiesOrThrow("new_ws"), equalTo(introducedWS)); assertThat( updated.getComponents().getTransformsOrThrow("second_sub"), equalTo(PTransform.getDefaultInstance())); assertThat( updated.getComponents().getPcollectionsOrThrow("intermediatePc").getUniqueName(), equalTo("intermediate_replacement")); assertThat( updated.getComponents().getTransformsOrThrow("first"), equalTo(p.getComponents().getTransformsOrThrow("first"))); assertThat( updated.getComponents().getCodersOrThrow("coder"), equalTo(p.getComponents().getCodersOrThrow("coder"))); assertThat(updated.getRootTransformIdsList(), equalTo(p.getRootTransformIdsList())); }