haze/src/image.rs

use bollard::models::CreateImageInfo;
use bollard::query_parameters::CreateImageOptions;
use bollard::Docker;
use futures_util::StreamExt;
use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
use miette::{IntoDiagnostic, Result, WrapErr};
use std::collections::HashMap;
use std::fmt::{Display, Formatter};
use std::str::FromStr;

#[derive(Debug, Copy, Clone, PartialOrd, PartialEq)]
pub struct ImageVersion {
    pub major: u8,
    pub minor: u8,
    pub patch: u8,
}

impl FromStr for ImageVersion {
    type Err = ();

    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
        let mut parts = s.split('.');
        let major = parts.next().ok_or(())?.parse().map_err(|_| ())?;
        let minor = parts.next().ok_or(())?.parse().map_err(|_| ())?;
        let patch = parts.next().ok_or(())?.parse().map_err(|_| ())?;
        Ok(ImageVersion {
            major,
            minor,
            patch,
        })
    }
}

impl Display for ImageVersion {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}.{}.{}", self.major, self.minor, self.patch)
    }
}

pub async fn image_exists(docker: &Docker, image: &str) -> bool {
    docker.inspect_image(image).await.is_ok()
}

pub async fn image_version(docker: &Docker, image: &str) -> Option<ImageVersion> {
    let labels = docker.inspect_image(image).await.ok()?.config?.labels?;
    let label = labels.get("nl.icewind.haze.version")?;
    ImageVersion::from_str(label).ok()
}

pub async fn update_image(docker: &Docker, image: &str) -> Result<()> {
    if image_exists(docker, image).await {
        force_pull_image(docker, image).await?;
    }
    Ok(())
}

pub async fn pull_image(docker: &Docker, image: &str) -> Result<()> {
    if !image_exists(docker, image).await {
        force_pull_image(docker, image).await?;
    }
    Ok(())
}

pub async fn force_pull_image(docker: &Docker, image: &str) -> Result<()> {
    println!("Pulling image {}", image);

    let mut info_stream = docker.create_image(
        Some(CreateImageOptions {
            from_image: if image.contains(':') {
                Some(image.to_string())
            } else {
                Some(format!("{}:latest", image))
            },
            ..Default::default()
        }),
        None,
        None,
    );

    let bar_style = ProgressStyle::with_template(
        "{spinner:.green} {msg} [{elapsed_precise}] [{bar:.cyan/blue}] {bytes:>12}/{total_bytes}",
    )
    .unwrap();
    let mut bars: HashMap<String, ProgressBar> = HashMap::new();
    let mp = MultiProgress::new();

    while let Some(info) = info_stream.next().await {
        let info: CreateImageInfo = info
            .into_diagnostic()
            .wrap_err_with(|| format!("Error while pulling image {}", image))?;
        if let (Some(id), Some(status), Some(progress)) =
            (info.id, info.status, info.progress_detail)
        {
            let bar = bars.entry(id.clone()).or_insert_with(|| {
                let bar = ProgressBar::new(progress.total.unwrap_or_default() as u64)
                    .with_style(bar_style.clone())
                    .with_message(format!("{id:20} - {status:10}"));
                mp.add(bar)
            });
            bar.set_message(format!("{id:10} - {status:20}"));
            if let Some(total) = progress.total {
                bar.set_length(total as u64);
            }
            if let Some(current) = progress.current {
                bar.set_position(current as u64);
            }
        }
    }
    Ok(())
}