mcai_docker/
image.rs

1use crate::{
2  container::ContainerBuilder,
3  error::{Error::Docker as DockerError, Result},
4};
5use async_std::prelude::StreamExt;
6use bollard::{
7  auth::DockerCredentials,
8  image::{CreateImageOptions, ListImagesOptions},
9  Docker,
10};
11use log::info;
12use std::collections::HashMap;
13
14#[derive(Default)]
15pub struct Image {
16  name: String,
17  credentials: Option<DockerCredentials>,
18}
19
20impl Image {
21  pub fn new(name: &str) -> Self {
22    Image {
23      name: name.to_string(),
24      ..Default::default()
25    }
26  }
27
28  pub fn with_credentials(self, credentials: DockerCredentials) -> Self {
29    self.with_optional_credentials(Some(credentials))
30  }
31
32  pub fn with_optional_credentials(mut self, credentials: Option<DockerCredentials>) -> Self {
33    self.credentials = credentials;
34    self
35  }
36
37  async fn pull_image(&self, docker: &Docker) -> Result<()> {
38    let filters = HashMap::from([("reference".to_string(), vec![self.name.to_string()])]);
39
40    let list_options = Some(ListImagesOptions {
41      filters,
42      ..Default::default()
43    });
44
45    let image_list = docker.list_images(list_options).await?;
46
47    if image_list.is_empty() {
48      let image_options = Some(CreateImageOptions {
49        from_image: self.name.to_string(),
50        ..Default::default()
51      });
52
53      let mut stream = docker.create_image(image_options.clone(), None, self.credentials.clone());
54
55      let mut pull_completed = false;
56
57      while !pull_completed {
58        match stream.next().await {
59          Some(Ok(info)) => {
60            match info.progress {
61              None => info!("{}", info.status.as_ref().unwrap()),
62              _ => info!(
63                "{}: {:?}kB/{:?}kB",
64                info.status.as_ref().unwrap(),
65                info
66                  .progress_detail
67                  .clone()
68                  .unwrap_or_default()
69                  .current
70                  .unwrap_or_default(),
71                info
72                  .progress_detail
73                  .unwrap_or_default()
74                  .total
75                  .unwrap_or_default()
76              ),
77            }
78
79            if info.status
80              == Some(format!(
81                "Status: Downloaded newer image for {}",
82                image_options.as_ref().unwrap().from_image
83              ))
84            {
85              pull_completed = true;
86            }
87          }
88          Some(Err(error)) => {
89            return Err(DockerError(error));
90          }
91          _ => {
92            return Ok(());
93          }
94        }
95      }
96    }
97    Ok(())
98  }
99
100  pub async fn build(&self, docker: &Docker) -> Result<ContainerBuilder> {
101    self.pull_image(docker).await?;
102    Ok(ContainerBuilder::new(&self.name))
103  }
104}