Examples with Struct org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.Struct used on opensource projects

Example 1 with Struct

use of org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.Struct in project beam by apache.

the class DefaultJobBundleFactoryTest method loadBalancesBundles.

@Test
public void loadBalancesBundles() throws Exception {
    PortablePipelineOptions portableOptions = PipelineOptionsFactory.as(PortablePipelineOptions.class);
    portableOptions.setSdkWorkerParallelism(2);
    portableOptions.setLoadBalanceBundles(true);
    Struct pipelineOptions = PipelineOptionsTranslation.toProto(portableOptions);
    try (DefaultJobBundleFactory bundleFactory = new DefaultJobBundleFactory(JobInfo.create("testJob", "testJob", "token", pipelineOptions), envFactoryProviderMap, stageIdGenerator, serverInfo)) {
        OutputReceiverFactory orf = mock(OutputReceiverFactory.class);
        StateRequestHandler srh = mock(StateRequestHandler.class);
        when(srh.getCacheTokens()).thenReturn(Collections.emptyList());
        StageBundleFactory sbf = bundleFactory.forStage(getExecutableStage(environment));
        RemoteBundle b1 = sbf.getBundle(orf, srh, BundleProgressHandler.ignored());
        verify(envFactory, Mockito.times(1)).createEnvironment(eq(environment), any());
        final RemoteBundle b2 = sbf.getBundle(orf, srh, BundleProgressHandler.ignored());
        verify(envFactory, Mockito.times(2)).createEnvironment(eq(environment), any());
        AtomicBoolean b2Closing = new AtomicBoolean(false);
        ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
        ScheduledFuture<Optional<Exception>> closingFuture = executor.schedule(() -> {
            try {
                b2Closing.compareAndSet(false, true);
                b2.close();
                return Optional.empty();
            } catch (Exception e) {
                return Optional.of(e);
            }
        }, 100, TimeUnit.MILLISECONDS);
        assertThat(b2Closing.get(), equalTo(false));
        // This call should block until closingFuture has finished closing b2 (100ms)
        RemoteBundle b3 = sbf.getBundle(orf, srh, BundleProgressHandler.ignored());
        // ensure the previous call waited for close
        assertThat(b2Closing.get(), equalTo(true));
        // Join closingFuture and check if an exception occurred
        Optional<Exception> closingException = closingFuture.get();
        if (closingException.isPresent()) {
            throw new AssertionError("Exception occurred while closing b2", closingException.get());
        }
        verify(envFactory, Mockito.times(2)).createEnvironment(eq(environment), any());
        b3.close();
        b1.close();
        executor.shutdown();
    }
}

Example 2 with Struct

use of org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.Struct in project beam by apache.

the class PipelineOptionsTranslation method toProto.

/**
 * Converts the provided {@link PipelineOptions} to a {@link Struct}.
 */
public static Struct toProto(PipelineOptions options) {
    Struct.Builder builder = Struct.newBuilder();
    try {
        // TODO: Officially define URNs for options and their scheme.
        JsonNode treeNode = MAPPER.valueToTree(options);
        JsonNode rootOptions = treeNode.get("options");
        Iterator<Map.Entry<String, JsonNode>> optionsEntries = rootOptions.fields();
        if (!optionsEntries.hasNext()) {
            // (observed with version 2.2.3)
            throw new RuntimeException("Unable to convert pipeline options, please check for outdated jackson-core version in the classpath.");
        }
        Map<String, TreeNode> optionsUsingUrns = new HashMap<>();
        while (optionsEntries.hasNext()) {
            Map.Entry<String, JsonNode> entry = optionsEntries.next();
            optionsUsingUrns.put("beam:option:" + CaseFormat.LOWER_CAMEL.to(CaseFormat.LOWER_UNDERSCORE, entry.getKey()) + ":v1", entry.getValue());
        }
        // The JSON format of a Protobuf Struct is the JSON object that is equivalent to that struct
        // (with values encoded in a standard json-codeable manner). See Beam PR 3719 for more.
        JsonFormat.parser().merge(MAPPER.writeValueAsString(optionsUsingUrns), builder);
        return builder.build();
    } catch (IOException e) {
        throw new RuntimeException("Failed to convert PipelineOptions to Protocol", e);
    }
}

Example 3 with Struct

use of org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.Struct in project beam by apache.

the class StaticGrpcProvisionServiceTest method returnsProvisionInfo.

