QueryablePipeline { public PTransformNode getProducer(PCollectionNode pcollection) { return (PTransformNode) Iterables.getOnlyElement(pipelineNetwork.predecessors(pcollection)); } private  QueryablePipeline(Collection<String> transformIds, Components components); static QueryablePipeline forPrimitivesIn(Components components); static QueryablePipeline forPipeline(RunnerApi.Pipeline p); static QueryablePipeline forTransforms(
      Collection<String> transformIds, Components components); Collection<PTransformNode> getTransforms(); Iterable<PTransformNode> getTopologicallyOrderedTransforms(); Set<PTransformNode> getRootTransforms(); PTransformNode getProducer(PCollectionNode pcollection); Set<PTransformNode> getPerElementConsumers(PCollectionNode pCollection); Set<PTransformNode> getSingletonConsumers(PCollectionNode pCollection); Set<PCollectionNode> getPerElementInputPCollections(PTransformNode ptransform); Set<PCollectionNode> getOutputPCollections(PTransformNode ptransform); Components getComponents(); Collection<SideInputReference> getSideInputs(PTransformNode transform); Collection<UserStateReference> getUserStates(PTransformNode transform); Collection<TimerReference> getTimers(PTransformNode transform); Optional<Environment> getEnvironment(PTransformNode parDo);  }
@Test public void getProducer() { Pipeline p = Pipeline.create(); PCollection<Long> longs = p.apply("BoundedRead", Read.from(CountingSource.upTo(100L))); PCollectionList.of(longs).and(longs).and(longs).apply("flatten", Flatten.pCollections()); Components components = PipelineTranslation.toProto(p).getComponents(); QueryablePipeline qp = QueryablePipeline.forPrimitivesIn(components); String impulseOutputName = getOnlyElement( PipelineNode.pTransform( "BoundedRead/Impulse", components.getTransformsOrThrow("BoundedRead/Impulse")) .getTransform() .getOutputsMap() .values()); PTransformNode impulseProducer = PipelineNode.pTransform( "BoundedRead/Impulse", components.getTransformsOrThrow("BoundedRead/Impulse")); PCollectionNode impulseOutput = PipelineNode.pCollection( impulseOutputName, components.getPcollectionsOrThrow(impulseOutputName)); String flattenOutputName = getOnlyElement( PipelineNode.pTransform("flatten", components.getTransformsOrThrow("flatten")) .getTransform() .getOutputsMap() .values()); PTransformNode flattenProducer = PipelineNode.pTransform("flatten", components.getTransformsOrThrow("flatten")); PCollectionNode flattenOutput = PipelineNode.pCollection( flattenOutputName, components.getPcollectionsOrThrow(flattenOutputName)); assertThat(qp.getProducer(impulseOutput), equalTo(impulseProducer)); assertThat(qp.getProducer(flattenOutput), equalTo(flattenProducer)); }
@Test public void getProducer() { Pipeline p = Pipeline.create(); PCollection<Long> longs = p.apply("BoundedRead", Read.from(CountingSource.upTo(100L))); PCollectionList.of(longs).and(longs).and(longs).apply("flatten", Flatten.pCollections()); Components components = PipelineTranslation.toProto(p).getComponents(); QueryablePipeline qp = QueryablePipeline.forPrimitivesIn(components); String longsOutputName = getOnlyElement( PipelineNode.pTransform("BoundedRead", components.getTransformsOrThrow("BoundedRead")) .getTransform() .getOutputsMap() .values()); PTransformNode longsProducer = PipelineNode.pTransform("BoundedRead", components.getTransformsOrThrow("BoundedRead")); PCollectionNode longsOutput = PipelineNode.pCollection( longsOutputName, components.getPcollectionsOrThrow(longsOutputName)); String flattenOutputName = getOnlyElement( PipelineNode.pTransform("flatten", components.getTransformsOrThrow("flatten")) .getTransform() .getOutputsMap() .values()); PTransformNode flattenProducer = PipelineNode.pTransform("flatten", components.getTransformsOrThrow("flatten")); PCollectionNode flattenOutput = PipelineNode.pCollection( flattenOutputName, components.getPcollectionsOrThrow(flattenOutputName)); assertThat(qp.getProducer(longsOutput), equalTo(longsProducer)); assertThat(qp.getProducer(flattenOutput), equalTo(flattenProducer)); }