Examples with AllowList org.apache.beam.sdk.expansion.service.JavaClassLookupTransformProvider.AllowList used on opensource projects

Example 1 with AllowList

use of org.apache.beam.sdk.expansion.service.JavaClassLookupTransformProvider.AllowList in project beam by apache.

the class ExpansionService method expand.

@VisibleForTesting
/*package*/
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 = createPipeline();
    boolean isUseDeprecatedRead = ExperimentalOptions.hasExperiment(pipelineOptions, "use_deprecated_read") || ExperimentalOptions.hasExperiment(pipelineOptions, "beam_fn_api_use_deprecated_read");
    if (!isUseDeprecatedRead) {
        ExperimentalOptions.addExperiment(pipeline.getOptions().as(ExperimentalOptions.class), "beam_fn_api");
        // TODO(BEAM-10670): Remove this when we address performance issue.
        ExperimentalOptions.addExperiment(pipeline.getOptions().as(ExperimentalOptions.class), "use_sdf_read");
    } else {
        LOG.warn("Using use_depreacted_read in portable runners is runner-dependent. The " + "ExpansionService will respect that, but if your runner does not have support for " + "native Read transform, your Pipeline will fail during Pipeline submission.");
    }
    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);
        }
    }));
    String urn = request.getTransform().getSpec().getUrn();
    TransformProvider transformProvider = null;
    if (getUrn(ExpansionMethods.Enum.JAVA_CLASS_LOOKUP).equals(urn)) {
        AllowList allowList = pipelineOptions.as(ExpansionServiceOptions.class).getJavaClassLookupAllowlist();
        assert allowList != null;
        transformProvider = new JavaClassLookupTransformProvider(allowList);
    } else {
        transformProvider = getRegisteredTransforms().get(urn);
        if (transformProvider == null) {
            throw new UnsupportedOperationException("Unknown urn: " + request.getTransform().getSpec().getUrn());
        }
    }
    List<String> classpathResources = transformProvider.getDependencies(request.getTransform().getSpec(), pipeline.getOptions());
    pipeline.getOptions().as(PortablePipelineOptions.class).setFilesToStage(classpathResources);
    Map<String, PCollection<?>> outputs = transformProvider.apply(pipeline, request.getTransform().getUniqueName(), request.getTransform().getSpec(), inputs);
    // Needed to find which transform was new...
    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);
        }
    }));
    if (isUseDeprecatedRead) {
        SplittableParDo.convertReadBasedSplittableDoFnsToPrimitiveReadsIfNecessary(pipeline);
    }
    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();
}

Example 2 with AllowList

use of org.apache.beam.sdk.expansion.service.JavaClassLookupTransformProvider.AllowList in project beam by apache.

the class JavaClassLookupTransformProviderTest method testEverythingAllowList.

@Test
public void testEverythingAllowList() {
    AllowList everything = AllowList.everything();
    AllowedClass allowedClass = everything.getAllowedClass("org.apache.beam.sdk.transforms.KvSwap");
    assertTrue(allowedClass.isAllowedBuilderMethod("builder"));
    assertTrue(allowedClass.isAllowedConstructorMethod("constructor"));
}

Example 3 with AllowList

use of org.apache.beam.sdk.expansion.service.JavaClassLookupTransformProvider.AllowList in project beam by apache.

the class JavaClassLookupTransformProviderTest method testNothingAllowList.

@Test
public void testNothingAllowList() {
    AllowList nothing = AllowList.nothing();
    RuntimeException thrown = assertThrows(RuntimeException.class, () -> nothing.getAllowedClass("org.apache.beam.sdk.transforms.KvSwap"));
    assertTrue(thrown.getMessage(), thrown.getMessage().contains("allow list does not enable"));
    assertTrue(thrown.getMessage(), thrown.getMessage().contains("org.apache.beam.sdk.transforms.KvSwap"));
}

Example 4 with AllowList

use of org.apache.beam.sdk.expansion.service.JavaClassLookupTransformProvider.AllowList in project beam by apache.

the class JavaClassLookupTransformProviderTest method testPackageAllowList.

@Test
public void testPackageAllowList() {
    AllowList allowList = AllowList.create(ALLOW_LIST_VERSION, Collections.singletonList(AllowedClass.create("good.package.*", Collections.singletonList("goodBuilder"), AllowedClass.WILDCARD)));
    assertThrows(RuntimeException.class, () -> allowList.getAllowedClass("bad.package.Transform"));
    AllowedClass allowedClass = allowList.getAllowedClass("good.package.Transform");
    assertTrue(allowedClass.isAllowedBuilderMethod("goodBuilder"));
    assertFalse(allowedClass.isAllowedBuilderMethod("badBuilder"));
    assertTrue(allowedClass.isAllowedConstructorMethod("anyConstructor"));
}