ExpansionService extends ExpansionServiceGrpc.ExpansionServiceImplBase implements AutoCloseable { @VisibleForTesting ExpansionApi.ExpansionResponse expand(ExpansionApi.ExpansionRequest request) { LOG.info( "Expanding '{}' with URN '{}'", request.getTransform().getUniqueName(), request.getTransform().getSpec().getUrn()); LOG.debug("Full transform: {}", request.getTransform()); Set<String> existingTransformIds = request.getComponents().getTransformsMap().keySet(); Pipeline pipeline = Pipeline.create(); ExperimentalOptions.addExperiment( pipeline.getOptions().as(ExperimentalOptions.class), "beam_fn_api"); ClassLoader classLoader = Environments.class.getClassLoader(); if (classLoader == null) { throw new RuntimeException( "Cannot detect classpath: classload is null (is it the bootstrap classloader?)"); } List<String> classpathResources = detectClassPathResourcesToStage(classLoader, pipeline.getOptions()); if (classpathResources.isEmpty()) { throw new IllegalArgumentException("No classpath elements found."); } LOG.debug("Staging to files from the classpath: {}", classpathResources.size()); pipeline.getOptions().as(PortablePipelineOptions.class).setFilesToStage(classpathResources); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(request.getComponents()).withPipeline(pipeline); Map<String, PCollection<?>> inputs = request.getTransform().getInputsMap().entrySet().stream() .collect( Collectors.toMap( Map.Entry::getKey, input -> { try { return rehydratedComponents.getPCollection(input.getValue()); } catch (IOException exn) { throw new RuntimeException(exn); } })); @Nullable TransformProvider transformProvider = getRegisteredTransforms().get(request.getTransform().getSpec().getUrn()); if (transformProvider == null) { throw new UnsupportedOperationException( "Unknown urn: " + request.getTransform().getSpec().getUrn()); } Map<String, PCollection<?>> outputs = transformProvider.apply( pipeline, request.getTransform().getUniqueName(), request.getTransform().getSpec(), inputs); SdkComponents sdkComponents = rehydratedComponents .getSdkComponents(Collections.emptyList()) .withNewIdPrefix(request.getNamespace()); sdkComponents.registerEnvironment( Environments.createOrGetDefaultEnvironment( pipeline.getOptions().as(PortablePipelineOptions.class))); Map<String, String> outputMap = outputs.entrySet().stream() .collect( Collectors.toMap( Map.Entry::getKey, output -> { try { return sdkComponents.registerPCollection(output.getValue()); } catch (IOException exn) { throw new RuntimeException(exn); } })); RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline, sdkComponents); String expandedTransformId = Iterables.getOnlyElement( pipelineProto.getRootTransformIdsList().stream() .filter(id -> !existingTransformIds.contains(id)) .collect(Collectors.toList())); RunnerApi.Components components = pipelineProto.getComponents(); RunnerApi.PTransform expandedTransform = components .getTransformsOrThrow(expandedTransformId) .toBuilder() .setUniqueName(expandedTransformId) .clearOutputs() .putAllOutputs(outputMap) .build(); LOG.debug("Expanded to {}", expandedTransform); return ExpansionApi.ExpansionResponse.newBuilder() .setComponents(components.toBuilder().removeTransforms(expandedTransformId)) .setTransform(expandedTransform) .addAllRequirements(pipelineProto.getRequirementsList()) .build(); }  @Override void expand(
      ExpansionApi.ExpansionRequest request,
      StreamObserver<ExpansionApi.ExpansionResponse> responseObserver); @Override void close(); static void main(String[] args);  }
@Test public void testConstruct() { Pipeline p = Pipeline.create(); p.apply(Impulse.create()); RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p); String inputPcollId = Iterables.getOnlyElement( Iterables.getOnlyElement(pipelineProto.getComponents().getTransformsMap().values()) .getOutputsMap() .values()); ExpansionApi.ExpansionRequest request = ExpansionApi.ExpansionRequest.newBuilder() .setComponents(pipelineProto.getComponents()) .setTransform( RunnerApi.PTransform.newBuilder() .setUniqueName(TEST_NAME) .setSpec(RunnerApi.FunctionSpec.newBuilder().setUrn(TEST_URN)) .putInputs("input", inputPcollId)) .setNamespace(TEST_NAMESPACE) .build(); ExpansionApi.ExpansionResponse response = expansionService.expand(request); RunnerApi.PTransform expandedTransform = response.getTransform(); assertEquals(TEST_NAMESPACE + TEST_NAME, expandedTransform.getUniqueName()); assertThat(expandedTransform.getInputsMap().values(), contains(inputPcollId)); assertThat(expandedTransform.getOutputsMap().keySet(), contains("output")); assertThat(expandedTransform.getSubtransformsCount(), greaterThan(0)); for (String subtransform : expandedTransform.getSubtransformsList()) { assertTrue(response.getComponents().containsTransforms(subtransform)); } Set<String> originalIds = allIds(request.getComponents()); for (String id : allIds(response.getComponents())) { assertTrue(id, id.startsWith(TEST_NAMESPACE) || originalIds.contains(id)); } }
@Test public void testConstructGenerateSequence() { ExternalTransforms.ExternalConfigurationPayload payload = encodeRow( Row.withSchema( Schema.of( Field.of("start", FieldType.INT64), Field.nullable("stop", FieldType.INT64))) .withFieldValue("start", 0L) .withFieldValue("stop", 1L) .build()); Pipeline p = Pipeline.create(); RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p); ExpansionApi.ExpansionRequest request = ExpansionApi.ExpansionRequest.newBuilder() .setComponents(pipelineProto.getComponents()) .setTransform( RunnerApi.PTransform.newBuilder() .setUniqueName(TEST_NAME) .setSpec( RunnerApi.FunctionSpec.newBuilder() .setUrn(GenerateSequence.External.URN) .setPayload(payload.toByteString()))) .setNamespace(TEST_NAMESPACE) .build(); ExpansionApi.ExpansionResponse response = expansionService.expand(request); RunnerApi.PTransform expandedTransform = response.getTransform(); assertEquals(TEST_NAMESPACE + TEST_NAME, expandedTransform.getUniqueName()); assertThat(expandedTransform.getInputsCount(), Matchers.is(0)); assertThat(expandedTransform.getOutputsCount(), Matchers.is(1)); assertThat(expandedTransform.getSubtransformsCount(), greaterThan(0)); }