Java Code Examples for org.apache.beam.model.pipeline.v1.RunnerApi#ArtifactInformation

private RunnerApi.Pipeline resolveDependencies(RunnerApi.Pipeline pipeline, String stagingToken) {
  Map<String, List<RunnerApi.ArtifactInformation>> resolvedDependencies =
      stagingService.getService().getStagedArtifacts(stagingToken);
  Map<String, RunnerApi.Environment> newEnvironments = new HashMap<>();
  for (Map.Entry<String, RunnerApi.Environment> entry :
      pipeline.getComponents().getEnvironmentsMap().entrySet()) {
    if (entry.getValue().getDependenciesCount() > 0 && resolvedDependencies == null) {
      throw new RuntimeException(
          "Artifact dependencies provided but not staged for " + entry.getKey());
    }
    newEnvironments.put(
        entry.getKey(),
        entry.getValue().getDependenciesCount() == 0
            ? entry.getValue()
            : entry
                .getValue()
                .toBuilder()
                .clearDependencies()
                .addAllDependencies(resolvedDependencies.get(entry.getKey()))
                .build());
  }
  RunnerApi.Pipeline.Builder builder = pipeline.toBuilder();
  builder.getComponentsBuilder().clearEnvironments().putAllEnvironments(newEnvironments);
  return builder.build();
}
 

/**
 * Stages all artifact into the jar file at outputStream, returning a new artifact pointing to its
 * location in the classpath.
 */
private RunnerApi.ArtifactInformation writeArtifact(
    RunnerApi.ArtifactInformation artifact, String jobName) throws IOException {
  String path = PortablePipelineJarUtils.getArtifactUri(jobName, UUID.randomUUID().toString());
  LOG.trace("Copying artifact {} to {}", artifact, path);
  outputStream.putNextEntry(new JarEntry(path));
  try (InputStream artifactStream = ArtifactRetrievalService.getArtifact(artifact)) {
    ByteStreams.copy(artifactStream, outputStream);
  }
  return artifact
      .toBuilder()
      .setTypeUrn(ArtifactRetrievalService.FILE_ARTIFACT_URN)
      .setTypePayload(
          RunnerApi.ArtifactFilePayload.newBuilder()
              .setPath(ClassLoaderFileSystem.SCHEMA + "://" + path)
              .build()
              .toByteString())
      .build();
}
 

@Override
public List<RunnerApi.ArtifactInformation> resolveArtifacts(
    List<RunnerApi.ArtifactInformation> artifacts) {
  for (ResolutionFn fn : Lists.reverse(fns)) {
    List<RunnerApi.ArtifactInformation> moreResolved = new ArrayList<>();
    for (RunnerApi.ArtifactInformation artifact : artifacts) {
      Optional<List<RunnerApi.ArtifactInformation>> resolved = fn.resolve(artifact);
      if (resolved.isPresent()) {
        moreResolved.addAll(resolved.get());
      } else {
        moreResolved.add(artifact);
      }
    }
    artifacts = moreResolved;
  }
  return artifacts;
}
 

@Test
public void testStageArtifacts() throws InterruptedException, ExecutionException {
  List<String> contentsList =
      ImmutableList.of("a", "bb", Strings.repeat("xyz", TEST_BUFFER_SIZE * 3 / 4));
  stagingService.registerJob(
      "stagingToken",
      ImmutableMap.of(
          "env1",
          Lists.transform(contentsList, FakeArtifactRetrievalService::resolvedArtifact),
          "env2",
          Lists.transform(contentsList, FakeArtifactRetrievalService::unresolvedArtifact)));
  ArtifactStagingService.offer(new FakeArtifactRetrievalService(), stagingStub, "stagingToken");
  Map<String, List<RunnerApi.ArtifactInformation>> staged =
      stagingService.getStagedArtifacts("stagingToken");
  assertEquals(2, staged.size());
  checkArtifacts(contentsList, staged.get("env1"));
  checkArtifacts(contentsList, staged.get("env2"));
}
 

