ragit/index/commands/
pull.rs

1use super::Index;
2use super::archive::BlockType;
3use crate::constant::{
4    CHUNK_DIR_NAME,
5    CONFIG_DIR_NAME,
6    FILE_INDEX_DIR_NAME,
7    II_DIR_NAME,
8    IMAGE_DIR_NAME,
9    INDEX_DIR_NAME,
10    INDEX_FILE_NAME,
11    METADATA_FILE_NAME,
12    PROMPT_DIR_NAME,
13};
14use crate::error::Error;
15use ragit_fs::{
16    exists,
17    join,
18    remove_dir_all,
19    rename,
20};
21use reqwest::Url;
22
23pub enum PullResult {
24    PulledArchives,
25    AlreadyUpToDate,
26}
27
28impl Index {
29    // TODO: `include_configs` and `include_prompts` are not thoroughly tested yet
30    pub async fn pull(
31        &self,
32        include_configs: bool,
33        include_prompts: bool,
34        ii: bool,
35        quiet: bool,
36    ) -> Result<PullResult, Error> {
37        let Some(repo_url) = self.repo_url.clone() else {
38            return Err(Error::NoRemoteToPullFrom);
39        };
40
41        // compare remote uid and local uid. if they're the same do nothing
42        let url = Url::parse(&repo_url)?;
43        let get_uid_url = url.join("uid")?;
44
45        match self.get_uid("pull", get_uid_url.as_str()).await {
46            Ok(remote_uid) => {
47                let self_uid = self.calculate_uid(false  /* force */)?;
48
49                if remote_uid == self_uid {
50                    return Ok(PullResult::AlreadyUpToDate);
51                }
52            },
53            Err(e) => {
54                if !quiet {
55                    eprintln!("Failed to get {get_uid_url}: {e:?}");
56                }
57            },
58        }
59
60        let mut tmp_no = 0;
61        let mut tmp_clone_dir = format!("tmp-clone-dir-{tmp_no}");
62
63        while exists(&tmp_clone_dir) {
64            tmp_no += 1;
65            tmp_clone_dir = format!("tmp-clone-dir-{tmp_no}");
66        }
67
68        let result = self.pull_worker(
69            repo_url,
70            &tmp_clone_dir,
71            include_configs,
72            include_prompts,
73            ii,
74            quiet,
75        ).await;
76
77        if exists(&tmp_clone_dir) {
78            remove_dir_all(&tmp_clone_dir)?;
79        }
80
81        result
82    }
83
84    async fn pull_worker(
85        &self,
86        repo_url: String,
87        tmp_clone_dir: &str,
88        include_configs: bool,
89        include_prompts: bool,
90        ii: bool,
91        quiet: bool,
92    ) -> Result<PullResult, Error> {
93        let cloned_blocks = Index::clone(repo_url, Some(tmp_clone_dir.to_string()), ii, quiet).await?;
94        let cloned_configs = cloned_blocks.get(&BlockType::Config).map(|n| *n).unwrap_or(0) > 1;
95        let cloned_prompts = cloned_blocks.get(&BlockType::Prompt).map(|n| *n).unwrap_or(0) > 1;
96
97        let curr_index_dir = join(
98            &self.root_dir,
99            INDEX_DIR_NAME,
100        )?;
101        let new_index_dir = join(
102            &tmp_clone_dir,
103            INDEX_DIR_NAME,
104        )?;
105
106        // If power goes down while moving `.ragit/files/`, you can run `rag check --recover` to recover
107        // If power goes down while moving `chunks/`, `images/` or `index.json`, you can run `rag pull` again to recover
108        // If power goes down while moving `meta.json`... you cannot tell whether something's wrong or not. That's a problem
109        for dir in [
110            FILE_INDEX_DIR_NAME,
111            CHUNK_DIR_NAME,
112            IMAGE_DIR_NAME,
113        ].iter() {
114            remove_dir_all(&join(&curr_index_dir, dir)?)?;
115            rename(
116                &join(&new_index_dir, dir)?,
117                &join(&curr_index_dir, dir)?,
118            )?;
119        }
120
121        rename(
122            &join(&new_index_dir, INDEX_FILE_NAME)?,
123            &join(&curr_index_dir, INDEX_FILE_NAME)?,
124        )?;
125
126        if exists(&join(&new_index_dir, METADATA_FILE_NAME)?) {
127            rename(
128                &join(&new_index_dir, METADATA_FILE_NAME)?,
129                &join(&curr_index_dir, METADATA_FILE_NAME)?,
130            )?;
131        }
132
133        if include_configs && cloned_configs {
134            remove_dir_all(&join(&curr_index_dir, CONFIG_DIR_NAME)?)?;
135            rename(
136                &join(&new_index_dir, CONFIG_DIR_NAME)?,
137                &join(&curr_index_dir, CONFIG_DIR_NAME)?,
138            )?;
139        }
140
141        if include_prompts && cloned_prompts {
142            remove_dir_all(&join(&curr_index_dir, PROMPT_DIR_NAME)?)?;
143            rename(
144                &join(&new_index_dir, PROMPT_DIR_NAME)?,
145                &join(&curr_index_dir, PROMPT_DIR_NAME)?,
146            )?;
147        }
148
149        if exists(&join(&new_index_dir, II_DIR_NAME)?) {
150            remove_dir_all(&join(&curr_index_dir, II_DIR_NAME)?)?;
151            rename(
152                &join(&new_index_dir, II_DIR_NAME)?,
153                &join(&curr_index_dir, II_DIR_NAME)?,
154            )?;
155        }
156
157        Ok(PullResult::PulledArchives)
158    }
159}
ragit/index/commands/pull.rs

ragit/index/commands/
pull.rs
