faucet_source_sftp/
stream.rs

1//! SFTP source stream executor.
2//!
3//! The one place that performs SFTP I/O. Construction is lazy — [`SftpSource::new`]
4//! only stores (and validates) the config; the SSH transport is opened on the
5//! first `fetch_*` / `stream_pages` call, so an unreachable host surfaces as a
6//! typed error on first poll rather than at construction time.
7
8use crate::config::{SftpFormat, SftpSourceConfig, glob_match};
9use async_trait::async_trait;
10use faucet_common_sftp::{SftpSession, connect};
11use faucet_core::{FaucetError, Stream, StreamPage};
12use serde_json::Value;
13use std::collections::HashMap;
14use std::pin::Pin;
15use tokio::io::AsyncBufReadExt;
16
17/// An SFTP source that lists and reads remote files.
18pub struct SftpSource {
19    config: SftpSourceConfig,
20}
21
22impl SftpSource {
23    /// Create a new SFTP source. Lazy: performs no I/O and never connects here.
24    /// The batch size is validated up front so a bad config fails fast.
25    pub fn new(config: SftpSourceConfig) -> Result<Self, FaucetError> {
26        faucet_core::validate_batch_size(config.batch_size)?;
27        Ok(Self { config })
28    }
29
30    /// Resolve the effective remote path, substituting parent-context tokens
31    /// when a non-empty context is supplied (parent/child matrix runs).
32    fn effective_path(&self, context: &HashMap<String, Value>) -> String {
33        if context.is_empty() {
34            self.config.path.clone()
35        } else {
36            faucet_core::util::substitute_context(&self.config.path, context)
37        }
38    }
39
40    /// List the files to read for `path`: the directory's regular files
41    /// (glob-filtered, sorted for a deterministic order) when `path` is a
42    /// directory, or `[path]` when it is a single file.
43    async fn resolve_files(
44        &self,
45        sftp: &SftpSession,
46        path: &str,
47    ) -> Result<Vec<String>, FaucetError> {
48        let meta = sftp
49            .metadata(path)
50            .await
51            .map_err(|e| FaucetError::Source(format!("SFTP stat '{path}' failed: {e}")))?;
52
53        if !meta.file_type().is_dir() {
54            return Ok(vec![path.to_string()]);
55        }
56
57        let entries = sftp
58            .read_dir(path)
59            .await
60            .map_err(|e| FaucetError::Source(format!("SFTP read_dir '{path}' failed: {e}")))?;
61
62        let mut files = Vec::new();
63        for entry in entries {
64            if !entry.file_type().is_file() {
65                continue;
66            }
67            let name = entry.file_name();
68            if let Some(pattern) = &self.config.glob
69                && !glob_match(pattern, &name)
70            {
71                continue;
72            }
73            files.push(entry.path());
74        }
75        files.sort();
76        Ok(files)
77    }
78
79    /// Read a whole remote file into a UTF-8 `String` (used by `raw_text` and
80    /// `json_array`, which cannot be parsed incrementally).
81    async fn read_file_text(sftp: &SftpSession, path: &str) -> Result<String, FaucetError> {
82        let bytes = sftp
83            .read(path)
84            .await
85            .map_err(|e| FaucetError::Source(format!("SFTP read '{path}' failed: {e}")))?;
86        String::from_utf8(bytes)
87            .map_err(|e| FaucetError::Source(format!("SFTP file '{path}' is not valid UTF-8: {e}")))
88    }
89}
90
91#[async_trait]
92impl faucet_core::Source for SftpSource {
93    async fn fetch_with_context(
94        &self,
95        context: &HashMap<String, Value>,
96    ) -> Result<Vec<Value>, FaucetError> {
97        use futures::StreamExt;
98        let mut all = Vec::new();
99        let mut pages = self.stream_pages(context, self.config.batch_size);
100        while let Some(page) = pages.next().await {
101            all.extend(page?.records);
102        }
103        Ok(all)
104    }
105
106    /// Stream records from the resolved remote files without buffering the full
107    /// scan. Each emitted [`StreamPage`] holds up to
108    /// [`SftpSourceConfig::batch_size`] records.
109    ///
110    /// - `jsonl` / `raw_text`: files are decoded incrementally (line-by-line
111    ///   for JSONL; one record per file for raw text) so client memory is
112    ///   bounded regardless of file size. Multi-file scans flatten — a page
113    ///   may carry records from more than one file.
114    /// - `json_array`: each file is buffered fully, then its records chunked.
115    ///
116    /// `batch_size = 0` emits one page per file. Every page carries
117    /// `bookmark: None` — the SFTP source is not resumable.
118    fn stream_pages<'a>(
119        &'a self,
120        context: &'a HashMap<String, Value>,
121        _batch_size: usize,
122    ) -> Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>> {
123        let batch_size = self.config.batch_size;
124
125        Box::pin(async_stream::try_stream! {
126            let sftp = connect(&self.config.connection).await?;
127            let path = self.effective_path(context);
128            let files = self.resolve_files(&sftp, &path).await?;
129            tracing::info!(path = %path, files = files.len(), "SFTP source listed files");
130
131            let chunk = if batch_size == 0 { usize::MAX } else { batch_size };
132            let initial_capacity = if batch_size == 0 { 1024 } else { batch_size };
133            let mut buffer: Vec<Value> = Vec::with_capacity(initial_capacity);
134            let mut total = 0usize;
135
136            for file in &files {
137                match self.config.format {
138                    SftpFormat::Jsonl => {
139                        let handle = sftp.open(file.as_str()).await.map_err(|e| {
140                            FaucetError::Source(format!("SFTP open '{file}' failed: {e}"))
141                        })?;
142                        let reader = tokio::io::BufReader::new(handle);
143                        let mut lines = reader.lines();
144                        let mut line_num = 0usize;
145                        while let Some(line) = lines.next_line().await.map_err(|e| {
146                            FaucetError::Source(format!("SFTP read '{file}' failed: {e}"))
147                        })? {
148                            line_num += 1;
149                            let trimmed = line.trim();
150                            if trimmed.is_empty() {
151                                continue;
152                            }
153                            let value: Value = serde_json::from_str(trimmed).map_err(|e| {
154                                FaucetError::Source(format!(
155                                    "SFTP JSON parse error in '{file}' at line {line_num}: {e}"
156                                ))
157                            })?;
158                            buffer.push(value);
159                            if batch_size != 0 && buffer.len() >= chunk {
160                                let page = std::mem::replace(
161                                    &mut buffer,
162                                    Vec::with_capacity(initial_capacity),
163                                );
164                                total += page.len();
165                                yield StreamPage { records: page, bookmark: None };
166                            }
167                        }
168                        if batch_size == 0 && !buffer.is_empty() {
169                            let page = std::mem::take(&mut buffer);
170                            total += page.len();
171                            yield StreamPage { records: page, bookmark: None };
172                        }
173                    }
174                    SftpFormat::RawText => {
175                        let text = Self::read_file_text(&sftp, file).await?;
176                        buffer.push(serde_json::json!({ "path": file, "content": text }));
177                        if batch_size == 0 {
178                            let page = std::mem::take(&mut buffer);
179                            total += page.len();
180                            yield StreamPage { records: page, bookmark: None };
181                        } else if buffer.len() >= chunk {
182                            let page = std::mem::replace(
183                                &mut buffer,
184                                Vec::with_capacity(initial_capacity),
185                            );
186                            total += page.len();
187                            yield StreamPage { records: page, bookmark: None };
188                        }
189                    }
190                    SftpFormat::JsonArray => {
191                        let text = Self::read_file_text(&sftp, file).await?;
192                        let value: Value = serde_json::from_str(&text).map_err(|e| {
193                            FaucetError::Source(format!("SFTP JSON parse error in '{file}': {e}"))
194                        })?;
195                        let array = match value {
196                            Value::Array(arr) => arr,
197                            other => Err(FaucetError::Source(format!(
198                                "SFTP expected JSON array in '{file}', got {}",
199                                value_type_name(&other)
200                            )))?,
201                        };
202                        if batch_size == 0 {
203                            if !buffer.is_empty() {
204                                let page = std::mem::take(&mut buffer);
205                                total += page.len();
206                                yield StreamPage { records: page, bookmark: None };
207                            }
208                            total += array.len();
209                            yield StreamPage { records: array, bookmark: None };
210                        } else {
211                            for record in array {
212                                buffer.push(record);
213                                if buffer.len() >= chunk {
214                                    let page = std::mem::replace(
215                                        &mut buffer,
216                                        Vec::with_capacity(initial_capacity),
217                                    );
218                                    total += page.len();
219                                    yield StreamPage { records: page, bookmark: None };
220                                }
221                            }
222                        }
223                    }
224                }
225            }
226
227            if !buffer.is_empty() {
228                let page = std::mem::take(&mut buffer);
229                total += page.len();
230                yield StreamPage { records: page, bookmark: None };
231            }
232
233            tracing::info!(total_records = total, files = files.len(), "SFTP source stream complete");
234        })
235    }
236
237    fn config_schema(&self) -> Value {
238        serde_json::to_value(faucet_core::schema_for!(SftpSourceConfig))
239            .expect("schema serialization")
240    }
241
242    fn connector_name(&self) -> &'static str {
243        "sftp"
244    }
245
246    fn dataset_uri(&self) -> String {
247        format!(
248            "sftp://{}:{}/{}",
249            self.config.connection.host,
250            self.config.connection.port,
251            self.config.path.trim_start_matches('/')
252        )
253    }
254}
255
256/// Return a human-readable name for a JSON value type.
257fn value_type_name(v: &Value) -> &'static str {
258    match v {
259        Value::Null => "null",
260        Value::Bool(_) => "boolean",
261        Value::Number(_) => "number",
262        Value::String(_) => "string",
263        Value::Array(_) => "array",
264        Value::Object(_) => "object",
265    }
266}
267
268#[cfg(test)]
269mod tests {
270    use super::*;
271    use faucet_common_sftp::SftpConnectionConfig;
272    use faucet_core::Source;
273
274    fn cfg() -> SftpSourceConfig {
275        SftpSourceConfig::new(SftpConnectionConfig::with_password("h", "u", "p"), "/data")
276    }
277
278    #[test]
279    fn new_is_lazy_and_validates_batch_size() {
280        // Valid config constructs with no I/O.
281        assert!(SftpSource::new(cfg()).is_ok());
282        // Out-of-range batch size is rejected up front.
283        let bad = cfg().with_batch_size(faucet_core::MAX_BATCH_SIZE + 1);
284        assert!(matches!(SftpSource::new(bad), Err(FaucetError::Config(_))));
285    }
286
287    #[test]
288    fn connector_name_is_sftp() {
289        let src = SftpSource::new(cfg()).unwrap();
290        assert_eq!(src.connector_name(), "sftp");
291    }
292
293    #[test]
294    fn dataset_uri_has_no_credentials() {
295        let src = SftpSource::new(cfg()).unwrap();
296        let uri = src.dataset_uri();
297        assert_eq!(uri, "sftp://h:22/data");
298        assert!(!uri.contains('p') || !uri.contains("password"));
299    }
300
301    #[test]
302    fn config_schema_is_valid_object() {
303        let src = SftpSource::new(cfg()).unwrap();
304        let schema = src.config_schema();
305        assert!(schema.is_object());
306        assert!(schema.get("properties").is_some());
307    }
308}
faucet_source_sftp/stream.rs

faucet_source_sftp/
stream.rs
