Examples with Pipeline org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline used on opensource projects

Example 96 with Pipeline

use of org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline in project beam by apache.

the class DataflowPipelineTranslatorTest method testNetworkConfig.

@Test
public void testNetworkConfig() throws IOException {
    final String testNetwork = "test-network";
    DataflowPipelineOptions options = buildPipelineOptions();
    options.setNetwork(testNetwork);
    Pipeline p = buildPipeline(options);
    p.traverseTopologically(new RecordingPipelineVisitor());
    SdkComponents sdkComponents = createSdkComponents(options);
    RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p, sdkComponents, true);
    Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, pipelineProto, sdkComponents, DataflowRunner.fromOptions(options), Collections.emptyList()).getJob();
    assertEquals(1, job.getEnvironment().getWorkerPools().size());
    assertEquals(testNetwork, job.getEnvironment().getWorkerPools().get(0).getNetwork());
}

Example 97 with Pipeline

use of org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline in project beam by apache.

the class DataflowRunnerTest method testStageArtifactWithoutStagedName.

@Test
public void testStageArtifactWithoutStagedName() throws IOException {
    DataflowPipelineOptions options = buildPipelineOptions();
    DataflowRunner runner = DataflowRunner.fromOptions(options);
    File temp1 = File.createTempFile("artifact1-", ".txt");
    temp1.deleteOnExit();
    File temp2 = File.createTempFile("artifact2-", ".txt");
    temp2.deleteOnExit();
    RunnerApi.ArtifactInformation fooLocalArtifact = RunnerApi.ArtifactInformation.newBuilder().setTypeUrn(BeamUrns.getUrn(RunnerApi.StandardArtifacts.Types.FILE)).setTypePayload(RunnerApi.ArtifactFilePayload.newBuilder().setPath(temp1.getAbsolutePath()).build().toByteString()).build();
    RunnerApi.ArtifactInformation barLocalArtifact = RunnerApi.ArtifactInformation.newBuilder().setTypeUrn(BeamUrns.getUrn(RunnerApi.StandardArtifacts.Types.FILE)).setTypePayload(RunnerApi.ArtifactFilePayload.newBuilder().setPath(temp2.getAbsolutePath()).build().toByteString()).build();
    RunnerApi.Pipeline pipeline = RunnerApi.Pipeline.newBuilder().setComponents(RunnerApi.Components.newBuilder().putEnvironments("env", RunnerApi.Environment.newBuilder().addAllDependencies(ImmutableList.of(fooLocalArtifact, barLocalArtifact)).build())).build();
    List<DataflowPackage> packages = runner.stageArtifacts(pipeline);
    for (DataflowPackage pkg : packages) {
        assertThat(pkg.getName(), matchesRegex("artifact[1,2]-.+\\.txt"));
    }
}

Example 98 with Pipeline

use of org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline in project beam by apache.

the class DataflowRunnerTest method testTransformTranslatorMissing.

@Test
public void testTransformTranslatorMissing() throws IOException {
    DataflowPipelineOptions options = buildPipelineOptions();
    Pipeline p = Pipeline.create(options);
    p.apply(Create.of(Arrays.asList(1, 2, 3))).apply(new TestTransform());
    thrown.expect(IllegalStateException.class);
    thrown.expectMessage(containsString("no translator registered"));
    SdkComponents sdkComponents = SdkComponents.create(options);
    RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p, sdkComponents, true);
    DataflowPipelineTranslator.fromOptions(options).translate(p, pipelineProto, sdkComponents, DataflowRunner.fromOptions(options), Collections.emptyList());
    ArgumentCaptor<Job> jobCaptor = ArgumentCaptor.forClass(Job.class);
    Mockito.verify(mockJobs).create(eq(PROJECT_ID), eq(REGION_ID), jobCaptor.capture());
    assertValidJob(jobCaptor.getValue());
}

Example 99 with Pipeline

use of org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline in project beam by apache.

the class DataflowRunnerTest method testResolveArtifacts.

