use std::collections::HashMap;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use faucet_core::write_mode::WriteMode;
use faucet_core::{FaucetError, Sink, Source, StreamPage, Value, async_trait};
use futures_core::Stream;
use serde_json::json;
pub struct CountingSource {
total: usize,
batch: usize,
resumable: bool,
start: Arc<Mutex<usize>>,
}
impl CountingSource {
pub fn new(total: usize, batch: usize) -> Self {
Self {
total,
batch,
resumable: true,
start: Arc::new(Mutex::new(0)),
}
}
pub fn non_resumable(total: usize, batch: usize) -> Self {
Self {
total,
batch,
resumable: false,
start: Arc::new(Mutex::new(0)),
}
}
}
#[async_trait]
impl Source for CountingSource {
async fn fetch_with_context(
&self,
_context: &HashMap<String, Value>,
) -> Result<Vec<Value>, FaucetError> {
let start = *self.start.lock().unwrap();
Ok((start..self.total).map(|i| json!({ "n": i })).collect())
}
fn stream_pages<'a>(
&'a self,
_context: &'a HashMap<String, Value>,
_batch_size: usize,
) -> Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>> {
let batch = if self.batch == 0 {
self.total.max(1)
} else {
self.batch
};
let total = self.total;
let start = (*self.start.lock().unwrap()).min(total);
Box::pin(async_stream::try_stream! {
let mut n = start;
if n >= total {
yield StreamPage { records: Vec::new(), bookmark: Some(json!({ "n": total })) };
return;
}
while n < total {
let end = (n + batch).min(total);
let records: Vec<Value> = (n..end).map(|i| json!({ "n": i })).collect();
n = end;
let bookmark = if n >= total { Some(json!({ "n": total })) } else { None };
yield StreamPage { records, bookmark };
}
})
}
fn connector_name(&self) -> &'static str {
"counting-source"
}
fn state_key(&self) -> Option<String> {
Some("conformance:counting".to_string())
}
async fn apply_start_bookmark(&self, bookmark: Value) -> Result<(), FaucetError> {
if !self.resumable {
return Ok(());
}
if let Some(n) = bookmark.get("n").and_then(|v| v.as_u64()) {
*self.start.lock().unwrap() = n as usize;
}
Ok(())
}
}
pub struct FailingSource;
#[async_trait]
impl Source for FailingSource {
async fn fetch_with_context(
&self,
_context: &HashMap<String, Value>,
) -> Result<Vec<Value>, FaucetError> {
Err(FaucetError::Source(
"unreachable endpoint (test double)".to_string(),
))
}
fn connector_name(&self) -> &'static str {
"failing-source"
}
}
pub struct PanickingSource;
#[async_trait]
impl Source for PanickingSource {
async fn fetch_with_context(
&self,
_context: &HashMap<String, Value>,
) -> Result<Vec<Value>, FaucetError> {
panic!("connector bug: unwrap() on a None value");
}
fn connector_name(&self) -> &'static str {
"panicking-source"
}
}
#[derive(Clone, Default)]
pub struct TestSink {
key_field: Option<String>,
idempotent: bool,
keyed: Arc<Mutex<HashMap<String, Value>>>,
appended: Arc<Mutex<Vec<Value>>>,
tokens: Arc<Mutex<HashMap<String, String>>>,
write_calls: Arc<Mutex<usize>>,
}
impl TestSink {
pub fn new() -> Self {
Self::default()
}
pub fn keyed(key_field: impl Into<String>) -> Self {
Self {
key_field: Some(key_field.into()),
..Self::default()
}
}
pub fn idempotent(key_field: impl Into<String>) -> Self {
Self {
key_field: Some(key_field.into()),
idempotent: true,
..Self::default()
}
}
pub fn len(&self) -> usize {
if self.key_field.is_some() {
self.keyed.lock().unwrap().len()
} else {
self.appended.lock().unwrap().len()
}
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn total_written(&self) -> usize {
*self.write_calls.lock().unwrap()
}
}
#[async_trait]
impl Sink for TestSink {
async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
*self.write_calls.lock().unwrap() += records.len();
match &self.key_field {
Some(field) => {
let mut map = self.keyed.lock().unwrap();
for r in records {
let key = r.get(field).map(|v| v.to_string()).ok_or_else(|| {
FaucetError::Sink(format!("record missing key `{field}`"))
})?;
map.insert(key, r.clone());
}
}
None => self
.appended
.lock()
.unwrap()
.extend(records.iter().cloned()),
}
Ok(records.len())
}
fn supports_idempotent_writes(&self) -> bool {
self.idempotent
}
fn dedups_by_key(&self) -> bool {
self.key_field.is_some()
}
fn supported_write_modes(&self) -> &'static [WriteMode] {
if self.key_field.is_some() {
&[WriteMode::Append, WriteMode::Upsert, WriteMode::Delete]
} else {
&[WriteMode::Append]
}
}
async fn write_batch_idempotent(
&self,
records: &[Value],
scope: &str,
token: &str,
) -> Result<usize, FaucetError> {
self.tokens
.lock()
.unwrap()
.insert(scope.to_string(), token.to_string());
self.write_batch(records).await
}
async fn last_committed_token(&self, scope: &str) -> Result<Option<String>, FaucetError> {
Ok(self.tokens.lock().unwrap().get(scope).cloned())
}
fn connector_name(&self) -> &'static str {
"test-sink"
}
}
#[derive(Clone, Default)]
pub struct LyingIdempotentSink {
appended: Arc<Mutex<Vec<Value>>>,
}
impl LyingIdempotentSink {
pub fn new() -> Self {
Self::default()
}
pub fn len(&self) -> usize {
self.appended.lock().unwrap().len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
#[async_trait]
impl Sink for LyingIdempotentSink {
async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
self.appended
.lock()
.unwrap()
.extend(records.iter().cloned());
Ok(records.len())
}
fn supports_idempotent_writes(&self) -> bool {
true }
fn connector_name(&self) -> &'static str {
"lying-idempotent-sink"
}
}
#[derive(Clone, Default)]
pub struct LyingKeyedSink {
appended: Arc<Mutex<Vec<Value>>>,
}
impl LyingKeyedSink {
pub fn new() -> Self {
Self::default()
}
pub fn len(&self) -> usize {
self.appended.lock().unwrap().len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
#[async_trait]
impl Sink for LyingKeyedSink {
async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
self.appended
.lock()
.unwrap()
.extend(records.iter().cloned());
Ok(records.len())
}
fn dedups_by_key(&self) -> bool {
true }
fn supported_write_modes(&self) -> &'static [WriteMode] {
&[WriteMode::Append, WriteMode::Upsert]
}
fn connector_name(&self) -> &'static str {
"lying-keyed-sink"
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use std::collections::HashMap;
#[tokio::test]
async fn counting_source_resumes_and_ignores_when_non_resumable() {
let s = CountingSource::new(5, 2);
assert_eq!(s.state_key().as_deref(), Some("conformance:counting"));
assert_eq!(s.connector_name(), "counting-source");
assert_eq!(
s.fetch_with_context(&HashMap::new()).await.unwrap().len(),
5
);
s.apply_start_bookmark(json!({ "n": 5 })).await.unwrap();
assert!(
s.fetch_with_context(&HashMap::new())
.await
.unwrap()
.is_empty()
);
let nr = CountingSource::non_resumable(5, 2);
nr.apply_start_bookmark(json!({ "n": 5 })).await.unwrap();
assert_eq!(
nr.fetch_with_context(&HashMap::new()).await.unwrap().len(),
5
);
}
#[tokio::test]
async fn test_sink_accessors() {
let s = TestSink::new();
assert!(s.is_empty());
s.write_batch(&[json!({ "id": 1 })]).await.unwrap();
assert!(!s.is_empty());
assert_eq!(s.len(), 1);
assert_eq!(s.total_written(), 1);
assert_eq!(s.connector_name(), "test-sink");
}
#[tokio::test]
async fn lying_idempotent_sink_never_persists_a_token() {
let s = LyingIdempotentSink::new();
assert!(s.is_empty());
assert!(s.supports_idempotent_writes());
assert_eq!(s.connector_name(), "lying-idempotent-sink");
s.write_batch_idempotent(&[json!({ "id": 1 })], "scope", "00000000000000000001")
.await
.unwrap();
assert_eq!(s.len(), 1);
assert!(s.last_committed_token("scope").await.unwrap().is_none());
}
#[tokio::test]
async fn lying_keyed_sink_appends_duplicates() {
let s = LyingKeyedSink::new();
assert!(s.is_empty());
assert!(s.dedups_by_key());
assert!(s.supported_write_modes().contains(&WriteMode::Upsert));
assert_eq!(s.connector_name(), "lying-keyed-sink");
s.write_batch(&[json!({ "id": 1 })]).await.unwrap();
s.write_batch(&[json!({ "id": 1 })]).await.unwrap();
assert_eq!(s.len(), 2, "lying keyed sink does not dedup");
}
#[tokio::test]
async fn failing_and_panicking_source_labels() {
assert_eq!(FailingSource.connector_name(), "failing-source");
assert_eq!(PanickingSource.connector_name(), "panicking-source");
assert!(FailingSource.fetch_all().await.is_err());
}
}