1use crate::accounting::WriteReport;
15use crate::capabilities::Capabilities;
16use crate::check::{CheckContext, CheckReport};
17use crate::cleanup::SeenKeys;
18use crate::dlq::BatchAtomicity;
19use crate::drift::SchemaEvolution;
20use crate::error::FaucetError;
21use crate::idempotency::SinkGuarantee;
22use crate::io::{CallContext, CallSlot};
23use crate::local_outputs::LocalOutput;
24use crate::observability::RoundtripRecorder;
25use crate::rollback::{RollbackOptions, RollbackOutcome};
26use crate::traits::{RowOutcome, Sink};
27use crate::write_mode::WriteMode;
28use async_trait::async_trait;
29use serde_json::Value;
30use std::collections::BTreeMap;
31use std::sync::Arc;
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35#[non_exhaustive]
36pub enum RowPolicy {
37 WholePage,
42 PerRow,
45}
46
47#[derive(Debug)]
49#[non_exhaustive]
50pub struct PageWrite {
51 pub report: WriteReport,
53 pub failed: Vec<(usize, FaucetError)>,
55}
56
57impl PageWrite {
58 pub fn accepted(n: usize) -> Self {
60 Self::reported(WriteReport::accepted(n))
61 }
62
63 pub fn reported(report: WriteReport) -> Self {
65 Self {
66 report,
67 failed: Vec::new(),
68 }
69 }
70
71 pub fn with_failures(mut self, failed: Vec<(usize, FaucetError)>) -> Self {
73 self.failed = failed;
74 self
75 }
76
77 pub fn from_outcomes(outcomes: Vec<RowOutcome>) -> Self {
79 let mut failed = Vec::new();
80 let mut ok = 0usize;
81 for (i, o) in outcomes.into_iter().enumerate() {
82 match o {
83 Ok(()) => ok += 1,
84 Err(e) => failed.push((i, e)),
85 }
86 }
87 Self::accepted(ok).with_failures(failed)
88 }
89}
90
91#[async_trait]
98pub trait StagingDriver: Send + Sync {
99 fn active(&self) -> bool;
102
103 async fn stage(&self, calls: &CallSlot) -> Result<(), FaucetError>;
106
107 async fn publish(&self, calls: &CallSlot) -> Result<(), FaucetError>;
110
111 async fn discard(&self, calls: &CallSlot) -> Result<(), FaucetError>;
114
115 async fn staged(&self, calls: &CallSlot) -> Result<bool, FaucetError>;
117
118 fn reports_staged(&self) -> bool {
123 true
124 }
125}
126
127#[async_trait]
130pub trait CommitTokenDriver: Send + Sync {
131 async fn last_token(
133 &self,
134 calls: &CallSlot,
135 scope: &str,
136 ) -> Result<Option<String>, FaucetError>;
137
138 async fn write_with_token(
141 &self,
142 calls: &CallSlot,
143 rows: &[Value],
144 scope: &str,
145 token: &str,
146 ) -> Result<WriteReport, FaucetError>;
147}
148
149#[async_trait]
151pub trait CleanupDriver: Send + Sync {
152 async fn delete_absent(
155 &self,
156 calls: &CallSlot,
157 scope: &BTreeMap<String, Value>,
158 seen: &SeenKeys,
159 ) -> Result<u64, FaucetError>;
160}
161
162#[async_trait]
164pub trait SchemaDriver: Send + Sync {
165 async fn current_schema(&self, calls: &CallSlot) -> Result<Option<Value>, FaucetError>;
168
169 fn evolvable(&self) -> bool {
171 false
172 }
173
174 async fn evolve(
176 &self,
177 calls: &CallSlot,
178 evolution: &SchemaEvolution,
179 ) -> Result<(), FaucetError> {
180 let _ = (calls, evolution);
181 Err(FaucetError::Sink(
182 "schema evolution is not supported".into(),
183 ))
184 }
185}
186
187#[async_trait]
189pub trait RollbackDriver: Send + Sync {
190 fn enabled(&self) -> bool {
194 true
195 }
196
197 async fn rollback_run(
199 &self,
200 calls: &CallSlot,
201 run_id: &str,
202 opts: &RollbackOptions,
203 ) -> Result<RollbackOutcome, FaucetError>;
204
205 async fn forget_run(&self, calls: &CallSlot, run_id: &str) -> Result<(), FaucetError>;
207
208 async fn rewind_commit_token(
210 &self,
211 calls: &CallSlot,
212 scope: &str,
213 token: Option<&str>,
214 ) -> Result<(), FaucetError>;
215}
216
217#[cfg(feature = "arrow")]
219#[async_trait]
220pub trait ColumnarDriver: Send + Sync {
221 async fn write_columnar(
223 &self,
224 calls: &CallSlot,
225 batch: &arrow::array::RecordBatch,
226 ) -> Result<usize, FaucetError>;
227}
228
229#[async_trait]
275pub trait SinkDriver: Send + Sync {
276 fn connector_name(&self) -> &'static str;
278
279 fn config_schema(&self) -> Value {
281 serde_json::json!({"type": "object", "properties": {}})
282 }
283
284 fn dataset_uri(&self) -> String {
286 format!("{}://unknown", self.connector_name())
287 }
288
289 fn call_context(&self) -> CallContext {
292 CallContext::new().without_timeout()
293 }
294
295 async fn check(
297 &self,
298 calls: &CallSlot,
299 ctx: &CheckContext,
300 ) -> Result<CheckReport, FaucetError> {
301 let _ = (calls, ctx);
302 Ok(CheckReport::not_implemented())
303 }
304
305 async fn write(
309 &self,
310 calls: &CallSlot,
311 rows: &[Value],
312 policy: RowPolicy,
313 ) -> Result<PageWrite, FaucetError>;
314
315 async fn flush(&self, calls: &CallSlot) -> Result<(), FaucetError> {
317 let _ = calls;
318 Ok(())
319 }
320
321 fn write_modes(&self) -> &'static [WriteMode] {
325 &[WriteMode::Append]
326 }
327
328 fn dedups_by_key(&self) -> bool {
331 false
332 }
333
334 fn write_is_replay_safe(&self) -> bool {
336 self.dedups_by_key()
337 }
338
339 fn batch_atomicity(&self) -> BatchAtomicity {
341 BatchAtomicity::BestEffort
342 }
343
344 fn staged_load(&self) -> bool {
346 false
347 }
348
349 fn capabilities(&self, derived: Capabilities) -> Capabilities {
353 derived
354 }
355
356 fn sink_guarantee(&self, derived: SinkGuarantee) -> SinkGuarantee {
360 derived
361 }
362
363 fn staging(&self) -> Option<&dyn StagingDriver> {
365 None
366 }
367
368 fn commit_tokens(&self) -> Option<&dyn CommitTokenDriver> {
370 None
371 }
372
373 fn cleanup(&self) -> Option<&dyn CleanupDriver> {
375 None
376 }
377
378 fn schema(&self) -> Option<&dyn SchemaDriver> {
380 None
381 }
382
383 fn rollback(&self) -> Option<&dyn RollbackDriver> {
385 None
386 }
387
388 #[cfg(feature = "arrow")]
390 fn columnar(&self) -> Option<&dyn ColumnarDriver> {
391 None
392 }
393
394 fn readback_source(&self) -> Option<(String, Value)> {
396 None
397 }
398
399 async fn complete_run(&self, calls: &CallSlot) -> Result<(), FaucetError> {
401 let _ = calls;
402 Ok(())
403 }
404
405 async fn local_outputs(&self) -> Vec<LocalOutput> {
407 Vec::new()
408 }
409
410 fn set_call_context(&self, ctx: CallContext) {
413 let _ = ctx;
414 }
415
416 fn set_roundtrip_recorder(&self, recorder: Arc<RoundtripRecorder>) {
419 let _ = recorder;
420 }
421}
422
423pub struct DriverSink<D> {
440 driver: D,
441 calls: CallSlot,
442}
443
444impl<D: SinkDriver> DriverSink<D> {
445 pub fn new(driver: D) -> Self {
454 let sink = Self::new_unchecked(driver);
455 if let Err(e) = sink.check_declarations() {
456 debug_assert!(false, "{e}");
457 tracing::error!(error = %e, "sink driver declarations disagree with its primitives");
458 }
459 sink
460 }
461
462 #[doc(hidden)]
464 pub fn new_unchecked(driver: D) -> Self {
465 let calls = CallSlot::new(driver.call_context());
466 Self { driver, calls }
467 }
468
469 pub fn driver(&self) -> &D {
471 &self.driver
472 }
473
474 pub fn driver_mut(&mut self) -> &mut D {
476 &mut self.driver
477 }
478
479 pub fn calls(&self) -> &CallSlot {
481 &self.calls
482 }
483
484 pub fn into_driver(self) -> D {
486 self.driver
487 }
488
489 pub fn check_declarations(&self) -> Result<(), FaucetError> {
494 let modes = self.driver.write_modes();
495 let name = self.driver.connector_name();
496 let staging = self.driver.staging();
497 if modes.contains(&WriteMode::Overwrite) != staging.is_some() {
498 return Err(FaucetError::Config(format!(
499 "sink driver '{name}': write mode overwrite is declared {} but a staging driver \
500 is {}",
501 if modes.contains(&WriteMode::Overwrite) {
502 "on"
503 } else {
504 "off"
505 },
506 if staging.is_some() {
507 "present"
508 } else {
509 "absent"
510 }
511 )));
512 }
513 if self.driver.dedups_by_key()
514 && !modes
515 .iter()
516 .any(|m| matches!(m, WriteMode::Upsert | WriteMode::Delete))
517 {
518 return Err(FaucetError::Config(format!(
519 "sink driver '{name}': it dedups by key but declares no keyed write mode"
520 )));
521 }
522 Ok(())
523 }
524
525 async fn whole_page(&self, records: &[Value]) -> Result<WriteReport, FaucetError> {
526 let w = self
527 .driver
528 .write(&self.calls, records, RowPolicy::WholePage)
529 .await?;
530 if let Some((_, e)) = w.failed.into_iter().min_by_key(|(i, _)| *i) {
531 return Err(e);
532 }
533 w.report.check(records.len())?;
534 Ok(w.report)
535 }
536}
537
538#[async_trait]
539impl<D: SinkDriver> Sink for DriverSink<D> {
540 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
541 Ok(self.whole_page(records).await?.accepted)
542 }
543
544 async fn write_batch_reported(&self, records: &[Value]) -> Result<WriteReport, FaucetError> {
545 self.whole_page(records).await
546 }
547
548 async fn write_batch_partial(&self, records: &[Value]) -> Result<Vec<RowOutcome>, FaucetError> {
549 let w = self
550 .driver
551 .write(&self.calls, records, RowPolicy::PerRow)
552 .await?;
553 let mut outcomes: Vec<RowOutcome> = records.iter().map(|_| Ok(())).collect();
554 let failed = w.failed.len();
555 for (i, e) in w.failed {
556 match outcomes.get_mut(i) {
557 Some(slot) => *slot = Err(e),
558 None => {
559 return Err(FaucetError::Sink(format!(
560 "sink driver '{}' reported row {i} failed in a page of {}",
561 self.driver.connector_name(),
562 records.len()
563 )));
564 }
565 }
566 }
567 w.report.check(records.len().saturating_sub(failed))?;
568 Ok(outcomes)
569 }
570
571 async fn flush(&self) -> Result<(), FaucetError> {
572 self.driver.flush(&self.calls).await
573 }
574
575 fn batch_atomicity(&self) -> BatchAtomicity {
576 self.driver.batch_atomicity()
577 }
578
579 #[cfg(feature = "arrow")]
580 fn supports_columnar(&self) -> bool {
581 self.driver.columnar().is_some()
582 }
583
584 #[cfg(feature = "arrow")]
585 async fn write_batch_columnar(
586 &self,
587 batch: &arrow::array::RecordBatch,
588 ) -> Result<usize, FaucetError> {
589 match self.driver.columnar() {
590 Some(c) => c.write_columnar(&self.calls, batch).await,
591 None => {
592 let rows = crate::columnar::record_batch_to_values(batch)?;
593 self.write_batch(&rows).await
594 }
595 }
596 }
597
598 fn supports_idempotent_writes(&self) -> bool {
599 self.driver.commit_tokens().is_some()
600 }
601
602 fn capabilities(&self) -> Capabilities {
603 self.driver.capabilities(Capabilities::of_sink(self))
604 }
605
606 fn sink_guarantee(&self) -> SinkGuarantee {
607 let derived = if self.supports_idempotent_writes() {
608 SinkGuarantee::AtomicWatermark
609 } else if self.driver.write_modes().contains(&WriteMode::Upsert) {
610 SinkGuarantee::KeyedUpsert
611 } else {
612 SinkGuarantee::AtLeastOnce
613 };
614 self.driver.sink_guarantee(derived)
615 }
616
617 fn write_batch_is_replay_safe(&self) -> bool {
618 self.driver.write_is_replay_safe()
619 }
620
621 fn dedups_by_key(&self) -> bool {
622 self.driver.dedups_by_key()
623 }
624
625 fn supported_write_modes(&self) -> &'static [WriteMode] {
626 self.driver.write_modes()
627 }
628
629 async fn current_schema(&self) -> Result<Option<Value>, FaucetError> {
630 match self.driver.schema() {
631 Some(s) => s.current_schema(&self.calls).await,
632 None => Ok(None),
633 }
634 }
635
636 fn supports_schema_evolution(&self) -> bool {
637 self.driver.schema().is_some_and(SchemaDriver::evolvable)
638 }
639
640 async fn evolve_schema(&self, evolution: &SchemaEvolution) -> Result<(), FaucetError> {
641 match self.driver.schema().filter(|s| s.evolvable()) {
642 Some(s) => s.evolve(&self.calls, evolution).await,
643 None => Err(FaucetError::Sink(format!(
644 "sink '{}' does not support schema evolution",
645 self.driver.connector_name()
646 ))),
647 }
648 }
649
650 fn supports_cleanup(&self) -> bool {
651 self.driver.cleanup().is_some()
652 }
653
654 fn supports_staged_load(&self) -> bool {
655 self.driver.staged_load()
656 }
657
658 async fn cleanup_scope(
659 &self,
660 scope: &BTreeMap<String, Value>,
661 seen: &SeenKeys,
662 ) -> Result<u64, FaucetError> {
663 match self.driver.cleanup() {
664 Some(c) => c.delete_absent(&self.calls, scope, seen).await,
665 None => Err(FaucetError::Sink(format!(
666 "sink '{}' does not support scoped cleanup",
667 self.driver.connector_name()
668 ))),
669 }
670 }
671
672 #[allow(deprecated)] async fn write_batch_idempotent(
674 &self,
675 records: &[Value],
676 scope: &str,
677 token: &str,
678 ) -> Result<usize, FaucetError> {
679 Ok(self
680 .write_batch_idempotent_reported(records, scope, token)
681 .await?
682 .accepted)
683 }
684
685 async fn write_batch_idempotent_reported(
686 &self,
687 records: &[Value],
688 scope: &str,
689 token: &str,
690 ) -> Result<WriteReport, FaucetError> {
691 match self.driver.commit_tokens() {
692 Some(c) => c.write_with_token(&self.calls, records, scope, token).await,
693 None => self.whole_page(records).await,
694 }
695 }
696
697 async fn last_committed_token(&self, scope: &str) -> Result<Option<String>, FaucetError> {
698 match self.driver.commit_tokens() {
699 Some(c) => c.last_token(&self.calls, scope).await,
700 None => Ok(None),
701 }
702 }
703
704 fn is_overwrite(&self) -> bool {
705 self.driver.staging().is_some_and(StagingDriver::active)
706 }
707
708 async fn begin_overwrite(&self) -> Result<(), FaucetError> {
709 self.staging("write_mode: overwrite")?
710 .stage(&self.calls)
711 .await
712 }
713
714 async fn commit_overwrite(&self) -> Result<(), FaucetError> {
715 self.staging("write_mode: overwrite")?
716 .publish(&self.calls)
717 .await
718 }
719
720 async fn abort_overwrite(&self) -> Result<(), FaucetError> {
721 match self.driver.staging() {
722 Some(s) => s.discard(&self.calls).await,
723 None => Ok(()),
724 }
725 }
726
727 async fn overwrite_staging_exists(&self) -> Result<Option<bool>, FaucetError> {
728 match self.driver.staging() {
729 Some(s) if s.reports_staged() => s.staged(&self.calls).await.map(Some),
730 _ => Ok(None),
731 }
732 }
733
734 fn supports_rollback(&self) -> bool {
735 self.driver.rollback().is_some_and(RollbackDriver::enabled)
736 }
737
738 async fn rollback_run(
739 &self,
740 run_id: &str,
741 opts: &RollbackOptions,
742 ) -> Result<RollbackOutcome, FaucetError> {
743 match self.driver.rollback() {
744 Some(r) => r.rollback_run(&self.calls, run_id, opts).await,
745 None => Err(FaucetError::Sink(format!(
746 "sink '{}' does not support rollback",
747 self.driver.connector_name()
748 ))),
749 }
750 }
751
752 async fn forget_run(&self, run_id: &str) -> Result<(), FaucetError> {
753 match self.driver.rollback() {
754 Some(r) => r.forget_run(&self.calls, run_id).await,
755 None => Ok(()),
756 }
757 }
758
759 async fn rewind_commit_token(
760 &self,
761 scope: &str,
762 token: Option<&str>,
763 ) -> Result<(), FaucetError> {
764 match self.driver.rollback() {
765 Some(r) => r.rewind_commit_token(&self.calls, scope, token).await,
766 None => Err(FaucetError::Sink(format!(
767 "sink '{}' does not support rewinding its commit token",
768 self.driver.connector_name()
769 ))),
770 }
771 }
772
773 fn readback_source(&self) -> Option<(String, Value)> {
774 self.driver.readback_source()
775 }
776
777 async fn complete_run(&self) -> Result<(), FaucetError> {
778 self.driver.complete_run(&self.calls).await
779 }
780
781 fn config_schema(&self) -> Value {
782 self.driver.config_schema()
783 }
784
785 fn connector_name(&self) -> &'static str {
786 self.driver.connector_name()
787 }
788
789 fn set_roundtrip_recorder(&self, recorder: Arc<RoundtripRecorder>) {
790 self.calls.install_recorder(Arc::clone(&recorder));
791 self.driver.set_roundtrip_recorder(recorder);
792 }
793
794 fn set_call_context(&self, ctx: CallContext) {
795 self.calls.install(ctx.clone());
796 self.driver.set_call_context(ctx);
797 }
798
799 fn dataset_uri(&self) -> String {
800 self.driver.dataset_uri()
801 }
802
803 async fn local_outputs(&self) -> Vec<LocalOutput> {
804 self.driver.local_outputs().await
805 }
806
807 async fn check(&self, ctx: &CheckContext) -> Result<CheckReport, FaucetError> {
808 self.driver.check(&self.calls, ctx).await
809 }
810}
811
812impl<D: SinkDriver> DriverSink<D> {
813 fn staging(&self, what: &str) -> Result<&dyn StagingDriver, FaucetError> {
814 self.driver.staging().ok_or_else(|| {
815 FaucetError::Sink(format!(
816 "sink '{}' does not support {what}",
817 self.driver.connector_name()
818 ))
819 })
820 }
821}