@Test
public void testResolveArtifacts() throws IOException {
    DataflowPipelineOptions options = buildPipelineOptions();
    DataflowRunner runner = DataflowRunner.fromOptions(options);
    String stagingLocation = options.getStagingLocation().replaceFirst("/$", "");
    RunnerApi.ArtifactInformation fooLocalArtifact = RunnerApi.ArtifactInformation.newBuilder().setTypeUrn(BeamUrns.getUrn(RunnerApi.StandardArtifacts.Types.FILE)).setTypePayload(RunnerApi.ArtifactFilePayload.newBuilder().setPath("/tmp/foo.jar").build().toByteString()).setRoleUrn(BeamUrns.getUrn(RunnerApi.StandardArtifacts.Roles.STAGING_TO)).setRolePayload(RunnerApi.ArtifactStagingToRolePayload.newBuilder().setStagedName("foo_staged.jar").build().toByteString()).build();
    RunnerApi.ArtifactInformation barLocalArtifact = RunnerApi.ArtifactInformation.newBuilder().setTypeUrn(BeamUrns.getUrn(RunnerApi.StandardArtifacts.Types.FILE)).setTypePayload(RunnerApi.ArtifactFilePayload.newBuilder().setPath("/tmp/bar.jar").build().toByteString()).setRoleUrn(BeamUrns.getUrn(RunnerApi.StandardArtifacts.Roles.STAGING_TO)).setRolePayload(RunnerApi.ArtifactStagingToRolePayload.newBuilder().setStagedName("bar_staged.jar").build().toByteString()).build();
    RunnerApi.Pipeline pipeline = RunnerApi.Pipeline.newBuilder().setComponents(RunnerApi.Components.newBuilder().putEnvironments("env", RunnerApi.Environment.newBuilder().addAllDependencies(ImmutableList.of(fooLocalArtifact, barLocalArtifact)).build())).build();
    RunnerApi.ArtifactInformation fooStagedArtifact = RunnerApi.ArtifactInformation.newBuilder().setTypeUrn(BeamUrns.getUrn(RunnerApi.StandardArtifacts.Types.URL)).setTypePayload(RunnerApi.ArtifactUrlPayload.newBuilder().setUrl(stagingLocation + "/foo_staged.jar").build().toByteString()).setRoleUrn(BeamUrns.getUrn(RunnerApi.StandardArtifacts.Roles.STAGING_TO)).setRolePayload(RunnerApi.ArtifactStagingToRolePayload.newBuilder().setStagedName("foo_staged.jar").build().toByteString()).build();
    RunnerApi.ArtifactInformation barStagedArtifact = RunnerApi.ArtifactInformation.newBuilder().setTypeUrn(BeamUrns.getUrn(RunnerApi.StandardArtifacts.Types.URL)).setTypePayload(RunnerApi.ArtifactUrlPayload.newBuilder().setUrl(stagingLocation + "/bar_staged.jar").build().toByteString()).setRoleUrn(BeamUrns.getUrn(RunnerApi.StandardArtifacts.Roles.STAGING_TO)).setRolePayload(RunnerApi.ArtifactStagingToRolePayload.newBuilder().setStagedName("bar_staged.jar").build().toByteString()).build();
    RunnerApi.Pipeline expectedPipeline = RunnerApi.Pipeline.newBuilder().setComponents(RunnerApi.Components.newBuilder().putEnvironments("env", RunnerApi.Environment.newBuilder().addAllDependencies(ImmutableList.of(fooStagedArtifact, barStagedArtifact)).build())).build();
    assertThat(runner.resolveArtifacts(pipeline), equalTo(expectedPipeline));
}

Example 100 with Pipeline

use of org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline in project beam by apache.

the class DataflowRunnerTest method testStageDuplicatedArtifacts.

@Test
public void testStageDuplicatedArtifacts() throws IOException {
    DataflowPipelineOptions options = buildPipelineOptions();
    DataflowRunner runner = DataflowRunner.fromOptions(options);
    File foo = File.createTempFile("foo-", ".txt");
    foo.deleteOnExit();
    File bar = File.createTempFile("bar-", ".txt");
    bar.deleteOnExit();
    RunnerApi.ArtifactInformation foo1LocalArtifact = RunnerApi.ArtifactInformation.newBuilder().setTypeUrn(BeamUrns.getUrn(RunnerApi.StandardArtifacts.Types.FILE)).setTypePayload(RunnerApi.ArtifactFilePayload.newBuilder().setPath(foo.getAbsolutePath()).build().toByteString()).setRoleUrn(BeamUrns.getUrn(RunnerApi.StandardArtifacts.Roles.STAGING_TO)).setRolePayload(RunnerApi.ArtifactStagingToRolePayload.newBuilder().setStagedName("foo_staged1.jar").build().toByteString()).build();
    RunnerApi.ArtifactInformation foo2LocalArtifact = RunnerApi.ArtifactInformation.newBuilder().setTypeUrn(BeamUrns.getUrn(RunnerApi.StandardArtifacts.Types.FILE)).setTypePayload(RunnerApi.ArtifactFilePayload.newBuilder().setPath(foo.getAbsolutePath()).build().toByteString()).setRoleUrn(BeamUrns.getUrn(RunnerApi.StandardArtifacts.Roles.STAGING_TO)).setRolePayload(RunnerApi.ArtifactStagingToRolePayload.newBuilder().setStagedName("foo_staged2.jar").build().toByteString()).build();
    RunnerApi.ArtifactInformation barLocalArtifact = RunnerApi.ArtifactInformation.newBuilder().setTypeUrn(BeamUrns.getUrn(RunnerApi.StandardArtifacts.Types.FILE)).setTypePayload(RunnerApi.ArtifactFilePayload.newBuilder().setPath(bar.getAbsolutePath()).build().toByteString()).setRoleUrn(BeamUrns.getUrn(RunnerApi.StandardArtifacts.Roles.STAGING_TO)).setRolePayload(RunnerApi.ArtifactStagingToRolePayload.newBuilder().setStagedName("bar_staged.jar").build().toByteString()).build();
    RunnerApi.Environment env1 = RunnerApi.Environment.newBuilder().addAllDependencies(ImmutableList.of(foo1LocalArtifact, barLocalArtifact)).build();
    RunnerApi.Environment env2 = RunnerApi.Environment.newBuilder().addAllDependencies(ImmutableList.of(foo2LocalArtifact, barLocalArtifact)).build();
    RunnerApi.Pipeline pipeline = RunnerApi.Pipeline.newBuilder().setComponents(RunnerApi.Components.newBuilder().putEnvironments("env1", env1).putEnvironments("env2", env2)).build();
    List<DataflowPackage> packages = runner.stageArtifacts(pipeline);
    List<String> packageNames = packages.stream().map(DataflowPackage::getName).collect(Collectors.toList());
    assertThat(packageNames.size(), equalTo(3));
    assertThat(packageNames, containsInAnyOrder("foo_staged1.jar", "foo_staged2.jar", "bar_staged.jar"));
}