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    /// The columns the read declares (#955), from the backend's describe of
210    /// it. Default: none.
211    async fn declared_columns(
212        &self,
213        calls: &CallSlot,
214        context: &HashMap<String, Value>,
215    ) -> Result<Option<crate::declared::DeclaredColumns>, FaucetError> {
216        let _ = (calls, context);
217        Ok(None)
218    }
219
220    async fn discover(
221        &self,
222        calls: &CallSlot,
223    ) -> Result<Vec<crate::discover::DatasetDescriptor>, FaucetError> {
224        let _ = calls;
225        Err(FaucetError::Source(format!(
226            "source '{}' does not support dataset discovery",
227            self.connector_name()
228        )))
229    }
230}
231
232/// A stream of raw byte batches a [`SourceDriver`] reads.
233pub type NativeStream<'a> =
234    Pin<Box<dyn Stream<Item = Result<crate::native::NativeBatch, FaucetError>> + Send + 'a>>;
235
236/// A stream of Arrow pages a [`SourceDriver`] reads.
237#[cfg(feature = "arrow")]
238pub type ColumnarStream<'a> =
239    Pin<Box<dyn Stream<Item = Result<crate::columnar::ColumnarPage, FaucetError>> + Send + 'a>>;
240
241/// Split a buffered read into pages of `batch_size` records (`0`: one page),
242/// the bookmark on the last page — the shape a source without a native paging
243/// primitive streams in. An empty read yields one empty page only when it
244/// carries a bookmark.
245///
246/// ```
247/// use faucet_core::driver::pages_of;
248/// use faucet_core::json;
249///
250/// let pages: Vec<_> = pages_of(vec![json!(1), json!(2), json!(3)], Some(json!(3)), 2).collect();
251/// assert_eq!(pages.len(), 2);
252/// assert_eq!(pages[0].bookmark, None);
253/// assert_eq!(pages[1].bookmark, Some(json!(3)));
254/// ```
255pub fn pages_of(
256    records: Vec<Value>,
257    bookmark: Option<Value>,
258    batch_size: usize,
259) -> impl Iterator<Item = StreamPage> + Send {
260    let chunk = if batch_size == 0 {
261        usize::MAX
262    } else {
263        batch_size
264    };
265    let total = records.len();
266    let empty = (total == 0 && bookmark.is_some()).then(|| StreamPage {
267        records: Vec::new(),
268        bookmark: bookmark.clone(),
269    });
270    let mut iter = records.into_iter();
271    let mut consumed = 0usize;
272    let pages = std::iter::from_fn(move || {
273        let batch: Vec<Value> = iter.by_ref().take(chunk).collect();
274        if batch.is_empty() {
275            return None;
276        }
277        consumed += batch.len();
278        Some(StreamPage {
279            bookmark: if consumed >= total {
280                bookmark.clone()
281            } else {
282                None
283            },
284            records: batch,
285        })
286    });
287    empty.into_iter().chain(pages)
288}
289
290/// A [`SourceDriver`] as a [`Source`] (RFC 0014): every fetch is the driver's
291/// page stream, and the run's call context reaches every read.
292pub struct DriverSource<D> {
293    driver: D,
294    calls: CallSlot,
295}
296
297impl<D: SourceDriver> DriverSource<D> {
298    /// The source over `driver`.
299    pub fn new(driver: D) -> Self {
300        let calls = CallSlot::new(driver.call_context());
301        Self { driver, calls }
302    }
303
304    /// The driver.
305    pub fn driver(&self) -> &D {
306        &self.driver
307    }
308
309    /// The driver, mutably.
310    pub fn driver_mut(&mut self) -> &mut D {
311        &mut self.driver
312    }
313
314    /// The call slot the reads receive.
315    pub fn calls(&self) -> &CallSlot {
316        &self.calls
317    }
318
319    /// Unwrap the driver.
320    pub fn into_driver(self) -> D {
321        self.driver
322    }
323
324    async fn drain(
325        &self,
326        context: &HashMap<String, Value>,
327    ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
328        if let Some(fetched) = self.driver.fetch(&self.calls, context).await {
329            return fetched;
330        }
331        let mut pages = self.driver.read(&self.calls, context, 0);
332        let mut rows = Vec::new();
333        let mut bookmark = None;
334        while let Some(page) = pages.next().await {
335            let page = page?;
336            rows.extend(page.records);
337            if page.bookmark.is_some() {
338                bookmark = page.bookmark;
339            }
340        }
341        Ok((rows, bookmark))
342    }
343}
344
345#[async_trait]
346impl<D: SourceDriver> Source for DriverSource<D> {
347    async fn fetch_with_context(
348        &self,
349        context: &HashMap<String, Value>,
350    ) -> Result<Vec<Value>, FaucetError> {
351        Ok(self.drain(context).await?.0)
352    }
353
354    async fn fetch_with_context_incremental(
355        &self,
356        context: &HashMap<String, Value>,
357    ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
358        self.drain(context).await
359    }
360
361    fn stream_pages<'a>(
362        &'a self,
363        context: &'a HashMap<String, Value>,
364        batch_size: usize,
365    ) -> PageStream<'a> {
366        self.driver.read(&self.calls, context, batch_size)
367    }
368
369    fn config_schema(&self) -> Value {
370        self.driver.config_schema()
371    }
372
373    fn state_key(&self) -> Option<String> {
374        self.driver.state_key()
375    }
376
377    async fn apply_start_bookmark(&self, bookmark: Value) -> Result<(), FaucetError> {
378        self.driver.resume_from(bookmark).await
379    }
380
381    fn supports_exactly_once(&self) -> bool {
382        self.driver.deterministic_replay()
383    }
384
385    fn consumes_destructively(&self) -> bool {
386        self.driver.consumes_destructively()
387    }
388
389    fn capabilities(&self) -> Capabilities {
390        self.driver.capabilities(Capabilities::of_source(self))
391    }
392
393    fn connector_name(&self) -> &'static str {
394        self.driver.connector_name()
395    }
396
397    fn set_roundtrip_recorder(&self, recorder: Arc<RoundtripRecorder>) {
398        self.calls.install_recorder(Arc::clone(&recorder));
399        self.driver.set_roundtrip_recorder(recorder);
400    }
401
402    fn set_call_context(&self, ctx: CallContext) {
403        self.calls.install(ctx.clone());
404        self.driver.set_call_context(ctx);
405    }
406
407    fn set_run_clock(&self, now: chrono::DateTime<chrono::Utc>) {
408        self.driver.set_run_clock(now);
409    }
410
411    fn native_output_formats(&self) -> &'static [crate::native::NativeFormat] {
412        self.driver.native_output_formats()
413    }
414
415    fn stream_native<'a>(
416        &'a self,
417        context: &'a HashMap<String, Value>,
418        format: crate::native::NativeFormat,
419        batch_size: usize,
420    ) -> NativeStream<'a> {
421        self.driver
422            .read_native(&self.calls, context, format, batch_size)
423    }
424
425    #[cfg(feature = "arrow")]
426    fn supports_columnar(&self) -> bool {
427        self.driver.supports_columnar()
428    }
429
430    #[cfg(feature = "arrow")]
431    fn stream_batches<'a>(
432        &'a self,
433        context: &'a HashMap<String, Value>,
434        batch_size: usize,
435    ) -> ColumnarStream<'a> {
436        self.driver.read_columnar(&self.calls, context, batch_size)
437    }
438
439    fn supports_discover(&self) -> bool {
440        self.driver.supports_discover()
441    }
442
443    async fn declared_columns(
444        &self,
445        context: &HashMap<String, Value>,
446    ) -> Result<Option<crate::declared::DeclaredColumns>, FaucetError> {
447        self.driver.declared_columns(&self.calls, context).await
448    }
449
450    async fn discover(&self) -> Result<Vec<crate::discover::DatasetDescriptor>, FaucetError> {
451        self.driver.discover(&self.calls).await
452    }
453
454    fn set_cancel(&self, cancel: crate::CancellationToken) -> bool {
455        self.driver.set_cancel(cancel)
456    }
457
458    fn dataset_uri(&self) -> String {
459        self.driver.dataset_uri()
460    }
461
462    async fn check(&self, ctx: &CheckContext) -> Result<CheckReport, FaucetError> {
463        if let Some(report) = self.driver.check(&self.calls, ctx).await? {
464            return Ok(report);
465        }
466        let empty = HashMap::new();
467        let start = std::time::Instant::now();
468        let mut pages = self.driver.read(&self.calls, &empty, 1);
469        let probe = match tokio::time::timeout(ctx.timeout, pages.next()).await {
470            Err(_) => Probe::fail("read", start.elapsed(), "timed out fetching first page"),
471            Ok(None) | Ok(Some(Ok(_))) => Probe::pass("read", start.elapsed()),
472            Ok(Some(Err(e))) => Probe::fail("read", start.elapsed(), e.to_string()),
473        };
474        Ok(CheckReport::single(probe))
475    }
476}
faucet_core/driver/source.rs

faucet_core/driver/
source.rs
