syntax = "proto2";

package mlflow;

import "databricks.proto";
import "scalapb/scalapb.proto";

option java_generate_equals_and_hash = true;
option java_package = "com.databricks.api.proto.mlflow";
option py_generic_services = true;
option (scalapb.options) = {flat_package: true};

service DatabricksMlflowArtifactsService {
  // Fetch credentials to read from the specified MLflow artifact location
  //
  // Note: Even if no artifacts exist at the specified artifact location, this API will
  // still provide read credentials as long as the format of the location is valid.
  // Callers must subsequently check for the existence of the artifacts using the appropriate
  // cloud storage APIs (as determined by the `ArtifactCredentialType` property of the response)
  rpc getCredentialsForRead(GetCredentialsForRead) returns (GetCredentialsForRead.Response) {
    option (rpc) = {
      endpoints: [
        {
          method: "POST"
          path: "/mlflow/artifacts/credentials-for-read"
          since: {
            major: 2
            minor: 0
          }
        }
      ]
      visibility: PUBLIC_UNDOCUMENTED
    };
  }

  // Fetch credentials to write to the specified MLflow artifact location
  rpc getCredentialsForWrite(GetCredentialsForWrite) returns (GetCredentialsForWrite.Response) {
    option (rpc) = {
      endpoints: [
        {
          method: "POST"
          path: "/mlflow/artifacts/credentials-for-write"
          since: {
            major: 2
            minor: 0
          }
        }
      ]
      visibility: PUBLIC_UNDOCUMENTED
    };
  }

  // Initiate a multipart upload and return presinged URLs for uploading parts and aborting the
  // initiated multipart upload.
  rpc createMultipartUpload(CreateMultipartUpload) returns (CreateMultipartUpload.Response) {
    option (rpc) = {
      endpoints: [
        {
          method: "POST"
          path: "/mlflow/artifacts/create-multipart-upload"
          since: {
            major: 2
            minor: 0
          }
        }
      ]
      visibility: PUBLIC_UNDOCUMENTED
    };
  }

  // Complete a multipart upload
  rpc completeMultipartUpload(CompleteMultipartUpload) returns (CompleteMultipartUpload.Response) {
    option (rpc) = {
      endpoints: [
        {
          method: "POST"
          path: "/mlflow/artifacts/complete-multipart-upload"
          since: {
            major: 2
            minor: 0
          }
        }
      ]
      visibility: PUBLIC_UNDOCUMENTED
    };
  }

  // Get a presigned URL to upload a part. Used when we need to retry failed part uploads.
  rpc getPresignedUploadPartUrl(GetPresignedUploadPartUrl) returns (GetPresignedUploadPartUrl.Response) {
    option (rpc) = {
      endpoints: [
        {
          method: "GET"
          path: "/mlflow/artifacts/get-presigned-upload-part-url"
          since: {
            major: 2
            minor: 0
          }
        }
      ]
      visibility: PUBLIC_UNDOCUMENTED
    };
  }

  rpc getCredentialsForTraceDataDownload(GetCredentialsForTraceDataDownload) returns (GetCredentialsForTraceDataDownload.Response) {
    option (rpc) = {
      endpoints: [
        {
          method: "GET"
          path: "/mlflow/traces/{request_id}/credentials-for-data-download"
          since: {
            major: 2
            minor: 0
          }
        }
      ]
      visibility: PUBLIC_UNDOCUMENTED
    };
  }

  rpc getCredentialsForTraceDataUpload(GetCredentialsForTraceDataUpload) returns (GetCredentialsForTraceDataUpload.Response) {
    option (rpc) = {
      endpoints: [
        {
          method: "GET"
          path: "/mlflow/traces/{request_id}/credentials-for-data-upload"
          since: {
            major: 2
            minor: 0
          }
        }
      ]
      visibility: PUBLIC_UNDOCUMENTED
    };
  }

  rpc getCredentialsForLoggedModelUpload(GetCredentialsForLoggedModelUpload) returns (GetCredentialsForLoggedModelUpload.Response) {
    option (rpc) = {
      endpoints: [
        {
          method: "POST"
          path: "/mlflow/logged-models/{model_id}/artifacts/credentials-for-upload"
          since: {
            major: 2
            minor: 0
          }
        }
      ]
      visibility: PUBLIC_UNDOCUMENTED
    };
  }

  rpc getCredentialsForLoggedModelDownload(GetCredentialsForLoggedModelDownload) returns (GetCredentialsForLoggedModelDownload.Response) {
    option (rpc) = {
      endpoints: [
        {
          method: "POST"
          path: "/mlflow/logged-models/{model_id}/artifacts/credentials-for-download"
          since: {
            major: 2
            minor: 0
          }
        }
      ]
      visibility: PUBLIC_UNDOCUMENTED
    };
  }
}

