1use crate::idempotency::{ReplayGuarantee, SinkGuarantee};
33use crate::traits::{Sink, Source};
34use crate::write_mode::{WriteMode, WriteSpec};
35use serde::{Deserialize, Serialize};
36use serde_json::Value;
37
38#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
41#[serde(rename_all = "snake_case")]
42#[non_exhaustive]
43pub enum SchemaReport {
44 #[default]
46 Undeclared,
47 Reported,
50 Schemaless,
52 DropsUnknownFields,
56}
57
58impl SchemaReport {
59 pub fn as_str(&self) -> &'static str {
61 match self {
62 SchemaReport::Undeclared => "undeclared",
63 SchemaReport::Reported => "reported",
64 SchemaReport::Schemaless => "schemaless",
65 SchemaReport::DropsUnknownFields => "drops_unknown_fields",
66 }
67 }
68}
69
70#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize)]
77#[non_exhaustive]
78pub struct Capabilities {
79 write_modes: Vec<WriteMode>,
80 idempotent_writes: bool,
81 keyed_upsert: bool,
82 keyed_dedup: bool,
83 cleanup: bool,
84 rollback: bool,
85 schema_evolution: bool,
86 schema_report: SchemaReport,
87 scoped_overwrite: bool,
88 direct_overwrite: bool,
89 staged_load: bool,
90 exactly_once: bool,
91 discover: bool,
92 shardable: bool,
93 range_read: bool,
94 range_digest: bool,
95 reports_lag: bool,
96 columnar: bool,
97 native: bool,
98 table_routing: bool,
99 read_side_effects: bool,
100 complete_reads: bool,
101}
102
103const MODE_ORDER: [WriteMode; 4] = [
104 WriteMode::Append,
105 WriteMode::Upsert,
106 WriteMode::Delete,
107 WriteMode::Overwrite,
108];
109
110fn mode_rank(m: &WriteMode) -> usize {
111 MODE_ORDER
112 .iter()
113 .position(|o| o == m)
114 .unwrap_or(MODE_ORDER.len())
115}
116
117impl Capabilities {
118 pub fn builder() -> CapabilitiesBuilder {
120 CapabilitiesBuilder(Capabilities::default())
121 }
122
123 pub fn to_builder(&self) -> CapabilitiesBuilder {
125 CapabilitiesBuilder(self.clone())
126 }
127
128 pub fn plain_source() -> Self {
130 Self::default()
131 }
132
133 pub fn append_only_sink() -> Self {
135 Self::builder().write_modes([WriteMode::Append]).build()
136 }
137
138 pub fn of_source<S: Source + ?Sized>(source: &S) -> Self {
143 #[cfg(feature = "arrow")]
144 let columnar = source.supports_columnar();
145 #[cfg(not(feature = "arrow"))]
146 let columnar = false;
147 Self::builder()
148 .exactly_once(source.replay_guarantee() == ReplayGuarantee::Deterministic)
149 .discover(source.supports_discover())
150 .shardable(source.is_shardable())
151 .columnar(columnar)
152 .native(!source.native_output_formats().is_empty())
153 .build()
154 }
155
156 pub fn of_sink<S: Sink + ?Sized>(sink: &S) -> Self {
159 #[cfg(feature = "arrow")]
160 let columnar = sink.supports_columnar();
161 #[cfg(not(feature = "arrow"))]
162 let columnar = false;
163 let modes = sink.supported_write_modes();
164 let schema_evolution = sink.supports_schema_evolution();
165 Self::builder()
166 .write_modes(modes.iter().copied())
167 .idempotent_writes(sink.supports_idempotent_writes())
168 .keyed_upsert(
169 modes.contains(&WriteMode::Upsert)
170 && sink.sink_guarantee() != SinkGuarantee::AtLeastOnce,
171 )
172 .keyed_dedup(sink.dedups_by_key())
173 .cleanup(sink.supports_cleanup())
174 .rollback(sink.supports_rollback())
175 .schema_evolution(schema_evolution)
176 .schema_report(if schema_evolution {
177 SchemaReport::Reported
178 } else {
179 SchemaReport::Undeclared
180 })
181 .staged_load(sink.supports_staged_load())
182 .columnar(columnar)
183 .native(!sink.native_load_capabilities().is_empty())
184 .build()
185 }
186
187 pub fn decorated_source<S: Source + ?Sized>(outer: &S, inner: &Capabilities) -> Self {
191 Self::of_source(outer)
192 .to_builder()
193 .range_read(inner.range_read)
194 .range_digest(inner.range_digest)
195 .reports_lag(inner.reports_lag)
196 .table_routing(inner.table_routing)
197 .read_side_effects(inner.read_side_effects)
198 .complete_reads(inner.complete_reads)
199 .build()
200 }
201
202 pub fn decorated_sink<S: Sink + ?Sized>(outer: &S, inner: &Capabilities) -> Self {
204 Self::of_sink(outer)
205 .to_builder()
206 .schema_report(inner.schema_report)
207 .scoped_overwrite(inner.scoped_overwrite)
208 .direct_overwrite(inner.direct_overwrite)
209 .build()
210 }
211
212 pub fn write_modes(&self) -> &[WriteMode] {
214 &self.write_modes
215 }
216
217 pub fn supports_write_mode(&self, mode: WriteMode) -> bool {
219 self.write_modes.contains(&mode)
220 }
221
222 pub fn overwrite(&self) -> bool {
224 self.supports_write_mode(WriteMode::Overwrite)
225 }
226
227 pub fn idempotent_writes(&self) -> bool {
229 self.idempotent_writes
230 }
231
232 pub fn keyed_upsert(&self) -> bool {
236 self.keyed_upsert
237 }
238
239 pub fn keyed_dedup(&self) -> bool {
241 self.keyed_dedup
242 }
243
244 pub fn cleanup(&self) -> bool {
246 self.cleanup
247 }
248
249 pub fn rollback(&self) -> bool {
251 self.rollback
252 }
253
254 pub fn schema_evolution(&self) -> bool {
256 self.schema_evolution
257 }
258
259 pub fn schema_report(&self) -> SchemaReport {
261 self.schema_report
262 }
263
264 pub fn scoped_overwrite(&self) -> bool {
266 self.scoped_overwrite
267 }
268
269 pub fn direct_overwrite(&self) -> bool {
271 self.direct_overwrite
272 }
273
274 pub fn staged_load(&self) -> bool {
276 self.staged_load
277 }
278
279 pub fn exactly_once(&self) -> bool {
281 self.exactly_once
282 }
283
284 pub fn discover(&self) -> bool {
286 self.discover
287 }
288
289 pub fn shardable(&self) -> bool {
291 self.shardable
292 }
293
294 pub fn range_read(&self) -> bool {
297 self.range_read
298 }
299
300 pub fn range_digest(&self) -> bool {
302 self.range_digest
303 }
304
305 pub fn reports_lag(&self) -> bool {
307 self.reports_lag
308 }
309
310 pub fn columnar(&self) -> bool {
312 self.columnar
313 }
314
315 pub fn native(&self) -> bool {
317 self.native
318 }
319
320 pub fn table_routing(&self) -> bool {
323 self.table_routing
324 }
325
326 pub fn read_side_effects(&self) -> bool {
330 self.read_side_effects
331 }
332
333 pub fn complete_reads(&self) -> bool {
337 self.complete_reads
338 }
339
340 pub fn replay_guarantee(&self) -> ReplayGuarantee {
342 if self.exactly_once {
343 ReplayGuarantee::Deterministic
344 } else {
345 ReplayGuarantee::NonDeterministic
346 }
347 }
348
349 pub fn sink_guarantee(&self) -> SinkGuarantee {
351 if self.idempotent_writes {
352 SinkGuarantee::AtomicWatermark
353 } else if self.keyed_upsert {
354 SinkGuarantee::KeyedUpsert
355 } else {
356 SinkGuarantee::AtLeastOnce
357 }
358 }
359}
360
361#[derive(Debug, Clone, Default)]
363pub struct CapabilitiesBuilder(Capabilities);
364
365macro_rules! flag {
366 ($($name:ident),* $(,)?) => {
367 $(
368 #[doc = concat!("Set `", stringify!($name), "`.")]
369 #[must_use]
370 pub fn $name(mut self, on: bool) -> Self {
371 self.0.$name = on;
372 self
373 }
374 )*
375 };
376}
377
378impl CapabilitiesBuilder {
379 #[must_use]
381 pub fn write_modes(mut self, modes: impl IntoIterator<Item = WriteMode>) -> Self {
382 let mut v: Vec<WriteMode> = Vec::new();
383 for m in modes {
384 if !v.contains(&m) {
385 v.push(m);
386 }
387 }
388 v.sort_by_key(mode_rank);
389 self.0.write_modes = v;
390 self
391 }
392
393 #[must_use]
395 pub fn schema_report(mut self, report: SchemaReport) -> Self {
396 self.0.schema_report = report;
397 self
398 }
399
400 flag!(
401 idempotent_writes,
402 keyed_upsert,
403 keyed_dedup,
404 cleanup,
405 rollback,
406 schema_evolution,
407 scoped_overwrite,
408 direct_overwrite,
409 staged_load,
410 exactly_once,
411 discover,
412 shardable,
413 range_read,
414 range_digest,
415 reports_lag,
416 columnar,
417 native,
418 table_routing,
419 read_side_effects,
420 complete_reads,
421 );
422
423 pub fn build(self) -> Capabilities {
425 self.0
426 }
427}
428
429#[derive(Clone, Copy)]
433#[non_exhaustive]
434pub struct ConnectorDescriptor {
435 kind: &'static str,
436 capabilities: fn(&Value) -> Capabilities,
437}
438
439impl ConnectorDescriptor {
440 pub const fn new(kind: &'static str, capabilities: fn(&Value) -> Capabilities) -> Self {
443 Self { kind, capabilities }
444 }
445
446 pub fn kind(&self) -> &'static str {
448 self.kind
449 }
450
451 pub fn capabilities(&self, config: &Value) -> Capabilities {
453 (self.capabilities)(config)
454 }
455
456 pub fn kind_capabilities(&self) -> Capabilities {
460 (self.capabilities)(&Value::Null)
461 }
462}
463
464impl std::fmt::Debug for ConnectorDescriptor {
465 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
466 f.debug_struct("ConnectorDescriptor")
467 .field("kind", &self.kind)
468 .finish_non_exhaustive()
469 }
470}
471
472pub fn plain_source(_config: &Value) -> Capabilities {
474 Capabilities::plain_source()
475}
476
477pub fn append_only_sink(_config: &Value) -> Capabilities {
479 Capabilities::append_only_sink()
480}
481
482pub mod config {
484 use super::*;
485
486 pub fn write_spec(config: &Value) -> WriteSpec {
489 let mut spec = WriteSpec::default();
490 if let Some(mode) = config.get("write_mode")
491 && let Ok(m) = serde_json::from_value::<WriteMode>(mode.clone())
492 {
493 spec.write_mode = m;
494 }
495 if let Some(Value::Array(keys)) = config.get("key") {
496 spec.key = keys
497 .iter()
498 .filter_map(|k| k.as_str().map(str::to_string))
499 .collect();
500 }
501 spec
502 }
503
504 pub fn dedups_by_key(config: &Value) -> bool {
506 write_spec(config).dedups_by_key()
507 }
508
509 pub fn str_field<'a>(config: &'a Value, field: &str) -> Option<&'a str> {
511 config.get(field).and_then(Value::as_str)
512 }
513
514 pub fn bool_field(config: &Value, field: &str, default: bool) -> bool {
516 config
517 .get(field)
518 .and_then(Value::as_bool)
519 .unwrap_or(default)
520 }
521
522 pub fn is_set(config: &Value, field: &str) -> bool {
524 config.get(field).is_some_and(|v| !v.is_null())
525 }
526
527 pub fn tag<'a>(config: &'a Value, field: &str) -> Option<&'a str> {
530 match config.get(field)? {
531 Value::String(s) => Some(s),
532 Value::Object(o) => o
533 .get("type")
534 .and_then(Value::as_str)
535 .or_else(|| (o.len() == 1).then(|| o.keys().next().map(String::as_str))?),
536 _ => None,
537 }
538 }
539}
540
541#[cfg(test)]
542mod tests {
543 use super::*;
544 use crate::FaucetError;
545 use async_trait::async_trait;
546 use serde_json::json;
547
548 struct Plain;
549 #[async_trait]
550 impl Source for Plain {
551 async fn fetch_with_context(
552 &self,
553 _: &std::collections::HashMap<String, Value>,
554 ) -> Result<Vec<Value>, FaucetError> {
555 Ok(vec![])
556 }
557 }
558 #[async_trait]
559 impl Sink for Plain {
560 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
561 Ok(r.len())
562 }
563 }
564
565 struct Rich;
566 #[async_trait]
567 impl Source for Rich {
568 async fn fetch_with_context(
569 &self,
570 _: &std::collections::HashMap<String, Value>,
571 ) -> Result<Vec<Value>, FaucetError> {
572 Ok(vec![])
573 }
574 fn supports_exactly_once(&self) -> bool {
575 true
576 }
577 fn supports_discover(&self) -> bool {
578 true
579 }
580 fn is_shardable(&self) -> bool {
581 true
582 }
583 fn native_output_formats(&self) -> &'static [crate::native::NativeFormat] {
584 &[crate::native::NativeFormat::NdJson]
585 }
586 }
587 #[async_trait]
588 impl Sink for Rich {
589 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
590 Ok(r.len())
591 }
592 fn supported_write_modes(&self) -> &'static [WriteMode] {
593 &[WriteMode::Overwrite, WriteMode::Upsert, WriteMode::Append]
594 }
595 fn supports_idempotent_writes(&self) -> bool {
596 true
597 }
598 fn dedups_by_key(&self) -> bool {
599 true
600 }
601 fn supports_cleanup(&self) -> bool {
602 true
603 }
604 fn supports_rollback(&self) -> bool {
605 true
606 }
607 fn supports_schema_evolution(&self) -> bool {
608 true
609 }
610 fn supports_staged_load(&self) -> bool {
611 true
612 }
613 }
614
615 #[tokio::test]
618 async fn descriptors_built_at_run_time_and_the_test_connectors() {
619 let d = ConnectorDescriptor::new("rich", |_| Sink::capabilities(&Rich));
620 assert_eq!(d.kind(), "rich");
621 assert!(d.kind_capabilities().idempotent_writes());
622 let rows = [json!({"a": 1})];
623 for sink in [&Plain as &dyn Sink, &Rich] {
624 assert_eq!(sink.write_batch(&rows).await.unwrap(), 1);
625 }
626 for source in [&Plain as &dyn Source, &Rich] {
627 assert!(source.fetch_all().await.unwrap().is_empty());
628 }
629 }
630
631 #[test]
632 fn plain_connectors_derive_nothing() {
633 assert_eq!(Source::capabilities(&Plain), Capabilities::plain_source());
634 assert_eq!(Sink::capabilities(&Plain), Capabilities::append_only_sink());
635 assert_eq!(plain_source(&Value::Null), Capabilities::plain_source());
636 assert_eq!(
637 append_only_sink(&Value::Null),
638 Capabilities::append_only_sink()
639 );
640 let s = Capabilities::append_only_sink();
641 assert_eq!(s.sink_guarantee(), SinkGuarantee::AtLeastOnce);
642 assert_eq!(s.replay_guarantee(), ReplayGuarantee::NonDeterministic);
643 assert!(!s.overwrite());
644 assert_eq!(s.schema_report(), SchemaReport::Undeclared);
645 }
646
647 #[test]
648 fn derivation_reads_every_trait_method() {
649 let src = Source::capabilities(&Rich);
650 assert!(src.exactly_once() && src.discover() && src.shardable() && src.native());
651 assert!(!src.columnar() && !src.range_read() && !src.range_digest());
652 assert!(!src.reports_lag());
653 assert_eq!(src.replay_guarantee(), ReplayGuarantee::Deterministic);
654
655 let sink = Sink::capabilities(&Rich);
656 assert_eq!(
657 sink.write_modes(),
658 &[WriteMode::Append, WriteMode::Upsert, WriteMode::Overwrite]
659 );
660 assert!(sink.idempotent_writes() && sink.keyed_upsert() && sink.keyed_dedup());
661 assert!(sink.cleanup() && sink.rollback() && sink.schema_evolution());
662 assert!(sink.staged_load() && sink.overwrite());
663 assert!(!sink.scoped_overwrite() && !sink.direct_overwrite() && !sink.native());
664 assert_eq!(sink.schema_report(), SchemaReport::Reported);
665 assert_eq!(sink.sink_guarantee(), SinkGuarantee::AtomicWatermark);
666 }
667
668 #[test]
669 fn decorated_keeps_inner_declarations_and_outer_derivations() {
670 let inner = Capabilities::builder()
671 .range_read(true)
672 .range_digest(true)
673 .reports_lag(true)
674 .table_routing(true)
675 .read_side_effects(true)
676 .complete_reads(true)
677 .schema_report(SchemaReport::DropsUnknownFields)
678 .scoped_overwrite(true)
679 .direct_overwrite(true)
680 .discover(true)
681 .build();
682 let src = Capabilities::decorated_source(&Plain, &inner);
683 assert!(src.range_read() && src.range_digest() && src.reports_lag());
684 assert!(src.table_routing() && src.read_side_effects() && src.complete_reads());
685 assert!(!src.discover(), "derived from the outer connector");
686 let sink = Capabilities::decorated_sink(&Plain, &inner);
687 assert_eq!(sink.schema_report(), SchemaReport::DropsUnknownFields);
688 assert!(sink.scoped_overwrite() && sink.direct_overwrite());
689 assert_eq!(sink.write_modes(), &[WriteMode::Append]);
690 }
691
692 #[test]
693 fn builder_sets_every_flag_and_round_trips() {
694 let c = Capabilities::builder()
695 .write_modes([WriteMode::Delete, WriteMode::Append, WriteMode::Delete])
696 .keyed_upsert(true)
697 .scoped_overwrite(true)
698 .direct_overwrite(true)
699 .range_read(true)
700 .range_digest(true)
701 .reports_lag(true)
702 .columnar(true)
703 .schema_report(SchemaReport::DropsUnknownFields)
704 .build();
705 assert_eq!(c.write_modes(), &[WriteMode::Append, WriteMode::Delete]);
706 assert!(c.supports_write_mode(WriteMode::Delete));
707 assert!(c.scoped_overwrite() && c.direct_overwrite());
708 assert!(c.range_read() && c.range_digest() && c.reports_lag() && c.columnar());
709 assert_eq!(c.sink_guarantee(), SinkGuarantee::KeyedUpsert);
710 assert_eq!(c.to_builder().build(), c);
711 let off = c.to_builder().columnar(false).build();
712 assert!(!off.columnar());
713 assert!(!c.table_routing() && !c.read_side_effects());
714 let cdc = c
715 .to_builder()
716 .table_routing(true)
717 .read_side_effects(true)
718 .build();
719 assert!(cdc.table_routing() && cdc.read_side_effects());
720 let v = serde_json::to_value(&c).unwrap();
721 assert_eq!(v["schema_report"], "drops_unknown_fields");
722 assert_eq!(v["write_modes"], json!(["append", "delete"]));
723 }
724
725 #[test]
726 fn schema_report_names() {
727 for (r, n) in [
728 (SchemaReport::Undeclared, "undeclared"),
729 (SchemaReport::Reported, "reported"),
730 (SchemaReport::Schemaless, "schemaless"),
731 (SchemaReport::DropsUnknownFields, "drops_unknown_fields"),
732 ] {
733 assert_eq!(r.as_str(), n);
734 assert_eq!(serde_json::to_value(r).unwrap(), json!(n));
735 }
736 }
737
738 fn rich(c: &Value) -> Capabilities {
739 Capabilities::append_only_sink()
740 .to_builder()
741 .keyed_dedup(config::dedups_by_key(c))
742 .build()
743 }
744
745 #[test]
746 fn descriptor_answers_from_config() {
747 const D: ConnectorDescriptor = ConnectorDescriptor::new("rich", rich);
748 assert_eq!(D.kind(), "rich");
749 assert_eq!(D.kind_capabilities(), Capabilities::append_only_sink());
750 assert!(
751 D.capabilities(&json!({"write_mode": "upsert", "key": ["id"]}))
752 .keyed_dedup()
753 );
754 assert!(
755 !D.capabilities(&json!({"write_mode": "upsert"}))
756 .keyed_dedup()
757 );
758 assert!(!D.capabilities(&json!({"write_mode": "${x}"})).keyed_dedup());
759 assert!(format!("{D:?}").contains("rich"));
760 }
761
762 #[test]
763 fn config_helpers_are_lenient() {
764 let c = json!({
765 "write_mode": "delete", "key": ["a", 1], "s": "x", "b": true,
766 "n": null, "m": {"auto_map": {}}, "t": {"type": "jsonb"}, "e": "plain"
767 });
768 let spec = config::write_spec(&c);
769 assert_eq!(spec.write_mode, WriteMode::Delete);
770 assert_eq!(spec.key, vec!["a".to_string()]);
771 assert_eq!(config::str_field(&c, "s"), Some("x"));
772 assert!(config::bool_field(&c, "b", false));
773 assert!(config::bool_field(&c, "missing", true));
774 assert!(!config::is_set(&c, "n") && config::is_set(&c, "s"));
775 assert_eq!(config::tag(&c, "m"), Some("auto_map"));
776 assert_eq!(config::tag(&c, "t"), Some("jsonb"));
777 assert_eq!(config::tag(&c, "e"), Some("plain"));
778 assert_eq!(config::tag(&c, "b"), None);
779 assert_eq!(config::tag(&c, "missing"), None);
780 assert_eq!(
781 config::write_spec(&Value::Null).write_mode,
782 WriteMode::Append
783 );
784 }
785}