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());
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();
}