/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/*
 * Protocol Buffers describing the Artifact API, for communicating with a runner
 * for artifact staging and retrieval over GRPC.
 */

syntax = "proto3";

package org.apache.beam.model.job_management.v1;

option go_package = "github.com/apache/beam/sdks/v2/go/pkg/beam/model/jobmanagement_v1;jobmanagement_v1";
option java_package = "org.apache.beam.model.jobmanagement.v1";
option java_outer_classname = "ArtifactApi";

import "org/apache/beam/model/pipeline/v1/beam_runner_api.proto";

// A service to retrieve artifacts for use in a Job.
service ArtifactRetrievalService {
  // Resolves the given artifact references into one or more replacement
  // artifact references (e.g. a Maven dependency into a (transitive) set
  // of jars.
  rpc ResolveArtifacts(ResolveArtifactsRequest) returns (ResolveArtifactsResponse);

  // Retrieves the given artifact as a stream of bytes.
  rpc GetArtifact(GetArtifactRequest) returns (stream GetArtifactResponse);

  // TODO(robertwb): Consider a MergeArtifacts rpc that takes a list of
  // lists of artifacts and tries to consistently resolve them.
}

// A service that allows the client to act as an ArtifactRetrievalService,
// for a particular job with the server initiating requests and receiving
// responses.
//
// A client calls the service with an ArtifactResponseWrapper that has the
// staging token set, and thereafter responds to the server's requests.
service ArtifactStagingService {
  rpc ReverseArtifactRetrievalService(stream ArtifactResponseWrapper)
      returns (stream ArtifactRequestWrapper);
}

// A request for artifact resolution.
message ResolveArtifactsRequest {
  // An (ordered) set of artifacts to (jointly) resolve.
  repeated org.apache.beam.model.pipeline.v1.ArtifactInformation artifacts = 1;

  // A set of artifact type urns that are understood by the requester.
  // An attempt should be made to resolve the artifacts in terms of these URNs,
  // but other URNs may be used as well with the understanding that they must
  // be fetch-able as bytes via GetArtifact.
  repeated string preferred_urns = 2;
}

// A response for artifact resolution.
message ResolveArtifactsResponse {
  // A full (ordered) set of replacements for the set of requested artifacts,
  // preferably in terms of the requested type URNs.  If there is no better
  // resolution, the original list is returned.
  repeated org.apache.beam.model.pipeline.v1.ArtifactInformation replacements = 1;
}

// A request to get an artifact.
message GetArtifactRequest {
  org.apache.beam.model.pipeline.v1.ArtifactInformation artifact = 1;
}

// Part of a response to getting an artifact.
message GetArtifactResponse {
  bytes data = 1;
}

// Wraps an ArtifactRetrievalService request for use in ReverseArtifactRetrievalService.
message ArtifactRequestWrapper {
  oneof request {
    ResolveArtifactsRequest resolve_artifact = 1000;
    GetArtifactRequest get_artifact = 1001;
  }
}

// Wraps an ArtifactRetrievalService response for use in ReverseArtifactRetrievalService.
message ArtifactResponseWrapper {
  // A token indicating which job these artifacts are being staged for.
  string staging_token = 1;

  // Whether this is the last response for this request (for those responses that
  // would typically be terminated by the end of the response stream.)
  bool is_last = 2;

  // The response itself.
  oneof response {
    ResolveArtifactsResponse resolve_artifact_response = 1000;
    GetArtifactResponse get_artifact_response = 1001;
  }
}


// Legacy artifact staging service for pipeline-level artifacts.

// A service to stage artifacts for use in a Job.
service LegacyArtifactStagingService {
  // Stage an artifact to be available during job execution. The first request must contain the
  // name of the artifact. All future requests must contain sequential chunks of the content of
  // the artifact.
  rpc PutArtifact(stream PutArtifactRequest) returns (PutArtifactResponse);

  // Commit the manifest for a Job. All artifacts must have been successfully uploaded
  // before this call is made.
  //
  // Throws error INVALID_ARGUMENT if not all of the members of the manifest are present
  rpc CommitManifest(CommitManifestRequest) returns (CommitManifestResponse);
}

// A service to retrieve artifacts for use in a Job.
service LegacyArtifactRetrievalService {
  // Get the manifest for the job
  rpc GetManifest(GetManifestRequest) returns (GetManifestResponse);

  // Get an artifact staged for the job. The requested artifact must be within the manifest
  rpc GetArtifact(LegacyGetArtifactRequest) returns (stream ArtifactChunk);
}

// An artifact identifier and associated metadata.
message ArtifactMetadata {
  // (Required) The name of the artifact.
  string name = 1;

  // (Optional) The Unix-like permissions of the artifact
  uint32 permissions = 2;

  // (Optional) The hex-encoded sha256 checksum of the artifact. Used, among other things, by
  // harness boot code to validate the integrity of the artifact.
  string sha256 = 4;
}

// A collection of artifacts.
message Manifest {
  repeated ArtifactMetadata artifact = 1;
}

// A manifest with location information.
message ProxyManifest {
  Manifest manifest = 1;
  message Location {
     string name = 1;
     string uri = 2;
  }
  repeated Location location = 2;
}

// A request to get the manifest of a Job.
message GetManifestRequest {
  // (Required) An opaque token representing the entirety of the staged artifacts.
  // Returned in CommitManifestResponse.
  string retrieval_token = 1;
}

// A response containing a job manifest.
message GetManifestResponse {
  Manifest manifest = 1;
}

// A request to get an artifact. The artifact must be present in the manifest for the job.
message LegacyGetArtifactRequest {
  // (Required) The name of the artifact to retrieve.
  string name = 1;
  // (Required) An opaque token representing the entirety of the staged artifacts.
  // Returned in CommitManifestResponse.
  string retrieval_token = 2;
}

// Part of an artifact.
message ArtifactChunk {
  bytes data = 1;
}

message PutArtifactMetadata {
  // (Required) A token for artifact staging session. This token can be obtained
  // from PrepareJob request in JobService
  string staging_session_token = 1;
  // (Required) The Artifact metadata.
  ArtifactMetadata metadata = 2;
}

// A request to stage an artifact.
message PutArtifactRequest {
  // (Required)
  oneof content {
    // The first message in a PutArtifact call must contain this field.
    PutArtifactMetadata metadata = 1;

    // A chunk of the artifact. All messages after the first in a PutArtifact call must contain a
    // chunk.
    ArtifactChunk data = 2;
  }
}

message PutArtifactResponse {
}

// A request to commit the manifest for a Job. All artifacts must have been successfully uploaded
// before this call is made.
message CommitManifestRequest {
  // (Required) The manifest to commit.
  Manifest manifest = 1;
  // (Required) A token for artifact staging session. This token can be obtained
  // from PrepareJob request in JobService
  string staging_session_token = 2;
}

// The result of committing a manifest.
message CommitManifestResponse {
  enum Constants {
    // Token indicating that no artifacts were staged and therefore no retrieval attempt is necessary.
    NO_ARTIFACTS_STAGED_TOKEN = 0 [(org.apache.beam.model.pipeline.v1.beam_constant) = "__no_artifacts_staged__"];
  }
  // (Required) An opaque token representing the entirety of the staged artifacts.
  // This can be used to retrieve the manifest and artifacts from an associated
  // LegacyArtifactRetrievalService.
  string retrieval_token = 1;
}