// The type of a given artifact access credential
enum ArtifactCredentialType {
  // The credential is an Azure Shared Access Signature URI. For more information, see
  // https://docs.microsoft.com/en-us/azure/storage/common/storage-sas-overview
  AZURE_SAS_URI = 1;

  // The credential is an AWS Presigned URL. For more information, see
  // https://docs.aws.amazon.com/AmazonS3/latest/dev/ShareObjectPreSignedURL.html
  AWS_PRESIGNED_URL = 2;

  // The credential is a GCP Signed URL. For more information, see
  // https://cloud.google.com/storage/docs/access-control/signed-urls
  GCP_SIGNED_URL = 3;

  // The credential is an Azure Shared Access Signature URI for ADLS.  For more
  // information see
  // https://docs.microsoft.com/en-us/rest/api/storageservices/data-lake-storage-gen2
  // and
  // https://docs.microsoft.com/en-us/azure/storage/common/storage-sas-overview
  AZURE_ADLS_GEN2_SAS_URI = 4;
}

message ArtifactCredentialInfo {
  // The ID of the MLflow Run containing the artifact that can be accessed
  // with the credential
  optional string run_id = 1;

  // The path, relative to the Run's artifact root location, of the artifact
  // that can be accessed with the credential
  optional string path = 2;

  // The signed URI credential that provides access to the artifact
  optional string signed_uri = 3;

  message HttpHeader {
    // The HTTP header name
    optional string name = 1;

    // The HTTP header value
    optional string value = 2;
  }

  // A collection of HTTP headers that should be specified when uploading to
  // or downloading from the specified `signed_uri`
  repeated HttpHeader headers = 4;

  // The type of the signed credential URI (e.g., an AWS presigned URL
  // or an Azure Shared Access Signature URI)
  optional ArtifactCredentialType type = 5;
}

// Credentials returned by GetCredentialsForLoggedModelUpload and GetCredentialsForLoggedModelDownload for
// writing and reading respectively to the artifact locations of a logged model.
// Contains an ArtifactCredentialInfo object for the logged model artifacts.
message LoggedModelArtifactCredential {
  // The ID of the MLflow logged model containing the artifact that can be accessed with the credential
  optional string model_id = 1;

  // The ArtifactCredentialInfo object for the logged model artifacts
  optional ArtifactCredentialInfo credential_info = 2;
}

message GetCredentialsForRead {
  option (scalapb.message).extends = "com.databricks.rpc.RPC[$this.Response]";

  // The ID of the MLflow Run for which to fetch artifact read credentials
  optional string run_id = 1 [(validate_required) = true];

  // The artifact paths, relative to the Run's artifact root location, for which to
  // fetch artifact read credentials. Must not be empty.
  repeated string path = 2;

  // Token specifying the page of credentials to fetch for large requests that require pagination
  optional string page_token = 3;

  message Response {
    // Credentials for reading from the specified artifact locations
    repeated ArtifactCredentialInfo credential_infos = 2;

    // Token used to fetch the next page of credentials for large requests that require pagination
    optional string next_page_token = 3;

    // Older versions of the MLflow client use this field to retrieve credentials via point fetches.
    reserved 1;
  }
}

message GetCredentialsForWrite {
  option (scalapb.message).extends = "com.databricks.rpc.RPC[$this.Response]";

  // The ID of the MLflow Run for which to fetch artifact write credentials
  optional string run_id = 1 [(validate_required) = true];

  // The artifact paths, relative to the Run's artifact root location, for which to
  // fetch artifact write credentials. Must not be empty.
  repeated string path = 2;

  // Token specifying the page of credentials to fetch for large requests that require pagination
  optional string page_token = 3;

  message Response {
    // Credentials for writing to the specified artifact locations
    repeated ArtifactCredentialInfo credential_infos = 2;

    // Token used to fetch the next page of credentials for large requests that require pagination
    optional string next_page_token = 3;

    // Older versions of the MLflow client use this field to retrieve credentials via point fetches.
    reserved 1;
  }
}