public static InputStream getArtifact(RunnerApi.ArtifactInformation artifact) throws IOException {
  switch (artifact.getTypeUrn()) {
    case FILE_ARTIFACT_URN:
      RunnerApi.ArtifactFilePayload payload =
          RunnerApi.ArtifactFilePayload.parseFrom(artifact.getTypePayload());
      return Channels.newInputStream(
          FileSystems.open(
              FileSystems.matchNewResource(payload.getPath(), false /* is directory */)));
    case EMBEDDED_ARTIFACT_URN:
      return RunnerApi.EmbeddedFilePayload.parseFrom(artifact.getTypePayload())
          .getData()
          .newInput();
    default:
      throw new UnsupportedOperationException(
          "Unexpected artifact type: " + artifact.getTypeUrn());
  }
}
 

private List<RunnerApi.ArtifactInformation> getDefaultArtifacts() {
  ImmutableList.Builder<String> pathsToStageBuilder = ImmutableList.builder();
  String windmillBinary =
      options.as(DataflowPipelineDebugOptions.class).getOverrideWindmillBinary();
  String dataflowWorkerJar = options.getDataflowWorkerJar();
  if (dataflowWorkerJar != null && !dataflowWorkerJar.isEmpty()) {
    // Put the user specified worker jar at the start of the classpath, to be consistent with the
    // built in worker order.
    pathsToStageBuilder.add("dataflow-worker.jar=" + dataflowWorkerJar);
  }
  pathsToStageBuilder.addAll(options.getFilesToStage());
  if (windmillBinary != null) {
    pathsToStageBuilder.add("windmill_main=" + windmillBinary);
  }
  return Environments.getArtifacts(pathsToStageBuilder.build());
}
 

private String getArtifact(RunnerApi.ArtifactInformation artifact) {
  ByteString all = ByteString.EMPTY;
  Iterator<ArtifactApi.GetArtifactResponse> response =
      retrievalBlockingStub.getArtifact(
          ArtifactApi.GetArtifactRequest.newBuilder().setArtifact(artifact).build());
  while (response.hasNext()) {
    all = all.concat(response.next().getData());
  }
  return all.toStringUtf8();
}
 

@Test
public void testResolveArtifacts() throws IOException {
  RunnerApi.ArtifactInformation artifact = fileArtifact(Paths.get("somePath"));
  ArtifactApi.ResolveArtifactsResponse resolved =
      retrievalBlockingStub.resolveArtifacts(
          ArtifactApi.ResolveArtifactsRequest.newBuilder().addArtifacts(artifact).build());
  assertEquals(1, resolved.getReplacementsCount());
  assertEquals(artifact, resolved.getReplacements(0));
}
 

public StoreArtifact(
    String stagingToken,
    String name,
    RunnerApi.ArtifactInformation originalArtifact,
    BlockingQueue<ByteString> bytesQueue,
    OverflowingSemaphore totalPendingBytes) {
  this.stagingToken = stagingToken;
  this.name = name;
  this.originalArtifact = originalArtifact;
  this.bytesQueue = bytesQueue;
  this.totalPendingBytes = totalPendingBytes;
}
 

private RunnerApi.Pipeline createEmptyPipeline(
    Iterable<RunnerApi.ArtifactInformation> dependencies) {
  return RunnerApi.Pipeline.newBuilder()
      .setComponents(
          RunnerApi.Components.newBuilder()
              .putEnvironments(
                  "env",
                  RunnerApi.Environment.newBuilder().addAllDependencies(dependencies).build()))
      .build();
}
 

public static RunnerApi.ArtifactInformation resolvedArtifact(String contents) {
  return RunnerApi.ArtifactInformation.newBuilder()
      .setTypeUrn("resolved")
      .setTypePayload(ByteString.copyFromUtf8(contents))
      .setRoleUrn(contents)
      .build();
}
 

