rskit_dataset/
source.rs

1//! Source trait — pull data from any origin.
2
3use crate::{DataItem, DatasetItem};
4use rskit_errors::AppResult;
5use std::pin::Pin;
6use tokio_util::sync::CancellationToken;
7
8/// Boxed stream of items emitted by a source.
9pub type BoxItemStream<T> = Pin<Box<dyn futures::Stream<Item = AppResult<T>> + Send + 'static>>;
10
11/// Boxed dataset stream of [`DataItem`] samples.
12pub type BoxDataStream = BoxItemStream<DataItem>;
13
14/// Protocol for dataset sources over any [`DatasetItem`] type.
15pub trait Source<T: DatasetItem>: Send + Sync {
16    /// Stable source identifier used in manifests.
17    fn name(&self) -> &str;
18
19    /// Human-readable source label.
20    fn display_name(&self) -> &str {
21        self.name().rsplit('/').next().unwrap_or(self.name())
22    }
23
24    /// Stream items from this source.
25    fn stream(self: Box<Self>, cancel: CancellationToken) -> BoxItemStream<T>;
26
27    /// Stable cache key describing this source's configured inputs.
28    fn cache_key(&self) -> serde_json::Value;
29    /// Optional maximum number of items this source will emit.
30    fn max_items(&self) -> Option<usize>;
31
32    /// Configure resume state before streaming.
33    fn set_resume_state(&mut self, _offset: usize, _already_fetched: usize) {}
34}
35
36#[cfg(test)]
37mod tests {
38    use futures::stream;
39    use serde_json::json;
40
41    use super::*;
42
43    struct StaticSource {
44        name: String,
45        resume: Option<(usize, usize)>,
46    }
47
48    struct DefaultResumeSource;
49
50    impl Source<DataItem> for DefaultResumeSource {
51        fn name(&self) -> &str {
52            "default"
53        }
54
55        fn stream(self: Box<Self>, _cancel: CancellationToken) -> BoxDataStream {
56            Box::pin(stream::empty())
57        }
58
59        fn cache_key(&self) -> serde_json::Value {
60            json!("default")
61        }
62
63        fn max_items(&self) -> Option<usize> {
64            None
65        }
66    }
67
68    impl Source<DataItem> for StaticSource {
69        fn name(&self) -> &str {
70            &self.name
71        }
72
73        fn stream(self: Box<Self>, _cancel: CancellationToken) -> BoxDataStream {
74            Box::pin(stream::empty())
75        }
76
77        fn cache_key(&self) -> serde_json::Value {
78            json!({"name": self.name})
79        }
80
81        fn max_items(&self) -> Option<usize> {
82            Some(0)
83        }
84
85        fn set_resume_state(&mut self, offset: usize, already_fetched: usize) {
86            self.resume = Some((offset, already_fetched));
87        }
88    }
89
90    #[test]
91    fn source_defaults_derive_display_name_and_allow_resume_state() {
92        let mut source = StaticSource {
93            name: "fixtures/images".to_string(),
94            resume: None,
95        };
96
97        assert_eq!(source.name(), "fixtures/images");
98        assert_eq!(source.display_name(), "images");
99        assert_eq!(source.cache_key(), json!({"name": "fixtures/images"}));
100        assert_eq!(source.max_items(), Some(0));
101
102        source.set_resume_state(10, 3);
103        assert_eq!(source.resume, Some((10, 3)));
104
105        let stream = Box::new(source).stream(CancellationToken::new());
106        drop(stream);
107
108        let mut default_resume = DefaultResumeSource;
109        default_resume.set_resume_state(1, 1);
110        assert_eq!(default_resume.display_name(), "default");
111        assert_eq!(default_resume.cache_key(), json!("default"));
112        assert_eq!(default_resume.max_items(), None);
113        let stream = Box::new(default_resume).stream(CancellationToken::new());
114        drop(stream);
115    }
116}
rskit_dataset/source.rs

rskit_dataset/
source.rs
