pomelo_data/
source.rs

1use crate::error::DataError;
2use crate::loader::PRICES_DIR;
3use std::fs;
4use std::path::{Path, PathBuf};
5
6/// Read-only byte store. `Ok(None)` means the key is absent (fail-soft).
7pub trait ObjectSource {
8    fn get(&self, key: &str) -> Result<Option<Vec<u8>>, DataError>;
9}
10
11/// Files under a root directory; `key` is a relative path (e.g. "prices/AAPL.csv.gz").
12pub struct LocalSource {
13    root: PathBuf,
14}
15
16impl LocalSource {
17    pub fn new(root: impl Into<PathBuf>) -> Self {
18        LocalSource { root: root.into() }
19    }
20}
21
22impl ObjectSource for LocalSource {
23    fn get(&self, key: &str) -> Result<Option<Vec<u8>>, DataError> {
24        match fs::read(self.root.join(key)) {
25            Ok(bytes) => Ok(Some(bytes)),
26            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
27            Err(e) => Err(DataError::Io(e.to_string())),
28        }
29    }
30}
31
32/// Write-side counterpart to [`ObjectSource`]. Kept separate so the read path
33/// (and OSS consumers) stay write-free; only the panel rebuild needs this.
34pub trait ObjectSink {
35    fn put(&self, key: &str, bytes: &[u8]) -> Result<(), DataError>;
36}
37
38impl ObjectSink for LocalSource {
39    fn put(&self, key: &str, bytes: &[u8]) -> Result<(), DataError> {
40        let path = self.root.join(key);
41        if let Some(parent) = path.parent() {
42            fs::create_dir_all(parent).map_err(|e| DataError::Io(e.to_string()))?;
43        }
44        fs::write(path, bytes).map_err(|e| DataError::Io(e.to_string()))
45    }
46}
47
48/// Symbols with a per-symbol price file under `root/prices`, sorted and
49/// de-duplicated. Recognizes `.csv.gz`, `.parquet`, and `.csv`; the loaders
50/// detect the actual format from content.
51pub fn list_symbols(root: &Path) -> std::io::Result<Vec<String>> {
52    // `.csv.gz` before `.csv` so a gzip file isn't mis-stripped to "<sym>.csv".
53    const EXTS: &[&str] = &[".csv.gz", ".parquet", ".csv"];
54    let mut syms = std::collections::BTreeSet::new();
55    let prices = root.join(PRICES_DIR);
56    if !prices.exists() {
57        return Ok(Vec::new());
58    }
59    for entry in fs::read_dir(prices)? {
60        let entry = entry?;
61        if !entry.file_type()?.is_file() {
62            continue;
63        }
64        if let Some(name) = entry.file_name().to_str() {
65            if let Some(sym) = EXTS.iter().find_map(|ext| name.strip_suffix(ext)) {
66                syms.insert(sym.to_string());
67            }
68        }
69    }
70    Ok(syms.into_iter().collect())
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76    use std::fs;
77
78    #[test]
79    fn local_source_reads_present_and_missing() {
80        let dir = std::env::temp_dir().join("pomelo_data_source_test");
81        fs::create_dir_all(&dir).unwrap();
82        fs::write(dir.join("hello.bin"), b"hi").unwrap();
83        let src = LocalSource::new(&dir);
84        assert_eq!(src.get("hello.bin").unwrap(), Some(b"hi".to_vec()));
85        assert_eq!(src.get("nope.bin").unwrap(), None);
86    }
87
88    #[test]
89    fn local_source_put_writes_and_creates_parents() {
90        let dir = std::env::temp_dir().join("pomelo_data_sink_test");
91        let _ = fs::remove_dir_all(&dir);
92        fs::create_dir_all(&dir).unwrap();
93        let src = LocalSource::new(&dir);
94        // a key with a missing parent dir ("panels/") must still write
95        src.put("panels/close.csv.gz", b"data").unwrap();
96        assert_eq!(
97            src.get("panels/close.csv.gz").unwrap(),
98            Some(b"data".to_vec())
99        );
100    }
101
102    #[test]
103    fn list_symbols_finds_and_dedups_price_stems() {
104        let dir = std::env::temp_dir().join("pomelo_data_list_symbols_test");
105        let _ = fs::remove_dir_all(&dir);
106        let prices = dir.join(PRICES_DIR);
107        fs::create_dir_all(&prices).unwrap();
108        fs::write(prices.join("AAPL.csv.gz"), b"x").unwrap();
109        fs::write(prices.join("MSFT.csv"), b"x").unwrap();
110        fs::write(prices.join("GOOG.parquet"), b"x").unwrap();
111        fs::write(prices.join("notes.txt"), b"x").unwrap();
112        assert_eq!(
113            list_symbols(&dir).unwrap(),
114            vec!["AAPL".to_string(), "GOOG".to_string(), "MSFT".to_string()]
115        );
116    }
117
118    #[test]
119    fn list_symbols_missing_prices_dir_is_empty() {
120        let dir = std::env::temp_dir().join("pomelo_data_list_symbols_missing_test");
121        let _ = fs::remove_dir_all(&dir);
122        fs::create_dir_all(&dir).unwrap();
123        assert_eq!(list_symbols(&dir).unwrap(), Vec::<String>::new());
124    }
125}
pomelo_data/source.rs

pomelo_data/
source.rs
