rskit_dataset/record/
reader.rs

1//! Streaming readers for structured dataset records.
2
3use std::path::{Path, PathBuf};
4
5use futures::stream;
6use rskit_errors::{AppError, AppResult, ErrorCode};
7use serde_json::Value;
8use tokio::sync::mpsc;
9
10use super::limits::{
11    MAX_CSV_RECORD_BYTES, MAX_JSON_ARRAY_BYTES, MAX_JSON_LINE_BYTES, validate_json_record,
12};
13use super::model::DatasetRecord;
14use super::ops::BoxRecordStream;
15
16/// Streaming reader for structured dataset records.
17pub trait DatasetReader: Send + Sync + 'static {
18    /// Convert this reader into a record stream.
19    fn stream(self: Box<Self>) -> BoxRecordStream;
20}
21
22/// CSV record reader backed by the `csv` crate.
23pub struct CsvReader {
24    path: PathBuf,
25}
26
27impl CsvReader {
28    /// Create a CSV reader for a local path.
29    #[must_use]
30    pub fn new(path: impl Into<PathBuf>) -> Self {
31        Self { path: path.into() }
32    }
33}
34
35impl DatasetReader for CsvReader {
36    fn stream(self: Box<Self>) -> BoxRecordStream {
37        let path = self.path;
38        stream_from_blocking_reader(move |tx| {
39            let file = match std::fs::File::open(&path) {
40                Ok(file) => file,
41                Err(error) => {
42                    send_record(
43                        &tx,
44                        Err(AppError::new(
45                            ErrorCode::Internal,
46                            format!("failed to open CSV dataset {}: {error}", path.display()),
47                        )),
48                    );
49                    return;
50                }
51            };
52            let mut reader = std::io::BufReader::new(file);
53
54            let headers = match read_line_bounded(&mut reader, MAX_CSV_RECORD_BYTES, "CSV header")
55                .and_then(|line| {
56                    line.map_or_else(
57                        || {
58                            Err(AppError::new(
59                                ErrorCode::InvalidInput,
60                                format!("CSV dataset {} is missing headers", path.display()),
61                            ))
62                        },
63                        |line| parse_csv_record(&line),
64                    )
65                }) {
66                Ok(headers) => headers,
67                Err(error) => {
68                    send_record(&tx, Err(error));
69                    return;
70                }
71            };
72
73            loop {
74                match read_line_bounded(&mut reader, MAX_CSV_RECORD_BYTES, "CSV record")
75                    .and_then(|line| line.map(|line| parse_csv_record(&line)).transpose())
76                {
77                    Ok(Some(raw)) => {
78                        let fields = headers
79                            .iter()
80                            .zip(raw.iter())
81                            .map(|(key, value)| (key.to_string(), Value::String(value.to_string())))
82                            .collect();
83                        if !send_record(&tx, Ok(DatasetRecord::new(fields))) {
84                            return;
85                        }
86                    }
87                    Ok(None) => return,
88                    Err(error) => {
89                        send_record(&tx, Err(error));
90                        return;
91                    }
92                }
93            }
94        })
95    }
96}
97
98/// Newline-delimited JSON record reader.
99pub struct JsonLinesReader {
100    path: PathBuf,
101}
102
103impl JsonLinesReader {
104    /// Create a JSON Lines reader for a local path.
105    #[must_use]
106    pub fn new(path: impl Into<PathBuf>) -> Self {
107        Self { path: path.into() }
108    }
109}
110
111impl DatasetReader for JsonLinesReader {
112    fn stream(self: Box<Self>) -> BoxRecordStream {
113        let path = self.path;
114        stream_from_blocking_reader(move |tx| {
115            let file = match std::fs::File::open(&path) {
116                Ok(file) => file,
117                Err(error) => {
118                    send_record(
119                        &tx,
120                        Err(AppError::new(
121                            ErrorCode::Internal,
122                            format!(
123                                "failed to open JSON Lines dataset {}: {error}",
124                                path.display()
125                            ),
126                        )),
127                    );
128                    return;
129                }
130            };
131            let mut reader = std::io::BufReader::new(file);
132            loop {
133                let record = match read_json_line_bounded(&mut reader) {
134                    Ok(Some(line)) => record_from_json_bytes(&line),
135                    Ok(None) => return,
136                    Err(error) => {
137                        send_record(&tx, Err(error));
138                        return;
139                    }
140                };
141                if !send_record(&tx, record) {
142                    return;
143                }
144            }
145        })
146    }
147}
148
149/// Bounded JSON array reader for small fixture-style datasets.
150pub struct JsonArrayReader {
151    path: PathBuf,
152}
153
154impl JsonArrayReader {
155    /// Create a JSON array reader for a local path.
156    #[must_use]
157    pub fn new(path: impl Into<PathBuf>) -> Self {
158        Self { path: path.into() }
159    }
160}
161
162impl DatasetReader for JsonArrayReader {
163    fn stream(self: Box<Self>) -> BoxRecordStream {
164        let path = self.path;
165        stream_from_blocking_reader(move |tx| {
166            let values =
167                match read_bounded_file(&path, MAX_JSON_ARRAY_BYTES as usize).and_then(|bytes| {
168                    serde_json::from_slice::<Vec<Value>>(&bytes).map_err(|error| {
169                        AppError::new(
170                            ErrorCode::InvalidInput,
171                            format!("failed to parse JSON dataset {}: {error}", path.display()),
172                        )
173                    })
174                }) {
175                    Ok(values) => values,
176                    Err(error) => {
177                        send_record(&tx, Err(error));
178                        return;
179                    }
180                };
181
182            for value in values {
183                if !send_record(&tx, record_from_value(value)) {
184                    return;
185                }
186            }
187        })
188    }
189}
190
191fn stream_from_blocking_reader(
192    producer: impl FnOnce(mpsc::Sender<AppResult<DatasetRecord>>) + Send + 'static,
193) -> BoxRecordStream {
194    if tokio::runtime::Handle::try_current().is_err() {
195        return Box::pin(stream::once(async {
196            Err(AppError::new(
197                ErrorCode::Internal,
198                "dataset readers require a Tokio runtime for blocking IO isolation",
199            ))
200        }));
201    }
202
203    let (tx, rx) = mpsc::channel(8);
204    let handle = tokio::task::spawn_blocking(move || producer(tx));
205    drop(handle);
206    Box::pin(stream::unfold(rx, |mut rx| async {
207        rx.recv().await.map(|item| (item, rx))
208    }))
209}
210
211fn send_record(
212    tx: &mpsc::Sender<AppResult<DatasetRecord>>,
213    item: AppResult<DatasetRecord>,
214) -> bool {
215    tx.blocking_send(item).is_ok()
216}
217
218fn record_from_json_bytes(line: &[u8]) -> AppResult<DatasetRecord> {
219    let value = serde_json::from_slice(line).map_err(|error| {
220        AppError::new(
221            ErrorCode::InvalidInput,
222            format!("failed to parse JSON record: {error}"),
223        )
224    })?;
225    record_from_value(value)
226}
227
228fn read_json_line_bounded(reader: &mut impl std::io::BufRead) -> AppResult<Option<Vec<u8>>> {
229    read_line_bounded(reader, MAX_JSON_LINE_BYTES, "JSON Lines record")
230}
231
232fn read_line_bounded(
233    reader: &mut impl std::io::BufRead,
234    max_bytes: usize,
235    label: &str,
236) -> AppResult<Option<Vec<u8>>> {
237    let mut line = Vec::new();
238
239    loop {
240        let available = reader.fill_buf().map_err(|error| {
241            AppError::new(
242                ErrorCode::Internal,
243                format!("failed to read JSON line: {error}"),
244            )
245        })?;
246        if available.is_empty() {
247            return if line.is_empty() {
248                Ok(None)
249            } else {
250                Ok(Some(line))
251            };
252        }
253
254        let consumed = match available.iter().position(|byte| *byte == b'\n') {
255            Some(pos) => {
256                let end = pos + 1;
257                append_line_chunk(&mut line, &available[..end], max_bytes, label)?;
258                end
259            }
260            None => {
261                append_line_chunk(&mut line, available, max_bytes, label)?;
262                available.len()
263            }
264        };
265        reader.consume(consumed);
266
267        if line.last() == Some(&b'\n') {
268            return Ok(Some(line));
269        }
270    }
271}
272
273fn append_line_chunk(
274    line: &mut Vec<u8>,
275    chunk: &[u8],
276    max_bytes: usize,
277    label: &str,
278) -> AppResult<()> {
279    if line.len().saturating_add(chunk.len()) > max_bytes {
280        return Err(AppError::new(
281            ErrorCode::InvalidInput,
282            format!("{label} exceeded max {max_bytes} bytes"),
283        ));
284    }
285    line.extend_from_slice(chunk);
286    Ok(())
287}
288
289fn parse_csv_record(line: &[u8]) -> AppResult<csv::StringRecord> {
290    let mut reader = csv::ReaderBuilder::new()
291        .has_headers(false)
292        .from_reader(line);
293    let mut records = reader.records();
294    match records.next() {
295        Some(Ok(record)) => Ok(record),
296        Some(Err(error)) => Err(AppError::new(
297            ErrorCode::InvalidInput,
298            format!("failed to parse CSV record: {error}"),
299        )),
300        None => Err(AppError::new(
301            ErrorCode::InvalidInput,
302            "CSV record was empty",
303        )),
304    }
305}
306
307fn read_bounded_file(path: &Path, max_bytes: usize) -> AppResult<Vec<u8>> {
308    use std::io::Read as _;
309
310    let mut file = std::fs::File::open(path).map_err(|error| {
311        AppError::new(
312            ErrorCode::Internal,
313            format!("failed to open JSON dataset {}: {error}", path.display()),
314        )
315    })?;
316    let mut bytes = Vec::new();
317    file.by_ref()
318        .take(max_bytes as u64 + 1)
319        .read_to_end(&mut bytes)
320        .map_err(|error| {
321            AppError::new(
322                ErrorCode::Internal,
323                format!("failed to read JSON dataset {}: {error}", path.display()),
324            )
325        })?;
326    if bytes.len() > max_bytes {
327        return Err(AppError::new(
328            ErrorCode::InvalidInput,
329            format!(
330                "JSON array dataset {} exceeded max {MAX_JSON_ARRAY_BYTES} bytes while reading",
331                path.display()
332            ),
333        ));
334    }
335    Ok(bytes)
336}
337
338fn record_from_value(value: Value) -> AppResult<DatasetRecord> {
339    validate_json_record(&value)?;
340    match value {
341        Value::Object(fields) => Ok(DatasetRecord::new(fields.into_iter().collect())),
342        _ => Err(AppError::new(
343            ErrorCode::InvalidInput,
344            "dataset record must be a JSON object",
345        )),
346    }
347}
348
349#[cfg(test)]
350mod tests {
351    use futures_util::StreamExt as _;
352    use serde_json::json;
353
354    use super::*;
355
356    #[test]
357    fn private_record_parsers_reject_invalid_shapes() {
358        assert!(record_from_json_bytes(b"not-json").is_err());
359        assert!(record_from_value(json!([1, 2])).is_err());
360        assert!(parse_csv_record(b"").is_err());
361        assert!(parse_csv_record(b"\xff").is_err());
362    }
363
364    #[test]
365    fn private_bounded_readers_report_io_errors() {
366        struct FailingRead;
367
368        impl std::io::Read for FailingRead {
369            fn read(&mut self, _buf: &mut [u8]) -> std::io::Result<usize> {
370                Err(std::io::Error::other("boom"))
371            }
372        }
373
374        impl std::io::BufRead for FailingRead {
375            fn fill_buf(&mut self) -> std::io::Result<&[u8]> {
376                Err(std::io::Error::other("boom"))
377            }
378
379            fn consume(&mut self, _amt: usize) {}
380        }
381
382        let mut reader = FailingRead;
383        assert_eq!(
384            read_json_line_bounded(&mut reader).unwrap_err().code(),
385            ErrorCode::Internal
386        );
387
388        let dir = tempfile::tempdir().unwrap();
389        assert_eq!(
390            read_bounded_file(dir.path(), 16).unwrap_err().code(),
391            ErrorCode::Internal
392        );
393    }
394
395    #[test]
396    fn blocking_readers_report_missing_tokio_runtime() {
397        let reader = JsonLinesReader::new("unused.jsonl");
398        let mut stream = Box::new(reader).stream();
399        let err = futures::executor::block_on(stream.next())
400            .unwrap()
401            .unwrap_err();
402        assert_eq!(err.code(), ErrorCode::Internal);
403        assert!(err.to_string().contains("Tokio runtime"));
404    }
405
406    #[tokio::test]
407    async fn missing_record_readers_emit_errors_in_stream() {
408        let missing = std::env::temp_dir().join("rskit-dataset-missing-record-file");
409        let readers: Vec<Box<dyn DatasetReader>> = vec![
410            Box::new(CsvReader::new(&missing)),
411            Box::new(JsonLinesReader::new(&missing)),
412            Box::new(JsonArrayReader::new(&missing)),
413        ];
414
415        for reader in readers {
416            let mut stream = reader.stream();
417            assert!(stream.next().await.unwrap().is_err());
418        }
419    }
420}
rskit_dataset/record/reader.rs

rskit_dataset/record/
reader.rs
