Skip to main content

spawn_db/store/pinner/
spawn.rs

1use super::Pinner;
2use anyhow::Result;
3use anyhow::{anyhow, Context};
4use async_trait::async_trait;
5use opendal::Operator;
6use std::collections::HashMap;
7
8#[derive(Debug)]
9pub struct Spawn {
10    files: Option<HashMap<String, String>>,
11    pin_path: String,
12    source_path: String,
13}
14
15impl Spawn {
16    pub fn new(pin_path: String, source_path: String) -> Result<Self> {
17        let store = Self {
18            files: None,
19            pin_path,
20            source_path,
21        };
22
23        Ok(store)
24    }
25
26    pub async fn new_with_root_hash(
27        pin_path: String,
28        source_path: String,
29        root_hash: &str,
30        object_store: &Operator,
31    ) -> Result<Self> {
32        let mut files = HashMap::new();
33        Self::read_root_hash(object_store, &pin_path, &mut files, "", root_hash).await?;
34
35        let store = Self {
36            files: Some(files),
37            pin_path: pin_path.clone(),
38            source_path,
39        };
40
41        Ok(store)
42    }
43
44    async fn read_root_hash(
45        object_store: &Operator,
46        store_path: &str,
47        files: &mut HashMap<String, String>,
48        base_path: &str,
49        root_hash: &str,
50    ) -> Result<()> {
51        let contents = super::read_hash_file(object_store, store_path, root_hash)
52            .await
53            .context("cannot read root file")?;
54        let tree: super::Tree = toml::from_str(&contents).context("failed to parse tree TOML")?;
55
56        for (_, entry) in tree.entries.iter().enumerate() {
57            match entry.kind {
58                super::EntryKind::Blob => {
59                    let full_name = if base_path.is_empty() {
60                        entry.name.clone()
61                    } else {
62                        format!("{}/{}", base_path, &entry.name)
63                    };
64                    let full_path = format!("{}/{}", store_path, super::hash_to_path(&entry.hash)?);
65                    files.insert(full_name, full_path);
66                }
67                super::EntryKind::Tree => {
68                    let new_base = if base_path.is_empty() {
69                        entry.name.clone()
70                    } else {
71                        format!("{}/{}", base_path, &entry.name)
72                    };
73                    Box::pin(Self::read_root_hash(
74                        object_store,
75                        store_path,
76                        files,
77                        &new_base,
78                        &entry.hash,
79                    ))
80                    .await?;
81                }
82            }
83        }
84
85        Ok(())
86    }
87}
88
89#[async_trait]
90impl Pinner for Spawn {
91    /// Returns the file from the store if it exists.
92    async fn load(&self, name: &str, object_store: &Operator) -> Result<Option<String>> {
93        // Borrow files from inside self.files, if not none:
94        let files = self
95            .files
96            .as_ref()
97            .ok_or(anyhow!("files not initialized, was a root hash specified?"))?;
98
99        if let Some(path) = files.get(name) {
100            match object_store.read(path).await {
101                Ok(get_result) => {
102                    let bytes = get_result.to_bytes();
103                    let contents = String::from_utf8(bytes.to_vec())?;
104                    Ok(Some(contents))
105                }
106                Err(_) => Ok(None),
107            }
108        } else {
109            Ok(None)
110        }
111    }
112
113    async fn snapshot(&mut self, object_store: &Operator) -> Result<String> {
114        super::snapshot(object_store, &self.pin_path, &self.source_path).await
115    }
116}