faucet_core/

traits.rs

//! Shared traits for faucet sources and sinks.

use crate::error::FaucetError;
use crate::pipeline::StreamPage;
use async_trait::async_trait;
use futures_core::Stream;
use serde_json::Value;
use std::pin::Pin;

/// A source fetches records from an external system.
#[async_trait]
pub trait Source: Send + Sync {
    /// Primary fetch method. Receives context from a parent source's records.
    ///
    /// An empty context map means this is a root source (no parent).
    /// Connectors that support being a child should use
    /// [`substitute_context()`](crate::util::substitute_context) to resolve
    /// `{placeholder}` tokens in their URL path, query parameters, headers,
    /// or body. Connectors that don't need parent context ignore the map.
    async fn fetch_with_context(
        &self,
        context: &std::collections::HashMap<String, Value>,
    ) -> Result<Vec<Value>, FaucetError>;

    /// Convenience: fetch with no parent context.
    async fn fetch_all(&self) -> Result<Vec<Value>, FaucetError> {
        self.fetch_with_context(&std::collections::HashMap::new())
            .await
    }

    /// Incremental fetch with parent context support.
    ///
    /// Returns the records and an optional bookmark value for incremental
    /// replication. The default delegates to `fetch_with_context` and
    /// returns `None` for the bookmark.
    async fn fetch_with_context_incremental(
        &self,
        context: &std::collections::HashMap<String, Value>,
    ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
        let records = self.fetch_with_context(context).await?;
        Ok((records, None))
    }

