faucet_core/driver/
source.rs

1//! The family-neutral source driver contract and its [`Source`] adapter.
2
3use crate::capabilities::Capabilities;
4use crate::check::{CheckContext, CheckReport, Probe};
5use crate::error::FaucetError;
6use crate::io::{CallContext, CallSlot};
7use crate::observability::RoundtripRecorder;
8use crate::pipeline::StreamPage;
9use crate::traits::Source;
10use async_trait::async_trait;
11use futures::StreamExt;
12use futures_core::Stream;
13use serde_json::Value;
14use std::collections::HashMap;
15use std::pin::Pin;
16use std::sync::Arc;
17
18/// A stream of pages a [`SourceDriver`] reads.
19pub type PageStream<'a> = Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>>;
20
21/// A dataset's read primitives (RFC 0014). Wrap it in a [`DriverSource`] to
22/// get a [`Source`].
23///
24/// [`read`](Self::read) streams the dataset in pages, each with the bookmark
25/// to resume after it (when there is one); every backend request goes through
26/// the [`CallSlot`] it is handed.
27///
28/// ```
29/// use faucet_core::driver::{DriverSource, PageStream, SourceDriver};
30/// use faucet_core::{CallSlot, Source, StreamPage, Value, async_trait, json};
31/// use std::collections::HashMap;
32///
33/// struct Numbers;
34///
35/// #[async_trait]
36/// impl SourceDriver for Numbers {
37///     fn connector_name(&self) -> &'static str {
38///         "numbers"
39///     }
40///     fn read<'a>(
41///         &'a self,
42///         _calls: &'a CallSlot,
43///         _context: &'a HashMap<String, Value>,
44///         _batch_size: usize,
45///     ) -> PageStream<'a> {
46///         Box::pin(faucet_core::async_stream::stream! {
47///             yield Ok(StreamPage { records: vec![json!(1), json!(2)], bookmark: Some(json!(2)) });
48///         })
49///     }
50/// }
51///
52/// # tokio::runtime::Runtime::new().unwrap().block_on(async {
53/// let source = DriverSource::new(Numbers);
54/// let (rows, bookmark) = source.fetch_all_incremental().await.unwrap();
55/// assert_eq!(rows.len(), 2);
56/// assert_eq!(bookmark, Some(json!(2)));
57/// # });
58/// ```
59#[async_trait]
60pub trait SourceDriver: Send + Sync {
61    /// The connector's name, the `connector` metric label.
62    fn connector_name(&self) -> &'static str;
63
64    /// The JSON Schema of the connector's config.
65    fn config_schema(&self) -> Value {
66        serde_json::json!({"type": "object", "properties": {}})
67    }
68
69    /// The dataset this source reads, for lineage (credential-free).
70    fn dataset_uri(&self) -> String {
71        format!("{}://unknown", self.connector_name())
72    }
73
74    /// The connector's own call defaults. The run's context is layered over
75    /// it.
76    fn call_context(&self) -> CallContext {
77        CallContext::new()
78    }
79
80    /// Stream the dataset in pages of about `batch_size` records (`0`: one
81    /// page). `context` holds the parent record of a fan-out row.
82    fn read<'a>(
83        &'a self,
84        calls: &'a CallSlot,
85        context: &'a HashMap<String, Value>,
86        batch_size: usize,
87    ) -> PageStream<'a>;
88
89    /// The state key bookmarks are stored under, when the source resumes.
90    fn state_key(&self) -> Option<String> {
91        None
92    }
93
94    /// Resume after `bookmark` on the next [`read`](Self::read).
95    async fn resume_from(&self, bookmark: Value) -> Result<(), FaucetError> {
96        let _ = bookmark;
97        Ok(())
98    }
99
100    /// Whether a resumed read replays exactly what follows the bookmark.
101    fn deterministic_replay(&self) -> bool {
102        false
103    }
104
105    /// Whether reading removes what was read (a queue).
106    fn consumes_destructively(&self) -> bool {
107        false
108    }
109
110    /// Take the run's cancellation token. Return `true` when the source then
111    /// yields what it buffered and ends, so the pipeline polls it to its end.
112    fn set_cancel(&self, cancel: crate::CancellationToken) -> bool {
113        let _ = cancel;
114        false
115    }
116
117    /// A non-mutating preflight probe. `None` (the default) reads one page.
118    async fn check(
119        &self,
120        calls: &CallSlot,
121        ctx: &CheckContext,
122    ) -> Result<Option<CheckReport>, FaucetError> {
123        let _ = (calls, ctx);
124        Ok(None)
125    }
126
127    /// Adjust the capabilities [`DriverSource`] derived. Default: unchanged.
128    fn capabilities(&self, derived: Capabilities) -> Capabilities {
129        derived
130    }
131
132    /// A buffered read with its own all-at-once path, returning the records
133    /// and the bookmark to resume after them. `None` (the default): the
134    /// buffered fetches drain [`read`](Self::read).
135    async fn fetch(
136        &self,
137        calls: &CallSlot,
138        context: &HashMap<String, Value>,
139    ) -> Option<Result<(Vec<Value>, Option<Value>), FaucetError>> {
140        let _ = (calls, context);
141        None
142    }
143
144    /// Receive the run clock (the instant `${now.*}` renders from).
145    fn set_run_clock(&self, now: chrono::DateTime<chrono::Utc>) {
146        let _ = now;
147    }
148
149    /// The run's call context, for a driver whose clients hold their own.
150    fn set_call_context(&self, ctx: CallContext) {
151        let _ = ctx;
152    }
153
154    /// The run's round-trip recorder, for a driver whose clients hold their
155    /// own.
156    fn set_roundtrip_recorder(&self, recorder: Arc<RoundtripRecorder>) {
157        let _ = recorder;
158    }
159
160    /// Wire formats [`read_native`](Self::read_native) can emit (#633).
161    fn native_output_formats(&self) -> &'static [crate::native::NativeFormat] {
162        &[]
163    }
164
165    /// Stream the dataset as raw bytes in a negotiated `format`.
166    fn read_native<'a>(
167        &'a self,
168        calls: &'a CallSlot,
169        context: &'a HashMap<String, Value>,
170        format: crate::native::NativeFormat,
171        batch_size: usize,
172    ) -> NativeStream<'a> {
173        let _ = (calls, context, format, batch_size);
174        let name = self.connector_name();
175        let err: Result<crate::native::NativeBatch, FaucetError> = Err(FaucetError::Source(
176            format!("source '{name}' does not support native byte streaming (stream_native)"),
177        ));
178        Box::pin(futures::stream::once(async move { err }))
179    }
180
181    /// Whether [`read_columnar`](Self::read_columnar) is implemented.
182    #[cfg(feature = "arrow")]
183    fn supports_columnar(&self) -> bool {
184        false
185    }
186
187    /// Stream the dataset as Arrow batches.
188    #[cfg(feature = "arrow")]
189    fn read_columnar<'a>(
190        &'a self,
191        calls: &'a CallSlot,
192        context: &'a HashMap<String, Value>,
193        batch_size: usize,
194    ) -> ColumnarStream<'a> {
195        let _ = (calls, context, batch_size);
196        let name = self.connector_name();
197        let err: Result<crate::columnar::ColumnarPage, FaucetError> = Err(FaucetError::Source(
198            format!("source '{name}' does not support columnar streaming (stream_batches)"),
199        ));
200        Box::pin(futures::stream::once(async move { err }))
201    }
202
203    /// Whether [`discover`](Self::discover) is implemented.
204    fn supports_discover(&self) -> bool {
205        false
206    }
207
208    /// Enumerate the datasets behind the connection (read-only, cheap).
209    async fn discover(
210        &self,
211        calls: &CallSlot,
212    ) -> Result<Vec<crate::discover::DatasetDescriptor>, FaucetError> {
213        let _ = calls;
214        Err(FaucetError::Source(format!(
215            "source '{}' does not support dataset discovery",
216            self.connector_name()
217        )))
218    }
219}
220
221/// A stream of raw byte batches a [`SourceDriver`] reads.
222pub type NativeStream<'a> =
223    Pin<Box<dyn Stream<Item = Result<crate::native::NativeBatch, FaucetError>> + Send + 'a>>;
224
225/// A stream of Arrow pages a [`SourceDriver`] reads.
226#[cfg(feature = "arrow")]
227pub type ColumnarStream<'a> =
228    Pin<Box<dyn Stream<Item = Result<crate::columnar::ColumnarPage, FaucetError>> + Send + 'a>>;
229
230/// Split a buffered read into pages of `batch_size` records (`0`: one page),
231/// the bookmark on the last page — the shape a source without a native paging
232/// primitive streams in. An empty read yields one empty page only when it
233/// carries a bookmark.
234///
235/// ```
236/// use faucet_core::driver::pages_of;
237/// use faucet_core::json;
238///
239/// let pages: Vec<_> = pages_of(vec![json!(1), json!(2), json!(3)], Some(json!(3)), 2).collect();
240/// assert_eq!(pages.len(), 2);
241/// assert_eq!(pages[0].bookmark, None);
242/// assert_eq!(pages[1].bookmark, Some(json!(3)));
243/// ```
244pub fn pages_of(
245    records: Vec<Value>,
246    bookmark: Option<Value>,
247    batch_size: usize,
248) -> impl Iterator<Item = StreamPage> + Send {
249    let chunk = if batch_size == 0 {
250        usize::MAX
251    } else {
252        batch_size
253    };
254    let total = records.len();
255    let empty = (total == 0 && bookmark.is_some()).then(|| StreamPage {
256        records: Vec::new(),
257        bookmark: bookmark.clone(),
258    });
259    let mut iter = records.into_iter();
260    let mut consumed = 0usize;
261    let pages = std::iter::from_fn(move || {
262        let batch: Vec<Value> = iter.by_ref().take(chunk).collect();
263        if batch.is_empty() {
264            return None;
265        }
266        consumed += batch.len();
267        Some(StreamPage {
268            bookmark: if consumed >= total {
269                bookmark.clone()
270            } else {
271                None
272            },
273            records: batch,
274        })
275    });
276    empty.into_iter().chain(pages)
277}
278
279/// A [`SourceDriver`] as a [`Source`] (RFC 0014): every fetch is the driver's
280/// page stream, and the run's call context reaches every read.
281pub struct DriverSource<D> {
282    driver: D,
283    calls: CallSlot,
284}
285
286impl<D: SourceDriver> DriverSource<D> {
287    /// The source over `driver`.
288    pub fn new(driver: D) -> Self {
289        let calls = CallSlot::new(driver.call_context());
290        Self { driver, calls }
291    }
292
293    /// The driver.
294    pub fn driver(&self) -> &D {
295        &self.driver
296    }
297
298    /// The driver, mutably.
299    pub fn driver_mut(&mut self) -> &mut D {
300        &mut self.driver
301    }
302
303    /// The call slot the reads receive.
304    pub fn calls(&self) -> &CallSlot {
305        &self.calls
306    }
307
308    /// Unwrap the driver.
309    pub fn into_driver(self) -> D {
310        self.driver
311    }
312
313    async fn drain(
314        &self,
315        context: &HashMap<String, Value>,
316    ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
317        if let Some(fetched) = self.driver.fetch(&self.calls, context).await {
318            return fetched;
319        }
320        let mut pages = self.driver.read(&self.calls, context, 0);
321        let mut rows = Vec::new();
322        let mut bookmark = None;
323        while let Some(page) = pages.next().await {
324            let page = page?;
325            rows.extend(page.records);
326            if page.bookmark.is_some() {
327                bookmark = page.bookmark;
328            }
329        }
330        Ok((rows, bookmark))
331    }
332}
333
334#[async_trait]
335impl<D: SourceDriver> Source for DriverSource<D> {
336    async fn fetch_with_context(
337        &self,
338        context: &HashMap<String, Value>,
339    ) -> Result<Vec<Value>, FaucetError> {
340        Ok(self.drain(context).await?.0)
341    }
342
343    async fn fetch_with_context_incremental(
344        &self,
345        context: &HashMap<String, Value>,
346    ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
347        self.drain(context).await
348    }
349
350    fn stream_pages<'a>(
351        &'a self,
352        context: &'a HashMap<String, Value>,
353        batch_size: usize,
354    ) -> PageStream<'a> {
355        self.driver.read(&self.calls, context, batch_size)
356    }
357
358    fn config_schema(&self) -> Value {
359        self.driver.config_schema()
360    }
361
362    fn state_key(&self) -> Option<String> {
363        self.driver.state_key()
364    }
365
366    async fn apply_start_bookmark(&self, bookmark: Value) -> Result<(), FaucetError> {
367        self.driver.resume_from(bookmark).await
368    }
369
370    fn supports_exactly_once(&self) -> bool {
371        self.driver.deterministic_replay()
372    }
373
374    fn consumes_destructively(&self) -> bool {
375        self.driver.consumes_destructively()
376    }
377
378    fn capabilities(&self) -> Capabilities {
379        self.driver.capabilities(Capabilities::of_source(self))
380    }
381
382    fn connector_name(&self) -> &'static str {
383        self.driver.connector_name()
384    }
385
386    fn set_roundtrip_recorder(&self, recorder: Arc<RoundtripRecorder>) {
387        self.calls.install_recorder(Arc::clone(&recorder));
388        self.driver.set_roundtrip_recorder(recorder);
389    }
390
391    fn set_call_context(&self, ctx: CallContext) {
392        self.calls.install(ctx.clone());
393        self.driver.set_call_context(ctx);
394    }
395
396    fn set_run_clock(&self, now: chrono::DateTime<chrono::Utc>) {
397        self.driver.set_run_clock(now);
398    }
399
400    fn native_output_formats(&self) -> &'static [crate::native::NativeFormat] {
401        self.driver.native_output_formats()
402    }
403
404    fn stream_native<'a>(
405        &'a self,
406        context: &'a HashMap<String, Value>,
407        format: crate::native::NativeFormat,
408        batch_size: usize,
409    ) -> NativeStream<'a> {
410        self.driver
411            .read_native(&self.calls, context, format, batch_size)
412    }
413
414    #[cfg(feature = "arrow")]
415    fn supports_columnar(&self) -> bool {
416        self.driver.supports_columnar()
417    }
418
419    #[cfg(feature = "arrow")]
420    fn stream_batches<'a>(
421        &'a self,
422        context: &'a HashMap<String, Value>,
423        batch_size: usize,
424    ) -> ColumnarStream<'a> {
425        self.driver.read_columnar(&self.calls, context, batch_size)
426    }
427
428    fn supports_discover(&self) -> bool {
429        self.driver.supports_discover()
430    }
431
432    async fn discover(&self) -> Result<Vec<crate::discover::DatasetDescriptor>, FaucetError> {
433        self.driver.discover(&self.calls).await
434    }
435
436    fn set_cancel(&self, cancel: crate::CancellationToken) -> bool {
437        self.driver.set_cancel(cancel)
438    }
439
440    fn dataset_uri(&self) -> String {
441        self.driver.dataset_uri()
442    }
443
444    async fn check(&self, ctx: &CheckContext) -> Result<CheckReport, FaucetError> {
445        if let Some(report) = self.driver.check(&self.calls, ctx).await? {
446            return Ok(report);
447        }
448        let empty = HashMap::new();
449        let start = std::time::Instant::now();
450        let mut pages = self.driver.read(&self.calls, &empty, 1);
451        let probe = match tokio::time::timeout(ctx.timeout, pages.next()).await {
452            Err(_) => Probe::fail("read", start.elapsed(), "timed out fetching first page"),
453            Ok(None) | Ok(Some(Ok(_))) => Probe::pass("read", start.elapsed()),
454            Ok(Some(Err(e))) => Probe::fail("read", start.elapsed(), e.to_string()),
455        };
456        Ok(CheckReport::single(probe))
457    }
458}
faucet_core/driver/source.rs

faucet_core/driver/
source.rs
