rskit_storage/
sink.rs

1//! File sink — destination for file output.
2
3use std::path::PathBuf;
4
5use bytes::Bytes;
6use rskit_errors::{AppError, AppResult, ErrorCode};
7use rskit_fs::async_io::file;
8
9use crate::{FileSource, TempFile};
10
11/// Destination for file output.
12pub enum FileSink {
13    /// Write to a local path.
14    Path(PathBuf),
15    /// Write to a managed temp file (returned to caller).
16    Temp,
17    /// Write to an in-memory buffer.
18    Memory,
19}
20
21impl FileSink {
22    /// Create a writer handle for this sink.
23    pub async fn writer(&self) -> AppResult<FileWriter> {
24        match self {
25            Self::Path(p) => {
26                file::create_parent_dir(p).await?;
27                Ok(FileWriter {
28                    inner: WriterInner::Path(p.clone()),
29                    buffer: Vec::new(),
30                })
31            }
32            Self::Temp => Ok(FileWriter {
33                inner: WriterInner::Temp,
34                buffer: Vec::new(),
35            }),
36            Self::Memory => Ok(FileWriter {
37                inner: WriterInner::Memory,
38                buffer: Vec::new(),
39            }),
40        }
41    }
42}
43
44impl std::fmt::Debug for FileSink {
45    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
46        match self {
47            Self::Path(p) => f.debug_tuple("Path").field(p).finish(),
48            Self::Temp => write!(f, "Temp"),
49            Self::Memory => write!(f, "Memory"),
50        }
51    }
52}
53
54enum WriterInner {
55    Path(PathBuf),
56    Temp,
57    Memory,
58}
59
60/// Handle for writing output. Finalize to get the resulting [`FileSource`].
61pub struct FileWriter {
62    inner: WriterInner,
63    buffer: Vec<u8>,
64}
65
66impl FileWriter {
67    /// Write all bytes to the output.
68    pub async fn write_all(&mut self, data: &[u8]) -> AppResult<()> {
69        self.buffer.extend_from_slice(data);
70        Ok(())
71    }
72
73    /// Write a stream of bytes to the output.
74    pub async fn write_stream(
75        &mut self,
76        mut stream: impl futures::Stream<Item = AppResult<Bytes>> + Unpin,
77    ) -> AppResult<()> {
78        use futures::StreamExt;
79        while let Some(chunk) = stream.next().await {
80            let chunk = chunk?;
81            self.buffer.extend_from_slice(&chunk);
82        }
83        Ok(())
84    }
85
86    /// Finalize the writer and return the resulting [`FileSource`].
87    pub async fn finalize(self) -> AppResult<FileSource> {
88        match self.inner {
89            WriterInner::Path(p) => {
90                file::write(&p, &self.buffer).await.map_err(|e| {
91                    AppError::new(
92                        ErrorCode::Internal,
93                        format!("failed to write {}: {e}", p.display()),
94                    )
95                })?;
96                Ok(FileSource::Path(p))
97            }
98            WriterInner::Temp => {
99                let tmp = TempFile::new()?;
100                file::write(tmp.path(), &self.buffer).await.map_err(|e| {
101                    AppError::new(
102                        ErrorCode::Internal,
103                        format!("failed to write temp file: {e}"),
104                    )
105                })?;
106                Ok(FileSource::Temp(tmp))
107            }
108            WriterInner::Memory => Ok(FileSource::Bytes(Bytes::from(self.buffer))),
109        }
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use bytes::Bytes;
116    use futures::stream;
117
118    use super::*;
119
120    #[test]
121    fn file_sink_debug_describes_destination_without_side_effects() {
122        assert_eq!(
123            format!("{:?}", FileSink::Path(PathBuf::from("out.bin"))),
124            "Path(\"out.bin\")"
125        );
126        assert_eq!(format!("{:?}", FileSink::Temp), "Temp");
127        assert_eq!(format!("{:?}", FileSink::Memory), "Memory");
128    }
129
130    #[tokio::test]
131    async fn memory_writer_collects_direct_and_streamed_bytes() {
132        let mut writer = FileSink::Memory.writer().await.unwrap();
133        writer.write_all(b"hello ").await.unwrap();
134        writer
135            .write_stream(stream::iter([
136                Ok(Bytes::from_static(b"stream")),
137                Ok(Bytes::from_static(b"!")),
138            ]))
139            .await
140            .unwrap();
141
142        let source = writer.finalize().await.unwrap();
143        assert_eq!(
144            source.read_all().await.unwrap(),
145            Bytes::from_static(b"hello stream!")
146        );
147    }
148
149    #[tokio::test]
150    async fn path_and_temp_writers_materialize_file_sources() {
151        let dir = tempfile::tempdir().unwrap();
152        let output = dir.path().join("nested").join("file.txt");
153        let mut path_writer = FileSink::Path(output.clone()).writer().await.unwrap();
154        path_writer.write_all(b"path").await.unwrap();
155        let path_source = path_writer.finalize().await.unwrap();
156        assert!(matches!(path_source, FileSource::Path(_)));
157        assert_eq!(tokio::fs::read(&output).await.unwrap(), b"path");
158
159        let mut temp_writer = FileSink::Temp.writer().await.unwrap();
160        temp_writer.write_all(b"temp").await.unwrap();
161        let temp_source = temp_writer.finalize().await.unwrap();
162        assert_eq!(
163            temp_source.read_all().await.unwrap(),
164            Bytes::from_static(b"temp")
165        );
166    }
167
168    #[tokio::test]
169    async fn writer_stream_forwards_chunk_errors() {
170        let mut writer = FileSink::Memory.writer().await.unwrap();
171        let error = AppError::new(ErrorCode::Internal, "chunk failed");
172        let result = writer
173            .write_stream(stream::iter([Err::<Bytes, _>(error)]))
174            .await;
175        assert_eq!(result.unwrap_err().code(), ErrorCode::Internal);
176    }
177}
rskit_storage/sink.rs

rskit_storage/
sink.rs