    /// Convenience: incremental fetch with no parent context.
    async fn fetch_all_incremental(&self) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
        self.fetch_with_context_incremental(&std::collections::HashMap::new())
            .await
    }

    /// Stream records page-by-page so the pipeline can write to the sink as
    /// pages arrive instead of buffering the full result set.
    ///
    /// `batch_size` is the *hint* the pipeline passes down; sources are free
    /// to use a larger or smaller native chunk (e.g. one page per HTTP
    /// response, one row-group per Parquet file) but should approximate it
    /// where feasible. The special value `batch_size = 0` means "do not
    /// batch — emit the entire result set in a single page." Sources that
    /// stream natively should treat `0` as "skip the chunking layer and
    /// yield one page after the underlying read completes" (useful for
    /// small lookup tables or for sinks like SQL `COPY` / BigQuery load
    /// jobs that prefer one large request).
    ///
    /// The default implementation fetches the full result set via
    /// [`fetch_with_context_incremental`](Self::fetch_with_context_incremental)
    /// and chunks it in memory by `batch_size`. The bookmark (when present)
    /// is attached to the *final* page so the pipeline only persists after
    /// the entire fetch has been written. Sources that can stream natively
    /// override this method and may emit per-page bookmarks (e.g. CDC).
    ///
    /// An empty result with a `Some(bookmark)` still yields one empty page
    /// carrying the bookmark, so incremental runs that produce no records
    /// still advance their checkpoint.
    fn stream_pages<'a>(
        &'a self,
        context: &'a std::collections::HashMap<String, Value>,
        batch_size: usize,
    ) -> Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>> {
        Box::pin(async_stream::try_stream! {
            let (records, bookmark) = self
                .fetch_with_context_incremental(context)
                .await?;
            let total = records.len();
            // batch_size == 0 means "no batching" — emit all records as one
            // page. Otherwise chunk into pages of size `batch_size`.
            let chunk = if batch_size == 0 { usize::MAX } else { batch_size };

            if total == 0 {
                if bookmark.is_some() {
                    yield StreamPage {
                        records: Vec::new(),
                        bookmark,
                    };
                }
                return;
            }

            let mut iter = records.into_iter();
            let mut consumed = 0usize;
            loop {
                let batch: Vec<Value> = iter.by_ref().take(chunk).collect();
                if batch.is_empty() {
                    break;
                }
                consumed += batch.len();
                let page_bookmark = if consumed >= total {
                    bookmark.clone()
                } else {
                    None
                };
                yield StreamPage {
                    records: batch,
                    bookmark: page_bookmark,
                };
            }
        })
    }

    /// Return a JSON Schema describing the configuration this source accepts.
    fn config_schema(&self) -> Value {
        serde_json::json!({"type": "object", "properties": {}})
    }

    /// Stable key under which this source's incremental-replication bookmark
    /// should be persisted in a [`StateStore`](crate::state::StateStore).
    ///
    /// Returning `Some(key)` opts this source into resumable runs: when the
    /// pipeline is configured with a state store via
    /// [`Pipeline::with_state_store`](crate::Pipeline::with_state_store), it
    /// reads the bookmark at `key` before fetching and writes the new
    /// bookmark back only after the sink confirms the batch was written.
    ///
    /// The default returns `None`, meaning the source is not persisted.
    /// Keys must satisfy [`validate_state_key`](crate::state::validate_state_key).
    fn state_key(&self) -> Option<String> {
        None
    }

    /// Apply a bookmark loaded from a [`StateStore`](crate::state::StateStore)
    /// as this run's starting point.
    ///
    /// The default implementation ignores the value, which keeps existing
    /// sources backwards-compatible. Sources that support incremental
    /// replication override this — typically by storing the value behind
    /// interior mutability and consulting it inside
    /// `fetch_with_context_incremental`.
    async fn apply_start_bookmark(&self, _bookmark: Value) -> Result<(), FaucetError> {
        Ok(())
    }

    /// Stable identifier used as the `connector` label on metrics and the
    /// `connector` attribute on spans. Defaults to the final segment of
    /// `std::any::type_name::<Self>()`, e.g. `"RestSource"`. Built-in
    /// connectors override with a short, friendly snake_case name (e.g.
    /// `"rest"`). Must return a non-empty string; observability decorators
    /// fall back to `"unknown"` in release builds if it is empty (and
    /// `debug_assert!` in debug builds).
    fn connector_name(&self) -> &'static str {
        crate::observability::strip_type_name(std::any::type_name::<Self>())
    }

    /// Run a fast, non-mutating preflight probe (used by `faucet doctor`).
    ///
    /// The default pulls a **single page** via
    /// [`stream_pages`](Self::stream_pages) and reports success/failure — it
    /// exercises the real read path (DNS, TLS, auth, the first request, the
    /// first-record decode) but never paginates the full dataset and never
    /// repeats. The page stream is dropped immediately after the first page.
    ///
    /// Sources whose first page *blocks* waiting for inbound data (webhook,
    /// websocket) or has *side effects* (CDC consuming WAL) override this with a
    /// cheaper, side-effect-free probe. Probe-level failures are returned as a
    /// [`ProbeStatus::Fail`](crate::check::ProbeStatus) inside `Ok(report)`.
    async fn check(
        &self,
        ctx: &crate::check::CheckContext,
    ) -> Result<crate::check::CheckReport, FaucetError> {
        use crate::check::{CheckReport, Probe};
        use futures::StreamExt;

        let empty = std::collections::HashMap::new();
        let start = std::time::Instant::now();
        let mut pages = self.stream_pages(&empty, 1);
        let probe = match tokio::time::timeout(ctx.timeout, pages.next()).await {
            Err(_) => Probe::fail("read", start.elapsed(), "timed out fetching first page"),
            Ok(None) | Ok(Some(Ok(_))) => Probe::pass("read", start.elapsed()),
            Ok(Some(Err(e))) => Probe::fail("read", start.elapsed(), e.to_string()),
        };
        Ok(CheckReport::single(probe))
    }
}

/// Per-row outcome from [`Sink::write_batch_partial`].
///
/// `Ok(())` — the row was durably written to the sink.
/// `Err(_)` — the row failed; the pipeline will route it to the DLQ when
/// one is configured.
pub type RowOutcome = Result<(), FaucetError>;

/// A sink writes records to an external system.
#[async_trait]
pub trait Sink: Send + Sync {
    /// Write a batch of records to the destination.
    ///
    /// Returns the number of records successfully written.
    async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError>;

    /// Flush any buffered data to the destination.
    ///
    /// The default implementation is a no-op (suitable for sinks that
    /// write immediately in `write_batch`).
    async fn flush(&self) -> Result<(), FaucetError> {
        Ok(())
    }