private String getArtifact(RunnerApi.ArtifactInformation artifact) {
  ByteString all = ByteString.EMPTY;
  Iterator<ArtifactApi.GetArtifactResponse> response =
      retrievalBlockingStub.getArtifact(
          ArtifactApi.GetArtifactRequest.newBuilder().setArtifact(artifact).build());
  while (response.hasNext()) {
    all = all.concat(response.next().getData());
  }
  return all.toStringUtf8();
}
 

public static RunnerApi.ArtifactInformation unresolvedArtifact(String contents) {
  return RunnerApi.ArtifactInformation.newBuilder()
      .setTypeUrn("unresolved")
      .setTypePayload(ByteString.copyFromUtf8(contents))
      .setRoleUrn(contents)
      .build();
}
 

@Test
public void testWriteArtifacts_copiesArtifactsIntoJar() throws IOException {
  RunnerApi.ArtifactInformation artifact =
      RunnerApi.ArtifactInformation.newBuilder()
          .setTypeUrn(ArtifactRetrievalService.EMBEDDED_ARTIFACT_URN)
          .setTypePayload(
              RunnerApi.EmbeddedFilePayload.newBuilder()
                  .setData(ByteString.copyFromUtf8("someData"))
                  .build()
                  .toByteString())
          .setRoleUrn("someRole")
          .build();
  RunnerApi.Pipeline pipeline =
      RunnerApi.Pipeline.newBuilder()
          .setComponents(
              RunnerApi.Components.newBuilder()
                  .putEnvironments(
                      "env", RunnerApi.Environment.newBuilder().addDependencies(artifact).build())
                  .build())
          .build();
  RunnerApi.Pipeline newPipeline = jarCreator.writeArtifacts(pipeline, "jobName");
  verify(outputStream, times(1)).putNextEntry(any());
  RunnerApi.ArtifactInformation newArtifact =
      newPipeline.getComponents().getEnvironmentsMap().get("env").getDependencies(0);
  assertEquals(artifact.getRoleUrn(), newArtifact.getRoleUrn());
  assertNotEquals(artifact.getTypePayload(), newArtifact.getTypePayload());
}
 

/**
 * Stages all dependencies in environment into the jar file at outputStream, returning a new
 * environment that references these artifacts as classpath artifacts.
 */
private RunnerApi.Environment writeArtifacts(RunnerApi.Environment environment, String jobName)
    throws IOException {
  RunnerApi.Environment.Builder result = environment.toBuilder();
  result.clearDependencies();
  for (RunnerApi.ArtifactInformation artifact : environment.getDependenciesList()) {
    result.addDependencies(writeArtifact(artifact, jobName));
  }
  return result.build();
}
 

private RunnerApi.ArtifactInformation fileArtifact(Path path) {
  return fileArtifact(path, "");
}
 

private void checkArtifacts(
    Collection<String> expectedContents, List<RunnerApi.ArtifactInformation> staged) {
  assertEquals(
      expectedContents, Lists.transform(staged, RunnerApi.ArtifactInformation::getRoleUrn));
  assertEquals(expectedContents, Lists.transform(staged, this::getArtifact));
}
 

/**
 * Registers a set of artifacts to be staged with this service.
 *
 * <p>A client (e.g. a Beam SDK) is expected to connect to this service with the given staging
 * token and offer resolution and retrieval of this set of artifacts.
 *
 * @param stagingToken a staging token for this job
 * @param artifacts all artifacts to stage, keyed by environment
 */
public void registerJob(
    String stagingToken, Map<String, List<RunnerApi.ArtifactInformation>> artifacts) {
  assert !toStage.containsKey(stagingToken);
  toStage.put(stagingToken, artifacts);
}
 

/**
 * Returns the rewritten artifacts associated with this job, keyed by environment.
 *
 * <p>This should be called after the client has finished offering artifacts.
 *
 * @param stagingToken a staging token for this job
 */
public Map<String, List<RunnerApi.ArtifactInformation>> getStagedArtifacts(String stagingToken) {
  toStage.remove(stagingToken);
  return staged.remove(stagingToken);
}
 

Optional<List<RunnerApi.ArtifactInformation>> resolve(RunnerApi.ArtifactInformation info);