ray-dashboard-sdk 0.4.0

use std::{path::Path, sync::Arc};

use tracing::debug;

use crate::{client::HttpClient, schema::packages::PackageUri, utils::packaging::DirectoryPackage};

/// Client for the Ray Packages API.
///
/// Obtained via [`RayDashboardClient::packages`](crate::RayDashboardClient::packages).
///
/// Packages are content-addressed and always stored under the `gcs` protocol, so methods
/// take and return full package URIs of the form `gcs://_ray_pkg_<sha1>.zip`.
#[derive(Debug, Clone)]
pub struct PackagesClient {
    http: Arc<HttpClient>,
}

impl PackagesClient {
    pub(crate) fn new(http: Arc<HttpClient>) -> Self {
        Self { http }
    }

    /// Check if a package exists on the Ray cluster
    pub async fn exists(&self, uri: &PackageUri) -> crate::Result<bool> {
        let path = format!("/api/packages/{}/{}", uri.protocol, uri.name);
        let request = self.http.request(reqwest::Method::GET, &path)?;

        match request.send().await {
            Ok(response) => match response.status() {
                reqwest::StatusCode::OK => Ok(true),
                reqwest::StatusCode::NOT_FOUND => Ok(false),
                _ => Err(crate::Error::Generic(format!(
                    "Unexpected status code checking package existence: {}",
                    response.status()
                ))),
            },
            Err(e) => Err(e.into()),
        }
    }

    /// Upload a package to the Ray cluster
    pub async fn upload(&self, package: &DirectoryPackage) -> crate::Result<()> {
        debug!(
            "Uploading package {} ({} bytes)",
            package.uri(),
            package.data().len()
        );

        let path = format!("/api/packages/{}/{}", package.uri().protocol, package.uri().name);
        let request = self.http.request(reqwest::Method::PUT, &path)?;

        let response = request
            .body(package.data().to_vec())
            .send()
            .await?
            .error_for_status()?;

        debug!("Package uploaded successfully: {}", response.status());
        Ok(())
    }

    /// Upload a directory as a package if it doesn't already exist, returning the
    /// full package URI
    pub async fn upload_directory(&self, directory: &Path) -> crate::Result<PackageUri> {
        let directory = directory.to_owned();
        let package = blocking::unblock(move || DirectoryPackage::new(&directory)).await?;

        if self.exists(package.uri()).await? {
            debug!("Package {} already exists, skipping upload", package.uri());
        } else {
            self.upload(&package).await?;
        }

        Ok(package.uri().clone())
    }
}