    /// Write a batch and report per-row outcomes.
    ///
    /// Sinks whose underlying API exposes per-row results (BigQuery
    /// `insertAll`, Elasticsearch `_bulk`) override this. The default
    /// implementation delegates to [`Self::write_batch`] and maps a single success
    /// onto a uniform all-`Ok(())` vector. An outer failure is bubbled up
    /// unchanged so the pipeline's DLQ router can apply its `on_batch_error`
    /// policy at a single decision point.
    async fn write_batch_partial(&self, records: &[Value]) -> Result<Vec<RowOutcome>, FaucetError> {
        self.write_batch(records).await?;
        Ok(records.iter().map(|_| Ok(())).collect())
    }

    /// Return a JSON Schema describing the configuration this sink accepts.
    ///
    /// The schema is auto-generated from the config struct using `schemars`.
    /// Callers can inspect it to discover required fields, types, defaults,
    /// and descriptions before constructing the sink.
    ///
    /// The default returns an empty object schema.
    fn config_schema(&self) -> Value {
        serde_json::json!({"type": "object", "properties": {}})
    }

    /// Stable identifier used as the `connector` label on metrics and the
    /// `connector` attribute on spans. See `Source::connector_name`.
    fn connector_name(&self) -> &'static str {
        crate::observability::strip_type_name(std::any::type_name::<Self>())
    }

    /// Run a fast, non-mutating preflight probe (used by `faucet doctor`).
    ///
    /// Unlike sources, a sink has no non-mutating "first page" equivalent
    /// (`write_batch` mutates the destination), so the default returns
    /// [`CheckReport::not_implemented`](crate::check::CheckReport::not_implemented).
    /// Built-in sinks override this with a connect / auth / metadata probe.
    ///
    /// The probe **MUST be idempotent and side-effect-free** — no inserts, no
    /// residual rows or objects — and must never put credentials or connection
    /// strings in a probe `reason`/`hint`.
    async fn check(
        &self,
        _ctx: &crate::check::CheckContext,
    ) -> Result<crate::check::CheckReport, FaucetError> {
        Ok(crate::check::CheckReport::not_implemented())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    // ── Mock Source ──────────────────────────────────────────────────────────

    struct MockSource {
        records: Vec<Value>,
    }

    #[async_trait]
    impl Source for MockSource {
        async fn fetch_with_context(
            &self,
            _context: &std::collections::HashMap<String, Value>,
        ) -> Result<Vec<Value>, FaucetError> {
            Ok(self.records.clone())
        }
    }

    struct IncrementalSource {
        records: Vec<Value>,
        bookmark: Value,
    }

    #[async_trait]
    impl Source for IncrementalSource {
        async fn fetch_with_context(
            &self,
            _context: &std::collections::HashMap<String, Value>,
        ) -> Result<Vec<Value>, FaucetError> {
            Ok(self.records.clone())
        }

        async fn fetch_with_context_incremental(
            &self,
            _context: &std::collections::HashMap<String, Value>,
        ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
            Ok((self.records.clone(), Some(self.bookmark.clone())))
        }
    }

    struct FailingSource;

    #[async_trait]
    impl Source for FailingSource {
        async fn fetch_with_context(
            &self,
            _context: &std::collections::HashMap<String, Value>,
        ) -> Result<Vec<Value>, FaucetError> {
            Err(FaucetError::Auth("no credentials".into()))
        }
    }

    // ── Mock Sink ───────────────────────────────────────────────────────────

    struct MockSink {
        written: std::sync::Mutex<Vec<Value>>,
    }

    impl MockSink {
        fn new() -> Self {
            Self {
                written: std::sync::Mutex::new(Vec::new()),
            }
        }
    }

    #[async_trait]
    impl Sink for MockSink {
        async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
            let mut w = self.written.lock().unwrap();
            w.extend(records.iter().cloned());
            Ok(records.len())
        }
    }

    struct FailingSink;

    #[async_trait]
    impl Sink for FailingSink {
        async fn write_batch(&self, _records: &[Value]) -> Result<usize, FaucetError> {
            Err(FaucetError::Sink("write failed".into()))
        }
    }

    // ── Source tests ────────────────────────────────────────────────────────

    #[tokio::test]
    async fn source_fetch_all_returns_records() {
        let source = MockSource {
            records: vec![json!({"id": 1}), json!({"id": 2})],
        };
        let records = source.fetch_all().await.unwrap();
        assert_eq!(records.len(), 2);
        assert_eq!(records[0]["id"], 1);
    }

    #[tokio::test]
    async fn source_fetch_all_empty() {
        let source = MockSource { records: vec![] };
        let records = source.fetch_all().await.unwrap();
        assert!(records.is_empty());
    }

    #[tokio::test]
    async fn source_default_incremental_returns_none_bookmark() {
        let source = MockSource {
            records: vec![json!({"id": 1})],
        };
        let (records, bookmark) = source.fetch_all_incremental().await.unwrap();
        assert_eq!(records.len(), 1);
        assert!(bookmark.is_none());
    }

    #[tokio::test]
    async fn source_custom_incremental_returns_bookmark() {
        let source = IncrementalSource {
            records: vec![json!({"id": 1})],
            bookmark: json!("2024-12-01"),
        };
        let (records, bookmark) = source.fetch_all_incremental().await.unwrap();
        assert_eq!(records.len(), 1);
        assert_eq!(bookmark, Some(json!("2024-12-01")));
    }

    #[tokio::test]
    async fn source_error_propagates() {
        let source = FailingSource;
        let result = source.fetch_all().await;
        assert!(result.is_err());
        assert!(matches!(result, Err(FaucetError::Auth(_))));
    }

    #[tokio::test]
    async fn source_as_trait_object() {
        let source: Box<dyn Source> = Box::new(MockSource {
            records: vec![json!({"id": 42})],
        });
        let records = source.fetch_all().await.unwrap();
        assert_eq!(records[0]["id"], 42);
    }

    // ── Sink tests ──────────────────────────────────────────────────────────

    #[tokio::test]
    async fn sink_write_batch_returns_count() {
        let sink = MockSink::new();
        let records = vec![json!({"id": 1}), json!({"id": 2}), json!({"id": 3})];
        let count = sink.write_batch(&records).await.unwrap();
        assert_eq!(count, 3);
    }

    #[tokio::test]
    async fn sink_write_batch_empty() {
        let sink = MockSink::new();
        let count = sink.write_batch(&[]).await.unwrap();
        assert_eq!(count, 0);
    }

    #[tokio::test]
    async fn sink_accumulates_records() {
        let sink = MockSink::new();
        sink.write_batch(&[json!({"a": 1})]).await.unwrap();
        sink.write_batch(&[json!({"b": 2})]).await.unwrap();
        let written = sink.written.lock().unwrap();
        assert_eq!(written.len(), 2);
    }

    #[tokio::test]
    async fn sink_default_flush_is_noop() {
        let sink = MockSink::new();
        assert!(sink.flush().await.is_ok());
    }

    #[tokio::test]
    async fn sink_error_propagates() {
        let sink = FailingSink;
        let result = sink.write_batch(&[json!({"id": 1})]).await;
        assert!(result.is_err());
        assert!(matches!(result, Err(FaucetError::Sink(_))));
    }

    #[tokio::test]
    async fn sink_as_trait_object() {
        let sink: Box<dyn Sink> = Box::new(MockSink::new());
        let count = sink.write_batch(&[json!({"id": 1})]).await.unwrap();
        assert_eq!(count, 1);
    }

    // ── stream_pages tests ──────────────────────────────────────────────────

    use crate::pipeline::DEFAULT_BATCH_SIZE;
    use futures::StreamExt;

    #[tokio::test]
    async fn default_stream_pages_chunks_records() {
        let source = MockSource {
            records: (0..5).map(|i| json!({"i": i})).collect(),
        };
        let ctx = std::collections::HashMap::new();
        let mut pages = source.stream_pages(&ctx, 2);
        let mut all = Vec::new();
        while let Some(page) = pages.next().await {
            all.push(page.unwrap());
        }
        // 5 records, batch_size=2 → pages of [2, 2, 1]
        assert_eq!(all.len(), 3);
        assert_eq!(all[0].records.len(), 2);
        assert_eq!(all[1].records.len(), 2);
        assert_eq!(all[2].records.len(), 1);
    }

    #[tokio::test]
    async fn default_stream_pages_attaches_bookmark_to_final_page_only() {
        let source = IncrementalSource {
            records: (0..5).map(|i| json!({"i": i})).collect(),
            bookmark: json!("v1"),
        };
        let ctx = std::collections::HashMap::new();
        let mut pages = source.stream_pages(&ctx, 2);
        let mut collected = Vec::new();
        while let Some(page) = pages.next().await {
            collected.push(page.unwrap());
        }
        assert_eq!(collected.len(), 3);
        assert!(collected[0].bookmark.is_none());
        assert!(collected[1].bookmark.is_none());
        assert_eq!(collected[2].bookmark, Some(json!("v1")));
    }

    #[tokio::test]
    async fn default_stream_pages_single_page_when_batch_size_exceeds_total() {
        let source = MockSource {
            records: vec![json!({"id": 1}), json!({"id": 2})],
        };
        let ctx = std::collections::HashMap::new();
        let mut pages = source.stream_pages(&ctx, 100);
        let mut collected = Vec::new();
        while let Some(page) = pages.next().await {
            collected.push(page.unwrap());
        }
        assert_eq!(collected.len(), 1);
        assert_eq!(collected[0].records.len(), 2);
    }

    #[tokio::test]
    async fn default_stream_pages_batch_size_zero_emits_single_page() {
        // batch_size = 0 is the "no batching" sentinel — yields every record
        // in one page regardless of total count.
        let source = MockSource {
            records: (0..50_000).map(|i| json!({"i": i})).collect(),
        };
        let ctx = std::collections::HashMap::new();
        let mut pages = source.stream_pages(&ctx, 0);
        let mut collected = Vec::new();
        while let Some(page) = pages.next().await {
            collected.push(page.unwrap());
        }
        assert_eq!(
            collected.len(),
            1,
            "batch_size=0 must emit exactly one page"
        );
        assert_eq!(collected[0].records.len(), 50_000);
    }

    #[tokio::test]
    async fn default_stream_pages_batch_size_zero_attaches_bookmark_to_sole_page() {
        let source = IncrementalSource {
            records: (0..3).map(|i| json!({"i": i})).collect(),
            bookmark: json!("v1"),
        };
        let ctx = std::collections::HashMap::new();
        let mut pages = source.stream_pages(&ctx, 0);
        let page = pages.next().await.unwrap().unwrap();
        assert_eq!(page.records.len(), 3);
        assert_eq!(page.bookmark, Some(json!("v1")));
        assert!(pages.next().await.is_none());
    }

    #[tokio::test]
    async fn default_stream_pages_empty_source_yields_no_pages() {
        let source = MockSource { records: vec![] };
        let ctx = std::collections::HashMap::new();
        let mut pages = source.stream_pages(&ctx, DEFAULT_BATCH_SIZE);
        assert!(pages.next().await.is_none());
    }

    #[tokio::test]
    async fn default_stream_pages_empty_source_with_bookmark_yields_single_empty_page() {
        let source = IncrementalSource {
            records: vec![],
            bookmark: json!("v0"),
        };
        let ctx = std::collections::HashMap::new();
        let mut pages = source.stream_pages(&ctx, DEFAULT_BATCH_SIZE);
        let mut collected = Vec::new();
        while let Some(page) = pages.next().await {
            collected.push(page.unwrap());
        }
        // One empty-records page that carries the bookmark, so the pipeline
        // still persists progress on otherwise-empty incremental runs.
        assert_eq!(collected.len(), 1);
        assert!(collected[0].records.is_empty());
        assert_eq!(collected[0].bookmark, Some(json!("v0")));
    }

    #[tokio::test]
    async fn default_stream_pages_propagates_fetch_errors() {
        let source = FailingSource;
        let ctx = std::collections::HashMap::new();
        let mut pages = source.stream_pages(&ctx, DEFAULT_BATCH_SIZE);
        let first = pages.next().await.unwrap();
        assert!(matches!(first, Err(FaucetError::Auth(_))));
    }

    #[test]
    fn source_default_connector_name_is_stripped_type_name() {
        // MockSource lives at `faucet_core::traits::tests::MockSource`; the
        // stripped type_name yields the trailing segment.
        let source = MockSource { records: vec![] };
        assert_eq!(source.connector_name(), "MockSource");
    }

    #[test]
    fn sink_default_connector_name_is_stripped_type_name() {
        let sink = MockSink::new();
        assert_eq!(sink.connector_name(), "MockSink");
    }

    // ── write_batch_partial tests ───────────────────────────────────────────

    #[tokio::test]
    async fn default_write_batch_partial_success_returns_all_ok() {
        let sink = MockSink::new();
        let records = vec![json!({"id": 1}), json!({"id": 2}), json!({"id": 3})];
        let outcomes = sink.write_batch_partial(&records).await.unwrap();
        assert_eq!(outcomes.len(), 3);
        assert!(outcomes.iter().all(|o| o.is_ok()));
        assert_eq!(sink.written.lock().unwrap().len(), 3);
    }

    #[tokio::test]
    async fn default_write_batch_partial_bubbles_outer_err() {
        let sink = FailingSink;
        let records = vec![json!({"id": 1}), json!({"id": 2})];
        let result = sink.write_batch_partial(&records).await;
        assert!(matches!(result, Err(FaucetError::Sink(_))));
    }

    #[tokio::test]
    async fn default_write_batch_partial_empty_returns_empty_vec() {
        let sink = MockSink::new();
        let outcomes = sink.write_batch_partial(&[]).await.unwrap();
        assert!(outcomes.is_empty());
    }

    #[tokio::test]
    async fn default_write_batch_partial_callable_through_trait_object() {
        let sink: Box<dyn Sink> = Box::new(MockSink::new());
        let records = vec![json!({"id": 1}), json!({"id": 2})];
        let outcomes = sink.write_batch_partial(&records).await.unwrap();
        assert_eq!(outcomes.len(), 2);
        assert!(outcomes.iter().all(|o| o.is_ok()));
    }

    // ── check() tests ─────────────────────────────────────────────────────────

    #[tokio::test]
    async fn source_default_check_pulls_first_page_and_passes() {
        let source = MockSource {
            records: vec![json!({"id": 1}), json!({"id": 2})],
        };
        let report = source
            .check(&crate::check::CheckContext::default())
            .await
            .unwrap();
        assert_eq!(report.failed_count(), 0);
        assert!(
            report
                .probes
                .iter()
                .any(|p| p.name == "read" && matches!(p.status, crate::check::ProbeStatus::Pass))
        );
    }

    #[tokio::test]
    async fn source_default_check_passes_on_empty_source() {
        let source = MockSource { records: vec![] };
        let report = source
            .check(&crate::check::CheckContext::default())
            .await
            .unwrap();
        // Reachable but empty is still a healthy source.
        assert_eq!(report.failed_count(), 0);
    }

    #[tokio::test]
    async fn source_default_check_fails_when_fetch_errors() {
        let source = FailingSource;
        let report = source
            .check(&crate::check::CheckContext::default())
            .await
            .unwrap();
        assert_eq!(report.failed_count(), 1);
        assert!(report.probes.iter().any(
            |p| p.name == "read" && matches!(p.status, crate::check::ProbeStatus::Fail { .. })
        ));
    }

    #[tokio::test]
    async fn sink_default_check_is_not_implemented_skip() {
        let sink = MockSink::new();
        let report = sink
            .check(&crate::check::CheckContext::default())
            .await
            .unwrap();
        assert_eq!(report.probes.len(), 1);
        assert!(matches!(
            report.probes[0].status,
            crate::check::ProbeStatus::Skip { .. }
        ));
    }

    #[tokio::test]
    async fn source_check_callable_through_trait_object() {
        let source: Box<dyn Source> = Box::new(MockSource {
            records: vec![json!({"id": 1})],
        });
        let report = source
            .check(&crate::check::CheckContext::default())
            .await
            .unwrap();
        assert_eq!(report.failed_count(), 0);
    }
}