faucet_sink_sftp/
sink.rs

1//! SFTP sink executor.
2//!
3//! Writes each `write_batch` chunk as a JSON Lines object under a remote
4//! directory. Writes are **atomic**: each object is uploaded to a hidden
5//! temporary name and then renamed to its final name, so a consumer watching
6//! the directory never observes a partially-written file. Construction is lazy
7//! — [`SftpSink::new`] performs no I/O; the SSH transport is opened on the
8//! first `write_batch` and reused for the life of the sink.
9
10use crate::config::SftpSinkConfig;
11use async_trait::async_trait;
12use faucet_common_sftp::{SftpSession, connect};
13use faucet_core::FaucetError;
14use russh_sftp::protocol::OpenFlags;
15use serde_json::Value;
16use tokio::io::AsyncWriteExt;
17use tokio::sync::Mutex;
18
19/// A sink that writes JSON records to an SFTP server as JSON Lines objects.
20pub struct SftpSink {
21    config: SftpSinkConfig,
22    /// Reused SFTP session, opened on first write. Behind a `Mutex` so writes
23    /// share one SSH connection instead of reconnecting per page.
24    session: Mutex<Option<SftpSession>>,
25}
26
27impl SftpSink {
28    /// Create a new SFTP sink. Lazy: performs no I/O and never connects here.
29    /// The batch size is validated up front so a bad config fails fast.
30    pub fn new(config: SftpSinkConfig) -> Result<Self, FaucetError> {
31        faucet_core::validate_batch_size(config.batch_size)?;
32        Ok(Self {
33            config,
34            session: Mutex::new(None),
35        })
36    }
37
38    /// Serialize a slice of records as JSON Lines bytes.
39    fn serialize_jsonl(records: &[Value]) -> Result<Vec<u8>, FaucetError> {
40        let mut buf: Vec<u8> = Vec::new();
41        for record in records {
42            let line = serde_json::to_vec(record)
43                .map_err(|e| FaucetError::Sink(format!("JSON serialization failed: {e}")))?;
44            buf.extend_from_slice(&line);
45            buf.push(b'\n');
46        }
47        Ok(buf)
48    }
49
50    /// Join the configured directory prefix with a file name using POSIX `/`.
51    fn join_path(&self, name: &str) -> String {
52        let dir = &self.config.path;
53        if dir.is_empty() {
54            name.to_string()
55        } else if dir.ends_with('/') {
56            format!("{dir}{name}")
57        } else {
58            format!("{dir}/{name}")
59        }
60    }
61
62    /// Generate the final object key: `{path}/{uuid}{ext}`.
63    fn final_key(&self) -> String {
64        let id = uuid::Uuid::new_v4();
65        self.join_path(&format!("{id}{}", self.config.file_extension))
66    }
67
68    /// Upload `body` to `final_key` atomically: write a temporary object and
69    /// rename it into place, so consumers never see a partial file.
70    async fn upload_atomic(
71        sftp: &SftpSession,
72        final_key: &str,
73        body: &[u8],
74    ) -> Result<(), FaucetError> {
75        let temp_key = format!("{final_key}.tmp");
76
77        let mut file = sftp
78            .open_with_flags(
79                temp_key.as_str(),
80                OpenFlags::CREATE | OpenFlags::WRITE | OpenFlags::TRUNCATE,
81            )
82            .await
83            .map_err(|e| {
84                FaucetError::Sink(format!("SFTP open '{temp_key}' for write failed: {e}"))
85            })?;
86
87        file.write_all(body)
88            .await
89            .map_err(|e| FaucetError::Sink(format!("SFTP write to '{temp_key}' failed: {e}")))?;
90        file.flush()
91            .await
92            .map_err(|e| FaucetError::Sink(format!("SFTP flush of '{temp_key}' failed: {e}")))?;
93        file.shutdown()
94            .await
95            .map_err(|e| FaucetError::Sink(format!("SFTP close of '{temp_key}' failed: {e}")))?;
96
97        sftp.rename(temp_key.as_str(), final_key)
98            .await
99            .map_err(|e| {
100                FaucetError::Sink(format!(
101                    "SFTP rename '{temp_key}' -> '{final_key}' failed: {e}"
102                ))
103            })?;
104
105        tracing::debug!(key = %final_key, "wrote SFTP object");
106        Ok(())
107    }
108}
109
110#[async_trait]
111impl faucet_core::Sink for SftpSink {
112    async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
113        if records.is_empty() {
114            return Ok(0);
115        }
116
117        let mut guard = self.session.lock().await;
118        if guard.is_none() {
119            let sftp = connect(&self.config.connection).await?;
120            // Best-effort: ensure the target directory exists. Ignore errors
121            // (it usually already exists; a real permission/path problem
122            // surfaces on the first write with a clear message).
123            if let Err(e) = sftp.create_dir(self.config.path.as_str()).await {
124                tracing::debug!(path = %self.config.path, error = %e, "SFTP create_dir (best-effort)");
125            }
126            *guard = Some(sftp);
127        }
128        let sftp = guard.as_ref().expect("session initialized above");
129
130        let chunks: Vec<&[Value]> = if self.config.batch_size == 0 {
131            vec![records]
132        } else {
133            records.chunks(self.config.batch_size).collect()
134        };
135
136        let files = chunks.len();
137        for chunk in chunks {
138            let body = Self::serialize_jsonl(chunk)?;
139            let key = self.final_key();
140            Self::upload_atomic(sftp, &key, &body).await?;
141        }
142
143        tracing::info!(records = records.len(), files, "SFTP batch write complete");
144        Ok(records.len())
145    }
146
147    fn config_schema(&self) -> Value {
148        serde_json::to_value(faucet_core::schema_for!(SftpSinkConfig))
149            .expect("schema serialization")
150    }
151
152    fn connector_name(&self) -> &'static str {
153        "sftp"
154    }
155
156    fn dataset_uri(&self) -> String {
157        format!(
158            "sftp://{}:{}/{}",
159            self.config.connection.host,
160            self.config.connection.port,
161            self.config.path.trim_start_matches('/')
162        )
163    }
164}
165
166#[cfg(test)]
167mod tests {
168    use super::*;
169    use faucet_common_sftp::SftpConnectionConfig;
170    use faucet_core::Sink;
171
172    fn cfg() -> SftpSinkConfig {
173        SftpSinkConfig::new(SftpConnectionConfig::with_password("h", "u", "p"), "/out")
174    }
175
176    #[test]
177    fn new_is_lazy_and_validates_batch_size() {
178        assert!(SftpSink::new(cfg()).is_ok());
179        let bad = cfg().with_batch_size(faucet_core::MAX_BATCH_SIZE + 1);
180        assert!(matches!(SftpSink::new(bad), Err(FaucetError::Config(_))));
181    }
182
183    #[test]
184    fn serialize_jsonl_is_newline_delimited() {
185        let records = vec![
186            serde_json::json!({"id": 1, "name": "Alice"}),
187            serde_json::json!({"id": 2, "name": "Bob"}),
188        ];
189        let bytes = SftpSink::serialize_jsonl(&records).unwrap();
190        let text = String::from_utf8(bytes).unwrap();
191        let lines: Vec<&str> = text.trim().split('\n').collect();
192        assert_eq!(lines.len(), 2);
193        let first: Value = serde_json::from_str(lines[0]).unwrap();
194        assert_eq!(first["id"], 1);
195    }
196
197    #[test]
198    fn serialize_jsonl_empty() {
199        assert!(SftpSink::serialize_jsonl(&[]).unwrap().is_empty());
200    }
201
202    #[test]
203    fn join_path_handles_trailing_slash() {
204        let sink = SftpSink::new(cfg()).unwrap();
205        assert_eq!(sink.join_path("f.jsonl"), "/out/f.jsonl");
206
207        let sink2 = SftpSink::new(SftpSinkConfig::new(
208            SftpConnectionConfig::with_password("h", "u", "p"),
209            "/out/",
210        ))
211        .unwrap();
212        assert_eq!(sink2.join_path("f.jsonl"), "/out/f.jsonl");
213    }
214
215    #[test]
216    fn final_key_uses_prefix_and_extension() {
217        let sink = SftpSink::new(cfg()).unwrap();
218        let key = sink.final_key();
219        assert!(key.starts_with("/out/"), "got {key}");
220        assert!(key.ends_with(".jsonl"), "got {key}");
221    }
222
223    #[test]
224    fn connector_name_is_sftp() {
225        let sink = SftpSink::new(cfg()).unwrap();
226        assert_eq!(sink.connector_name(), "sftp");
227    }
228
229    #[test]
230    fn dataset_uri_has_no_credentials() {
231        let sink = SftpSink::new(cfg()).unwrap();
232        assert_eq!(sink.dataset_uri(), "sftp://h:22/out");
233    }
234
235    #[test]
236    fn append_only_capabilities() {
237        let sink = SftpSink::new(cfg()).unwrap();
238        assert!(!sink.supports_idempotent_writes());
239        assert!(!sink.dedups_by_key());
240        assert!(
241            sink.supported_write_modes()
242                .contains(&faucet_core::write_mode::WriteMode::Append)
243        );
244    }
245}
faucet_sink_sftp/sink.rs

faucet_sink_sftp/
sink.rs
