Skip to main content

dspk_hub/
lib.rs

1//! Verified download and local cache of published artifacts: data files and model weights that
2//! dspk crates use but do not ship (licences, size).
3//!
4//! - [`Artifact`]: what to fetch — id, family, file name, URI (`hf://`, `gh://`, `zenodo://`,
5//!   `https://`, `file://`, see [`providers`]), expected SHA-256 and size.
6//! - [`Hub`]: [`Hub::pull`] downloads into the [`HubCache`] and verifies hash and size before the
7//!   file is committed (an artifact without a hash is refused); [`Hub::status`] reads the cache
8//!   index; [`Hub::verify`] re-hashes.
9//! - [`catalog`]: which artifacts exist, where they come from and whom to credit, as data only.
10//!   Each domain is a feature (`synapse`: electrophysiology artifacts such as Kilosort4's
11//!   universal templates); [`pull`] fetches a catalog entry by id.
12//!
13//! The crate depends on no other dspk crate: the catalogs describe files, the crates that read
14//! them (sorters, models) live elsewhere. All network access of the workspace's libraries is here.
15
16pub mod cache;
17pub mod catalog;
18pub mod download;
19pub mod providers;
20
21use std::path::PathBuf;
22
23use anyhow::{Result, bail};
24
25pub use cache::{ArtifactStatus, CachedArtifactRecord, HubCache, HubCacheIndex, compute_file_sha256};
26pub use catalog::{Catalog, CatalogEntry, Credit, CreditCode, CreditPaper};
27pub use download::{DownloadOutcome, RemoteLinkCheck, check_remote_link, download_and_verify};
28pub use providers::{ProviderKind, resolve_gh_uri, resolve_hf_uri, resolve_weights_uri, resolve_zenodo_uri};
29
30/// A published file to fetch and verify.
31#[derive(Debug, Clone, PartialEq, Eq)]
32pub struct Artifact {
33    /// Unique id `"<family>/<slug>"`.
34    pub id: String,
35    pub family: String,
36    /// File name in the cache (e.g. `wTEMP.npz`).
37    pub file_name: String,
38    /// Download URI.
39    pub url: String,
40    /// Lowercase hex SHA-256.
41    pub sha256: String,
42    pub size_bytes: u64,
43}
44
45impl Artifact {
46    /// Last segment of the id (the cache sub-directory).
47    pub fn slug(&self) -> &str {
48        self.id.rsplit('/').next().unwrap_or(&self.id)
49    }
50}
51
52/// Local state and integrity of one artifact.
53#[derive(Debug, Clone)]
54pub struct ArtifactReport {
55    pub id: String,
56    pub status: ArtifactStatus,
57    pub local_path: PathBuf,
58    pub expected_sha256: String,
59    pub actual_sha256: Option<String>,
60    pub remote_check: Option<RemoteLinkCheck>,
61}
62
63/// Downloads into and reads from a [`HubCache`].
64#[derive(Debug, Clone)]
65pub struct Hub {
66    cache: HubCache,
67}
68
69impl Hub {
70    /// Cache at `$DSPK_HUB_DIR` or the user cache directory.
71    pub fn from_env_or_default() -> Self {
72        Self { cache: HubCache::from_env_or_default() }
73    }
74
75    pub fn with_cache(cache: HubCache) -> Self {
76        Self { cache }
77    }
78
79    pub fn cache(&self) -> &HubCache {
80        &self.cache
81    }
82
83    /// Local path of `artifact` (whether or not it is installed).
84    pub fn path(&self, artifact: &Artifact) -> PathBuf {
85        self.cache.path(artifact)
86    }
87
88    pub fn status(&self, artifact: &Artifact) -> ArtifactStatus {
89        self.cache.status(artifact)
90    }
91
92    /// Downloads and verifies `artifact` unless it is already installed (or `force`), returning
93    /// its local path.
94    pub fn pull(&self, artifact: &Artifact, force: bool) -> Result<PathBuf> {
95        let dest = self.cache.path(artifact);
96        if !force && self.cache.status(artifact) == ArtifactStatus::Installed {
97            return Ok(dest);
98        }
99        let outcome = download_and_verify(&artifact.url, &dest, &artifact.sha256, artifact.size_bytes)?;
100        self.cache.record_installed(artifact, &outcome.sha256, outcome.bytes_written)?;
101        Ok(dest)
102    }
103
104    /// Re-hashes the cached file (and, with `check_link`, probes the remote URL). A corrupted
105    /// file is an error.
106    pub fn verify(&self, artifact: &Artifact, check_link: bool) -> Result<ArtifactReport> {
107        let (status, actual_sha256) = self.cache.verify(artifact)?;
108        if status == ArtifactStatus::Corrupted {
109            bail!(
110                "corrupted cached file for '{}': expected sha256 {} and {} bytes, got {}",
111                artifact.id,
112                artifact.sha256,
113                artifact.size_bytes,
114                actual_sha256.as_deref().unwrap_or("<unreadable>")
115            );
116        }
117        let remote_check = if check_link { Some(check_remote_link(&artifact.url)?) } else { None };
118        Ok(ArtifactReport {
119            id: artifact.id.clone(),
120            status,
121            local_path: self.cache.path(artifact),
122            expected_sha256: artifact.sha256.clone(),
123            actual_sha256,
124            remote_check,
125        })
126    }
127
128    pub fn remove(&self, artifact: &Artifact) -> Result<bool> {
129        self.cache.remove(artifact)
130    }
131
132    pub fn clean(&self) -> Result<u64> {
133        self.cache.clean_all()
134    }
135}
136
137/// Downloads (unless installed) and verifies catalog entry `id` of [`Catalog::load_default`]
138/// into the default cache ([`Hub::from_env_or_default`]), returning its local path.
139pub fn pull(id: &str) -> Result<PathBuf> {
140    let catalog = Catalog::load_default()?;
141    let entry = catalog.get(id).ok_or_else(|| anyhow::anyhow!("unknown artifact id '{id}' (`dspk hub list` shows the catalog)"))?;
142    Hub::from_env_or_default().pull(&entry.artifact(), false)
143}
144
145#[cfg(test)]
146mod tests {
147    use super::*;
148    use sha2::{Digest, Sha256};
149
150    #[test]
151    fn pull_verify_status_and_remove_a_local_artifact() {
152        let root = std::env::temp_dir().join(format!("dspk_hub_test_{}", std::process::id()));
153        let _ = std::fs::remove_dir_all(&root);
154        std::fs::create_dir_all(&root).unwrap();
155        let bytes = b"LOCAL_TEST_ARTIFACT_BYTES";
156        let src = root.join("source.bin");
157        std::fs::write(&src, bytes).unwrap();
158        let artifact = Artifact {
159            id: "test/local-v1".into(),
160            family: "test".into(),
161            file_name: "local.bin".into(),
162            url: format!("file://{}", src.display()),
163            sha256: hex::encode(Sha256::digest(bytes)),
164            size_bytes: bytes.len() as u64,
165        };
166        let hub = Hub::with_cache(HubCache::new(root.join("cache")));
167        assert_eq!(hub.status(&artifact), ArtifactStatus::Available);
168        let path = hub.pull(&artifact, false).unwrap();
169        assert!(path.exists());
170        assert_eq!(hub.status(&artifact), ArtifactStatus::Installed);
171        assert!(hub.verify(&artifact, true).unwrap().remote_check.unwrap().reachable);
172
173        // A wrong size or an empty hash is refused
174        let wrong_size = Artifact { size_bytes: 1, id: "test/wrong-size".into(), ..artifact.clone() };
175        assert!(hub.pull(&wrong_size, false).is_err());
176        let no_hash = Artifact { sha256: String::new(), id: "test/no-hash".into(), ..artifact.clone() };
177        assert!(hub.pull(&no_hash, false).is_err());
178
179        assert!(hub.remove(&artifact).unwrap());
180        assert_eq!(hub.status(&artifact), ArtifactStatus::Available);
181        let _ = std::fs::remove_dir_all(&root);
182    }
183}