@Test
public void returnsProvisionInfo() throws Exception {
    Struct options = Struct.newBuilder().putFields("foo", Value.newBuilder().setBoolValue(true).build()).putFields("bar", Value.newBuilder().setNumberValue(2.5).build()).putFields("baz", Value.newBuilder().setListValue(ListValue.newBuilder().addValues(Value.newBuilder().setStructValue(Struct.newBuilder().putFields("spam", Value.newBuilder().setNullValue(NullValue.NULL_VALUE).build()))).build()).build()).build();
    ProvisionInfo info = ProvisionInfo.newBuilder().setPipelineOptions(options).build();
    GrpcFnServer<StaticGrpcProvisionService> server = GrpcFnServer.allocatePortAndCreateFor(StaticGrpcProvisionService.create(info, GrpcContextHeaderAccessorProvider.getHeaderAccessor()), InProcessServerFactory.create());
    ProvisionServiceBlockingStub stub = ProvisionServiceGrpc.newBlockingStub(InProcessChannelBuilder.forName(server.getApiServiceDescriptor().getUrl()).intercept(AddHarnessIdInterceptor.create("test_worker")).build());
    GetProvisionInfoResponse provisionResponse = stub.getProvisionInfo(GetProvisionInfoRequest.getDefaultInstance());
    assertThat(provisionResponse.getInfo(), equalTo(info));
}

Example 4 with Struct

use of org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.Struct in project beam by apache.

the class InMemoryJobService method prepare.

@Override
public void prepare(PrepareJobRequest request, StreamObserver<PrepareJobResponse> responseObserver) {
    try {
        LOG.trace("{} {}", PrepareJobRequest.class.getSimpleName(), request);
        // insert preparation
        String preparationId = String.format("%s_%s", request.getJobName(), UUID.randomUUID().toString());
        Struct pipelineOptions = request.getPipelineOptions();
        if (pipelineOptions == null) {
            throw new NullPointerException("Encountered null pipeline options.");
        }
        LOG.trace("PIPELINE OPTIONS {} {}", pipelineOptions.getClass(), pipelineOptions);
        JobPreparation preparation = JobPreparation.builder().setId(preparationId).setPipeline(request.getPipeline()).setOptions(pipelineOptions).build();
        JobPreparation previous = preparations.putIfAbsent(preparationId, preparation);
        if (previous != null) {
            // this should never happen with a UUID
            String errMessage = String.format("A job with the preparation ID \"%s\" already exists.", preparationId);
            StatusException exception = Status.NOT_FOUND.withDescription(errMessage).asException();
            responseObserver.onError(exception);
            return;
        }
        String stagingSessionToken = stagingServiceTokenProvider.apply(preparationId);
        stagingSessionTokens.putIfAbsent(preparationId, stagingSessionToken);
        stagingService.getService().registerJob(stagingSessionToken, Maps.transformValues(request.getPipeline().getComponents().getEnvironmentsMap(), RunnerApi.Environment::getDependenciesList));
        // send response
        PrepareJobResponse response = PrepareJobResponse.newBuilder().setPreparationId(preparationId).setArtifactStagingEndpoint(stagingServiceDescriptor).setStagingSessionToken(stagingSessionToken).build();
        responseObserver.onNext(response);
        responseObserver.onCompleted();
    } catch (Exception e) {
        LOG.error("Could not prepare job with name {}", request.getJobName(), e);
        responseObserver.onError(Status.INTERNAL.withCause(e).asException());
    }
}

Example 5 with Struct

use of org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.Struct in project grpc-java by grpc.

the class ClientXdsClientDataTest method parseHttpFilter_typedStructMigration.

@Test
public void parseHttpFilter_typedStructMigration() {
    filterRegistry.register(new TestFilter());
    Struct rawStruct = Struct.newBuilder().putFields("name", Value.newBuilder().setStringValue("default").build()).build();
    HttpFilter httpFilter = HttpFilter.newBuilder().setIsOptional(true).setTypedConfig(Any.pack(com.github.udpa.udpa.type.v1.TypedStruct.newBuilder().setTypeUrl("test-url").setValue(rawStruct).build())).build();
    FilterConfig config = ClientXdsClient.parseHttpFilter(httpFilter, filterRegistry, true).getStruct();
    assertThat(((SimpleFilterConfig) config).getConfig()).isEqualTo(rawStruct);
    HttpFilter httpFilterNewTypeStruct = HttpFilter.newBuilder().setIsOptional(true).setTypedConfig(Any.pack(TypedStruct.newBuilder().setTypeUrl("test-url").setValue(rawStruct).build())).build();
    config = ClientXdsClient.parseHttpFilter(httpFilterNewTypeStruct, filterRegistry, true).getStruct();
    assertThat(((SimpleFilterConfig) config).getConfig()).isEqualTo(rawStruct);
}