surrealml_core/storage/
stream_adapter.rs

1//! Stream adapter for file system
2use std::error::Error;
3use std::fs::File;
4use std::io::Read;
5use std::pin::Pin;
6
7use bytes::Bytes;
8use futures::stream::Stream;
9use futures::task::{Context, Poll};
10// For hyper 1.x compatibility
11use hyper::body::Frame;
12
13use crate::errors::error::{SurrealError, SurrealErrorStatus};
14use crate::safe_eject;
15
16/// Stream adapter for file system.
17///
18/// # Arguments
19/// * `chunk_size` - The size of the chunks to read from the file.
20/// * `file_pointer` - The pointer to the file to be streamed
21pub struct StreamAdapter {
22	chunk_size: usize,
23	file_pointer: File,
24}
25
26impl StreamAdapter {
27	/// Creates a new `StreamAdapter` struct.
28	///
29	/// # Arguments
30	/// * `chunk_size` - The size of the chunks to read from the file.
31	/// * `file_path` - The path to the file to be streamed
32	///
33	/// # Returns
34	/// A new `StreamAdapter` struct.
35	pub fn new(chunk_size: usize, file_path: String) -> Result<Self, SurrealError> {
36		let file_pointer = safe_eject!(File::open(file_path), SurrealErrorStatus::NotFound);
37		Ok(StreamAdapter {
38			chunk_size,
39			file_pointer,
40		})
41	}
42}
43
44impl Stream for StreamAdapter {
45	type Item = Result<Frame<Bytes>, Box<dyn Error + Send + Sync>>;
46
47	/// Polls the next chunk from the file.
48	///
49	/// # Arguments
50	/// * `self` - The `StreamAdapter` struct.
51	/// * `cx` - The context of the task to enable the task to be woken up and polled again using
52	///   the waker.
53	///
54	/// # Returns
55	/// A poll containing the next chunk from the file.
56	fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
57		let mut buffer = vec![0u8; self.chunk_size];
58		let bytes_read = self.file_pointer.read(&mut buffer)?;
59
60		buffer.truncate(bytes_read);
61		if buffer.is_empty() {
62			return Poll::Ready(None);
63		}
64		Poll::Ready(Some(Ok(Frame::data(buffer.into()))))
65	}
66}
surrealml_core/storage/stream_adapter.rs

surrealml_core/storage/
stream_adapter.rs