message CreateMultipartUpload {
  option (scalapb.message).extends = "com.databricks.rpc.RPC[$this.Response]";

  // Run ID
  optional string run_id = 1 [(validate_required) = true];

  // Artifact path, relative to the Run's artifact root location (e.g. "path/to/file")
  optional string path = 2;

  // Number of file parts (chunks of data) to upload in the initiated multipart upload
  optional int64 num_parts = 3 [(validate_required) = true];

  message Response {
    // ID identifying the initiated multipart upload
    optional string upload_id = 1;

    // Credentials for uploading file parts in the initiated multipart upload
    repeated ArtifactCredentialInfo upload_credential_infos = 2;

    // Why doesn't the response contain a presigned URL for completing the multipart upload?
    // As reported in https://github.com/aws/aws-sdk-java/issues/1537, AWS Java SDK v1
    // doesn't support generating a presigned URL for completing a multipart upload.

    // Credential for aborting the initiated multipart upload
    optional ArtifactCredentialInfo abort_credential_info = 3;
  }
}

message PartEtag {
  optional int64 part_number = 1;

  optional string etag = 2;
}

message CompleteMultipartUpload {
  option (scalapb.message).extends = "com.databricks.rpc.RPC[$this.Response]";

  // Run ID
  optional string run_id = 1 [(validate_required) = true];

  // Artifact path, relative to the Run's artifact root location (e.g. "path/to/file")
  optional string path = 2;

  // ID identifying the multipart upload to complete
  optional string upload_id = 3 [(validate_required) = true];

  // A list of file parts uploaded in the multipart upload to complete
  repeated PartEtag part_etags = 4;

  message Response {}
}

message GetPresignedUploadPartUrl {
  option (scalapb.message).extends = "com.databricks.rpc.RPC[$this.Response]";

  // Run ID
  optional string run_id = 1 [(validate_required) = true];

  // Atifact path, relative to the Run's artifact root location (e.g. "path/to/file")
  optional string path = 2;

  // ID identifying the multipart upload in which the part is uploaded
  optional string upload_id = 3 [(validate_required) = true];

  // Part number
  optional int64 part_number = 4 [(validate_required) = true];

  message Response {
    // Credential for uploading the part
    optional ArtifactCredentialInfo upload_credential_info = 1;
  }
}

message GetCredentialsForTraceDataDownload {
  option (scalapb.message).extends = "com.databricks.rpc.RPC[$this.Response]";

  optional string request_id = 1 [(validate_required) = true];

  // Optional relative path within the trace artifact directory.
  // If not specified, defaults to "traces.json".
  // Used for downloading individual attachment files (e.g., "attachments/<attachment_id>").
  optional string path = 2;

  message Response {
    optional ArtifactCredentialInfo credential_info = 1;
  }
}

message GetCredentialsForTraceDataUpload {
  option (scalapb.message).extends = "com.databricks.rpc.RPC[$this.Response]";

  optional string request_id = 1 [(validate_required) = true];

  // Optional relative path within the trace artifact directory.
  // If not specified, defaults to "traces.json".
  // Used for uploading individual attachment files (e.g., "attachments/<attachment_id>").
  optional string path = 2;

  message Response {
    optional ArtifactCredentialInfo credential_info = 1;
  }
}

message GetCredentialsForLoggedModelUpload {
  option (scalapb.message).extends = "com.databricks.rpc.RPC[$this.Response]";

  // The ID of the LoggedModel for which to fetch artifact write credentials
  optional string model_id = 1 [(validate_required) = true];

  // The artifact paths, relative to the LoggedModel's artifact root location,
  // for which to fetch artifact write credentials
  repeated string paths = 2;

  // Token specifying the page of credentials to fetch for large requests that require pagination
  optional string page_token = 3;

  message Response {
    // Credentials for writing to the specified artifact locations
    repeated LoggedModelArtifactCredential credentials = 1;

    // Token used to fetch the next page of credentials for large requests that require pagination
    optional string next_page_token = 2;
  }
}

message GetCredentialsForLoggedModelDownload {
  option (scalapb.message).extends = "com.databricks.rpc.RPC[$this.Response]";

  // The ID of the LoggedModel for which to fetch artifact read credentials
  optional string model_id = 1 [(validate_required) = true];

  // The artifact paths, relative to the LoggedModel's artifact root location,
  // for which to fetch artifact read credentials
  repeated string paths = 2;

  // Token specifying the page of credentials to fetch for large requests that require pagination
  optional string page_token = 3;

  message Response {
    // Credentials for reading the specified artifact locations
    repeated LoggedModelArtifactCredential credentials = 1;

    // Token used to fetch the next page of credentials for large requests that require pagination
    optional string next_page_token = 2;
  }
}