Skip to main content

laminar_connectors/connector/
mod.rs

1//! Connector traits — async `SourceConnector` / `SinkConnector`.
2
3mod contracts;
4mod coordinated_commit;
5mod sink;
6mod source;
7mod source_batch;
8mod task_tracking;
9
10pub use contracts::{
11    DeliveryGuarantee, SinkConsistency, SinkContract, SinkInputMode, SinkTopology,
12    SourceConsistency, SourceContract, SourceInputMode, SourceRowPositionCapability,
13    SourceTopology,
14};
15pub use coordinated_commit::{
16    CoordinatedAbortBatch, CoordinatedAbortCleaner, CoordinatedAbortDescriptor,
17    CoordinatedAbortEntry, CoordinatedCommitBatch, CoordinatedCommitContext,
18    CoordinatedCommitCursor, CoordinatedCommitNamespace, CoordinatedCommitPayload,
19    CoordinatedCommitter, MAX_COORDINATED_COMMIT_BATCH_BYTES, MAX_COORDINATED_COMMIT_BATCH_ENTRIES,
20    MAX_COORDINATED_COMMIT_PAYLOAD_BYTES,
21};
22pub use sink::{SinkConnector, SinkRuntimeContext, WriteResult};
23pub use source::{
24    SourceCheckpointUnavailablePolicy, SourceConnector, SourceDrainOutcome, SourceDrainRequest,
25    SourceDrainResolution, SourcePosition, SourceStart,
26};
27pub use source_batch::{
28    schema_with_source_mutations_and_row_positions, schema_with_source_row_positions,
29    source_mutations, source_mutations_routed, source_row_positions, strip_source_mutations,
30    strip_source_mutations_routed, strip_source_row_positions, SourceBatch, SourceBatchCursor,
31    SourceMutation, SourceMutationView, SourceRowPositionRef, SourceRowPositionView,
32    SourceRowPositions, SOURCE_MUTATION_COLUMN, SOURCE_ORDER_KEY_COLUMN, SOURCE_PARTITION_COLUMN,
33    SOURCE_SUB_OFFSET_COLUMN,
34};
35pub use task_tracking::{
36    ConnectorCancellationPolicy, ConnectorTaskAdmission, ConnectorTaskGuard, ConnectorTaskOwner,
37    ConnectorTaskTracker,
38};
39
40#[cfg(test)]
41#[allow(clippy::cast_possible_wrap)]
42mod tests {
43    use super::*;
44    use arrow_array::{ArrayRef, BinaryArray, Int64Array, RecordBatch, UInt32Array, UInt8Array};
45    use arrow_schema::{DataType, Field, Schema, SchemaRef};
46    use async_trait::async_trait;
47    use std::sync::Arc;
48
49    use crate::checkpoint::SourceCheckpoint;
50    use crate::config::ConnectorConfig;
51    use crate::error::ConnectorError;
52
53    fn test_schema() -> SchemaRef {
54        Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]))
55    }
56
57    fn test_batch(n: usize) -> RecordBatch {
58        #[allow(clippy::cast_possible_wrap)]
59        let ids: Vec<i64> = (0..n as i64).collect();
60        RecordBatch::try_new(test_schema(), vec![Arc::new(Int64Array::from(ids))]).unwrap()
61    }
62
63    #[test]
64    fn connector_task_generation_terminates_after_owner_and_every_guard() {
65        let (owner, tracker) = ConnectorTaskOwner::new();
66        let first = owner.track().expect("live generation");
67        let second = owner.track().expect("live generation");
68
69        assert!(!tracker.is_terminated());
70        drop(owner);
71        assert!(!tracker.is_terminated());
72        drop(first);
73        assert!(!tracker.is_terminated());
74        drop(second);
75        assert!(tracker.is_terminated());
76    }
77
78    #[test]
79    fn connector_task_admission_is_sealed_by_owner_drop() {
80        let (owner, tracker) = ConnectorTaskOwner::new();
81        let admission = owner.admission();
82        let admission_clone = admission.clone();
83        let admitted = admission.track().expect("live generation");
84
85        drop(owner);
86
87        assert!(admission.track().is_none());
88        assert!(admission_clone.track().is_none());
89        assert!(!tracker.is_terminated());
90        drop(admitted);
91        assert!(tracker.is_terminated());
92    }
93
94    #[test]
95    fn connector_task_admission_does_not_retain_generation_state() {
96        let (owner, tracker) = ConnectorTaskOwner::new();
97        let admission = owner.admission();
98
99        drop(owner);
100        drop(tracker);
101
102        assert!(admission.inner.upgrade().is_none());
103        assert!(admission.track().is_none());
104    }
105
106    #[tokio::test]
107    async fn connector_task_wait_wakes_every_tracker_clone() {
108        let (owner, tracker) = ConnectorTaskOwner::new();
109        let guard = owner.track().expect("live generation");
110        let first = tokio::spawn({
111            let tracker = tracker.clone();
112            async move { tracker.wait_terminated().await }
113        });
114        let second = tokio::spawn({
115            let tracker = tracker.clone();
116            async move { tracker.wait_terminated().await }
117        });
118
119        drop(owner);
120        assert!(!tracker.is_terminated());
121        drop(guard);
122
123        tokio::time::timeout(std::time::Duration::from_secs(1), async {
124            first.await.expect("first waiter task");
125            second.await.expect("second waiter task");
126        })
127        .await
128        .expect("tracker waiters must wake");
129        tracker.wait_terminated().await;
130    }
131
132    #[test]
133    fn test_source_batch() {
134        let batch = SourceBatch::new(test_batch(10));
135        assert_eq!(batch.num_rows(), 10);
136        assert!(batch.row_positions().is_none());
137        assert!(batch.mutations().is_none());
138    }
139
140    #[test]
141    fn source_row_positions_reject_nulls_and_misalignment() {
142        let null_partition = BinaryArray::from(vec![Some(&b"p0"[..]), None]);
143        let order = BinaryArray::from(vec![&b"0"[..], &b"1"[..]]);
144        let sub_offset = UInt32Array::from(vec![0, 0]);
145        assert!(
146            SourceRowPositions::try_new(null_partition, order.clone(), sub_offset.clone()).is_err()
147        );
148
149        let positions = SourceRowPositions::try_new(
150            BinaryArray::from(vec![&b"p0"[..], &b"p0"[..]]),
151            order,
152            sub_offset,
153        )
154        .unwrap();
155        assert!(SourceBatch::positioned(test_batch(1), positions).is_err());
156    }
157
158    #[test]
159    fn source_batch_validates_and_canonicalizes_mutations() {
160        let mixed = SourceBatch::new(test_batch(2))
161            .with_mutations(vec![SourceMutation::Put, SourceMutation::Tombstone])
162            .unwrap();
163        assert_eq!(
164            mixed.mutations(),
165            Some(&[SourceMutation::Put, SourceMutation::Tombstone][..])
166        );
167
168        let puts = SourceBatch::new(test_batch(2))
169            .with_mutations(vec![SourceMutation::Put; 2])
170            .unwrap();
171        assert!(puts.mutations().is_none());
172        assert!(SourceBatch::new(test_batch(2))
173            .with_mutations(vec![SourceMutation::Tombstone])
174            .is_err());
175    }
176
177    #[test]
178    fn source_metadata_round_trip_is_sparse_and_zero_copy() {
179        let records = test_batch(2);
180        assert!(source_row_positions(&records).unwrap().is_none());
181        let positioned_schema = schema_with_source_row_positions(&records.schema()).unwrap();
182        let mutation_schema =
183            schema_with_source_mutations_and_row_positions(&records.schema()).unwrap();
184        let positions = SourceRowPositions::try_new(
185            BinaryArray::from(vec![&b"p0"[..], &b"p0"[..]]),
186            BinaryArray::from(vec![&b"0"[..], &b"1"[..]]),
187            UInt32Array::from(vec![0, 0]),
188        )
189        .unwrap();
190        let encoded = SourceBatch::positioned(records.clone(), positions.clone())
191            .unwrap()
192            .with_mutations(vec![SourceMutation::Put, SourceMutation::Tombstone])
193            .unwrap()
194            .into_records_with_metadata(
195                SourceRowPositionCapability::OrderedDeterministic,
196                &positioned_schema,
197                &mutation_schema,
198            )
199            .unwrap();
200        let mutations = source_mutations(&encoded).unwrap().unwrap();
201        assert_eq!(mutations.len(), 2);
202        assert!(!mutations.is_empty());
203        assert_eq!(mutations.get(0), Some(SourceMutation::Put));
204        assert_eq!(mutations.get(1), Some(SourceMutation::Tombstone));
205        let row_positions = source_row_positions(&encoded).unwrap().unwrap();
206        assert_eq!(row_positions.len(), 2);
207        assert!(!row_positions.is_empty());
208        assert_eq!(
209            row_positions.get(1),
210            Some(SourceRowPositionRef {
211                partition: b"p0",
212                order_key: b"1",
213                sub_offset: 0,
214            })
215        );
216        assert_eq!(row_positions.get(2), None);
217        assert_eq!(
218            encoded.schema().field(records.num_columns()).name(),
219            SOURCE_MUTATION_COLUMN
220        );
221
222        let positioned = strip_source_mutations(&encoded).unwrap();
223        assert_eq!(positioned.schema(), positioned_schema);
224        assert!(Arc::ptr_eq(positioned.column(0), records.column(0)));
225
226        let routed_put = encoded.slice(0, 1);
227        assert!(source_mutations(&routed_put).is_err());
228        assert_eq!(
229            source_mutations_routed(&routed_put)
230                .unwrap()
231                .unwrap()
232                .get(0),
233            Some(SourceMutation::Put)
234        );
235        let routed_visible = Arc::clone(routed_put.column(0));
236        let routed_positioned = strip_source_mutations_routed(&routed_put).unwrap();
237        assert!(Arc::ptr_eq(&routed_visible, routed_positioned.column(0)));
238
239        let stripped = strip_source_row_positions(&encoded).unwrap();
240        assert_eq!(stripped.schema(), records.schema());
241        assert_eq!(stripped.num_rows(), records.num_rows());
242        assert!(Arc::ptr_eq(stripped.column(0), records.column(0)));
243
244        let puts = SourceBatch::positioned(records.clone(), positions)
245            .unwrap()
246            .with_mutations(vec![SourceMutation::Put; 2])
247            .unwrap()
248            .into_records_with_metadata(
249                SourceRowPositionCapability::OrderedDeterministic,
250                &positioned_schema,
251                &mutation_schema,
252            )
253            .unwrap();
254        assert!(Arc::ptr_eq(&puts.schema(), &positioned_schema));
255        assert!(puts.column_by_name(SOURCE_MUTATION_COLUMN).is_none());
256        assert!(source_mutations(&puts).unwrap().is_none());
257    }
258
259    #[test]
260    fn source_metadata_rejects_collisions_and_malformed_batches() {
261        let collision = Arc::new(Schema::new(vec![Field::new(
262            "__SOURCE_MUTATION",
263            DataType::UInt8,
264            false,
265        )]));
266        assert!(schema_with_source_row_positions(&collision).is_err());
267        assert!(schema_with_source_mutations_and_row_positions(&collision).is_err());
268
269        let records = test_batch(2);
270        let positioned_schema = schema_with_source_row_positions(&records.schema()).unwrap();
271        let mutation_schema =
272            schema_with_source_mutations_and_row_positions(&records.schema()).unwrap();
273        let positions = SourceRowPositions::try_new(
274            BinaryArray::from(vec![&b"p0"[..], &b"p0"[..]]),
275            BinaryArray::from(vec![&b"0"[..], &b"1"[..]]),
276            UInt32Array::from(vec![0, 0]),
277        )
278        .unwrap();
279        let encoded = SourceBatch::positioned(records.clone(), positions)
280            .unwrap()
281            .with_mutations(vec![SourceMutation::Put, SourceMutation::Tombstone])
282            .unwrap()
283            .into_records_with_metadata(
284                SourceRowPositionCapability::OrderedDeterministic,
285                &positioned_schema,
286                &mutation_schema,
287            )
288            .unwrap();
289        let mutation_index = records.num_columns();
290
291        let malformed = |field: Field, array: ArrayRef| {
292            let mut fields = encoded.schema().fields().to_vec();
293            fields[mutation_index] = Arc::new(field);
294            let mut columns = encoded.columns().to_vec();
295            columns[mutation_index] = array;
296            RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap()
297        };
298        let wrong_type = malformed(
299            Field::new(SOURCE_MUTATION_COLUMN, DataType::Int64, false),
300            Arc::new(Int64Array::from(vec![0, 1])),
301        );
302        assert!(source_mutations(&wrong_type).is_err());
303
304        let null = malformed(
305            Field::new(SOURCE_MUTATION_COLUMN, DataType::UInt8, true),
306            Arc::new(UInt8Array::from(vec![Some(0), None])),
307        );
308        assert!(strip_source_mutations(&null).is_err());
309
310        let unknown = malformed(
311            Field::new(SOURCE_MUTATION_COLUMN, DataType::UInt8, false),
312            Arc::new(UInt8Array::from(vec![0, 2])),
313        );
314        assert!(source_mutations(&unknown).is_err());
315        assert!(strip_source_mutations(&unknown).is_err());
316
317        let all_put = malformed(
318            Field::new(SOURCE_MUTATION_COLUMN, DataType::UInt8, false),
319            Arc::new(UInt8Array::from(vec![0, 0])),
320        );
321        assert!(strip_source_mutations(&all_put).is_err());
322
323        let mut fields = encoded.schema().fields().to_vec();
324        let mutation_field = fields.remove(mutation_index);
325        fields.push(mutation_field);
326        let mut columns = encoded.columns().to_vec();
327        let mutation_column = columns.remove(mutation_index);
328        columns.push(mutation_column);
329        let misplaced = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap();
330        assert!(source_mutations(&misplaced).is_err());
331    }
332
333    #[test]
334    fn test_write_result() {
335        let result = WriteResult::new(100, 5000);
336        assert_eq!(result.records_written, 100);
337        assert_eq!(result.bytes_written, 5000);
338    }
339
340    #[test]
341    fn source_drain_request_requires_canonical_round() {
342        let round = laminar_core::checkpoint::AssignmentDrainId {
343            predecessor_version: 7,
344            target_version: 8,
345            digest: [9; 32],
346        };
347        assert_eq!(SourceDrainRequest::new(round).unwrap().round, round);
348        assert!(
349            SourceDrainRequest::new(laminar_core::checkpoint::AssignmentDrainId {
350                predecessor_version: 8,
351                target_version: 8,
352                digest: [9; 32],
353            })
354            .is_err()
355        );
356    }
357
358    #[test]
359    fn source_contract_defaults_fail_closed() {
360        let contract = SourceContract::default();
361        assert_eq!(contract.consistency, SourceConsistency::Ephemeral);
362        assert_eq!(contract.topology, SourceTopology::Singleton);
363        assert_eq!(contract.input_mode, SourceInputMode::AppendOnly);
364        assert_eq!(
365            contract.row_positions,
366            SourceRowPositionCapability::Unavailable
367        );
368        assert!(!contract.supports_replay());
369        assert!(!contract.requires_checkpointing());
370        assert!(!contract.is_exact_delivery_certified());
371    }
372
373    #[test]
374    fn commit_coupled_sources_are_replayable_and_require_checkpoints() {
375        let contract = SourceContract::new(
376            SourceConsistency::CommitCoupled,
377            SourceTopology::NodeLocalIngress,
378            SourceInputMode::FullChangelog,
379        );
380        assert!(contract.supports_replay());
381        assert!(contract.requires_checkpointing());
382        assert_eq!(contract.input_mode, SourceInputMode::FullChangelog);
383    }
384
385    #[test]
386    fn source_start_rejects_split_and_zero_resume_before_connector_start() {
387        use laminar_core::checkpoint::CheckpointAttempt;
388
389        for attempt in [CheckpointAttempt::new(7, 8), CheckpointAttempt::new(0, 0)] {
390            let error = SourceStart::new(
391                ConnectorConfig::new("test"),
392                SourcePosition::Resume {
393                    attempt,
394                    checkpoint: SourceCheckpoint::new(),
395                },
396                DeliveryGuarantee::AtLeastOnce,
397            )
398            .unwrap_err();
399            assert!(matches!(
400                error,
401                ConnectorError::ConfigurationError(message)
402                    if message.contains("one nonzero canonical checkpoint ID")
403            ));
404        }
405    }
406
407    #[test]
408    fn source_start_accepts_initial_and_exposes_validated_parts() {
409        let mut config = ConnectorConfig::new("test");
410        config.set("endpoint", "local");
411        let request = SourceStart::new(
412            config,
413            SourcePosition::Initial,
414            DeliveryGuarantee::BestEffort,
415        )
416        .unwrap();
417
418        let (config, position, delivery) = request.into_parts();
419        assert_eq!(config.get("endpoint"), Some("local"));
420        assert!(matches!(position, SourcePosition::Initial));
421        assert_eq!(delivery, DeliveryGuarantee::BestEffort);
422    }
423
424    #[test]
425    fn sink_contract_defaults_fail_closed() {
426        let contract = SinkContract::default();
427        assert_eq!(contract.consistency, SinkConsistency::Ephemeral);
428        assert_eq!(contract.topology, SinkTopology::Singleton);
429        assert_eq!(contract.input_mode, SinkInputMode::AppendOnly);
430        assert!(!contract.input_mode.accepts_full_changelog());
431    }
432
433    #[test]
434    fn coordinated_namespace_is_bounded_stable_and_sink_scoped() {
435        use laminar_core::checkpoint::checkpoint_manifest::PipelineIdentity;
436        const DEPLOYMENT: &str = "018f0000-0000-7000-8000-000000000001";
437
438        let first =
439            CoordinatedCommitNamespace::try_new(PipelineIdentity::empty(), DEPLOYMENT, "orders")
440                .unwrap();
441        let same =
442            CoordinatedCommitNamespace::try_new(PipelineIdentity::empty(), DEPLOYMENT, "orders")
443                .unwrap();
444        let other =
445            CoordinatedCommitNamespace::try_new(PipelineIdentity::empty(), DEPLOYMENT, "audit")
446                .unwrap();
447        let other_deployment = CoordinatedCommitNamespace::try_new(
448            PipelineIdentity::empty(),
449            "018f0000-0000-7000-8000-000000000002",
450            "orders",
451        )
452        .unwrap();
453
454        assert_eq!(first.external_key(), same.external_key());
455        assert_eq!(
456            first.external_key(),
457            "ldb-c3-ab2cbc74e7e5dc7bbbde532272f6e7998fba65710e88151f4d1dd24e97fa0b56"
458        );
459        assert_ne!(first.external_key(), other.external_key());
460        assert_ne!(first.external_key(), other_deployment.external_key());
461        assert_eq!(first.external_key().len(), "ldb-c3-".len() + 64);
462        assert!(first
463            .external_key()
464            .bytes()
465            .all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'-'));
466    }
467
468    #[test]
469    fn coordinated_namespace_rejects_ambiguous_identity() {
470        const DEPLOYMENT: &str = "018f0000-0000-7000-8000-000000000001";
471
472        use laminar_core::checkpoint::checkpoint_manifest::{
473            PipelineIdentity, PIPELINE_IDENTITY_VERSION,
474        };
475        let malformed = PipelineIdentity {
476            canonical_version: PIPELINE_IDENTITY_VERSION,
477            sha256: "NOT-A-DIGEST".into(),
478        };
479        assert!(CoordinatedCommitNamespace::try_new(malformed, DEPLOYMENT, "orders").is_err());
480        let future_version = PipelineIdentity {
481            canonical_version: PIPELINE_IDENTITY_VERSION + 1,
482            sha256: PipelineIdentity::empty().sha256,
483        };
484        assert!(CoordinatedCommitNamespace::try_new(future_version, DEPLOYMENT, "orders").is_err());
485        assert!(
486            CoordinatedCommitNamespace::try_new(PipelineIdentity::empty(), DEPLOYMENT, "").is_err()
487        );
488        assert!(CoordinatedCommitNamespace::try_new(
489            PipelineIdentity::empty(),
490            "not-a-uuid",
491            "orders"
492        )
493        .is_err());
494    }
495
496    #[test]
497    fn coordinated_batch_fingerprint_covers_the_exact_ordered_cut() {
498        use laminar_core::checkpoint::checkpoint_manifest::PipelineIdentity;
499        use laminar_core::checkpoint::CheckpointAttempt;
500
501        let namespace = CoordinatedCommitNamespace::try_new(
502            PipelineIdentity::empty(),
503            "018f0000-0000-7000-8000-000000000001",
504            "orders",
505        )
506        .unwrap();
507        let attempt = CheckpointAttempt::new(8, 108);
508        let batch = CoordinatedCommitBatch {
509            namespace,
510            expected_predecessor: CoordinatedCommitCursor {
511                checkpoint_id: 107,
512                fencing_token: 3,
513            },
514            fencing_token: 4,
515            target: attempt,
516            entries: vec![CoordinatedCommitPayload {
517                attempt,
518                participant_id: 7,
519                payload: None,
520            }],
521        };
522        let expected = batch.exact_fingerprint();
523        assert_eq!(expected, batch.clone().exact_fingerprint());
524        assert_eq!(
525            expected,
526            [
527                0x58, 0x5e, 0xa2, 0xf8, 0xdf, 0x8c, 0xd0, 0x3a, 0xca, 0x01, 0xed, 0x95, 0x7e, 0xe6,
528                0xc1, 0xf0, 0x36, 0x41, 0x28, 0xbd, 0xd6, 0xe1, 0xb6, 0xfa, 0x21, 0xbe, 0xf7, 0x4f,
529                0x97, 0xf7, 0xa5, 0xc4,
530            ]
531        );
532
533        let mut variants = Vec::new();
534        let mut variant = batch.clone();
535        variant.namespace = CoordinatedCommitNamespace::try_new(
536            PipelineIdentity::empty(),
537            "018f0000-0000-7000-8000-000000000001",
538            "audit",
539        )
540        .unwrap();
541        variants.push(variant);
542        let mut variant = batch.clone();
543        variant.expected_predecessor.checkpoint_id -= 1;
544        variants.push(variant);
545        let mut variant = batch.clone();
546        variant.fencing_token += 1;
547        variants.push(variant);
548        let mut variant = batch.clone();
549        variant.target.epoch += 1;
550        variants.push(variant);
551        let mut variant = batch.clone();
552        variant.entries[0].attempt.checkpoint_id += 1;
553        variants.push(variant);
554        let mut variant = batch.clone();
555        variant.entries[0].participant_id += 1;
556        variants.push(variant);
557        let mut variant = batch;
558        variant.entries[0].payload = Some(Vec::new());
559        variants.push(variant);
560
561        assert!(variants
562            .into_iter()
563            .all(|variant| variant.exact_fingerprint() != expected));
564    }
565
566    fn valid_coordinated_batch() -> CoordinatedCommitBatch {
567        use laminar_core::checkpoint::checkpoint_manifest::PipelineIdentity;
568        use laminar_core::checkpoint::CheckpointAttempt;
569
570        let target = CheckpointAttempt::canonical(102);
571        CoordinatedCommitBatch {
572            namespace: CoordinatedCommitNamespace::try_new(
573                PipelineIdentity::empty(),
574                "018f0000-0000-7000-8000-000000000001",
575                "orders",
576            )
577            .unwrap(),
578            expected_predecessor: CoordinatedCommitCursor {
579                checkpoint_id: 101,
580                fencing_token: 1,
581            },
582            fencing_token: 2,
583            target,
584            entries: vec![CoordinatedCommitPayload {
585                attempt: target,
586                participant_id: 1,
587                payload: None,
588            }],
589        }
590    }
591
592    #[test]
593    fn coordinated_abort_batch_accepts_only_one_exact_attempt() {
594        use laminar_core::checkpoint::CheckpointAttempt;
595
596        let commit = valid_coordinated_batch();
597        let mut batch = CoordinatedAbortBatch {
598            namespace: commit.namespace,
599            fencing_token: commit.fencing_token,
600            target: commit.target,
601            entries: commit
602                .entries
603                .into_iter()
604                .map(|entry| CoordinatedAbortEntry {
605                    attempt: entry.attempt,
606                    participant_id: entry.participant_id,
607                    descriptor: CoordinatedAbortDescriptor::Prepared(entry.payload),
608                    artifact_intent: None,
609                })
610                .collect(),
611        };
612        batch.validate_shape().unwrap();
613
614        batch.entries.insert(
615            0,
616            CoordinatedAbortEntry {
617                attempt: CheckpointAttempt::canonical(101),
618                participant_id: 1,
619                descriptor: CoordinatedAbortDescriptor::Open,
620                artifact_intent: None,
621            },
622        );
623        assert!(batch
624            .validate_shape()
625            .unwrap_err()
626            .contains("exact target attempt"));
627        batch.entries.remove(0);
628        batch.fencing_token = 0;
629        assert!(batch
630            .validate_shape()
631            .unwrap_err()
632            .contains("fencing token"));
633    }
634
635    #[test]
636    fn coordinated_batch_rejects_noncanonical_target_before_other_shape_checks() {
637        use laminar_core::checkpoint::CheckpointAttempt;
638
639        for target in [
640            CheckpointAttempt::new(102, 103),
641            CheckpointAttempt::new(0, 0),
642        ] {
643            let mut batch = valid_coordinated_batch();
644            batch.target = target;
645            let error = batch.validate_shape().unwrap_err();
646            assert!(
647                error.contains("target must use one nonzero canonical checkpoint ID"),
648                "unexpected validation error: {error}"
649            );
650        }
651    }
652
653    #[test]
654    fn coordinated_batch_rejects_a_mutated_namespace() {
655        let mut batch = valid_coordinated_batch();
656        batch.namespace.sink_id.clear();
657
658        assert!(batch
659            .validate_shape()
660            .unwrap_err()
661            .contains("sink id cannot be empty"));
662    }
663
664    #[test]
665    fn coordinated_batch_rejects_noncanonical_entry_before_other_shape_checks() {
666        use laminar_core::checkpoint::CheckpointAttempt;
667
668        for attempt in [
669            CheckpointAttempt::new(101, 102),
670            CheckpointAttempt::new(0, 0),
671        ] {
672            let mut batch = valid_coordinated_batch();
673            batch.entries[0].attempt = attempt;
674            let error = batch.validate_shape().unwrap_err();
675            assert!(
676                error.contains("canonical checkpoint ID"),
677                "unexpected validation error: {error}"
678            );
679        }
680        let mut batch = valid_coordinated_batch();
681        batch.entries[0].participant_id = 0;
682        assert!(batch
683            .validate_shape()
684            .unwrap_err()
685            .contains("nonzero participant"));
686    }
687
688    #[test]
689    fn coordinated_batch_rejects_cursor_rollback_and_unproven_overlap() {
690        use laminar_core::checkpoint::checkpoint_manifest::PipelineIdentity;
691        use laminar_core::checkpoint::CheckpointAttempt;
692
693        let first = CheckpointAttempt::canonical(108);
694        let target = CheckpointAttempt::canonical(110);
695        let batch = CoordinatedCommitBatch {
696            namespace: CoordinatedCommitNamespace::try_new(
697                PipelineIdentity::empty(),
698                "018f0000-0000-7000-8000-000000000001",
699                "orders",
700            )
701            .unwrap(),
702            expected_predecessor: CoordinatedCommitCursor {
703                checkpoint_id: 107,
704                fencing_token: 3,
705            },
706            fencing_token: 4,
707            target,
708            entries: vec![
709                CoordinatedCommitPayload {
710                    attempt: first,
711                    participant_id: 1,
712                    payload: None,
713                },
714                CoordinatedCommitPayload {
715                    attempt: target,
716                    participant_id: 1,
717                    payload: None,
718                },
719            ],
720        };
721
722        let cursor = |checkpoint_id, fencing_token| {
723            Some(CoordinatedCommitCursor {
724                checkpoint_id,
725                fencing_token,
726            })
727        };
728        assert!(batch.validate_observed_cursor(cursor(106, 3)).is_err());
729        assert!(batch.validate_observed_cursor(cursor(109, 3)).is_err());
730        assert!(batch.validate_observed_cursor(cursor(107, 2)).is_err());
731        assert!(batch.validate_observed_cursor(cursor(107, 3)).is_ok());
732        assert!(batch.validate_observed_cursor(cursor(108, 3)).is_ok());
733        assert!(batch.validate_observed_cursor(cursor(110, 4)).is_ok());
734        assert!(batch.validate_observed_cursor(cursor(110, 3)).is_err());
735        assert!(batch.validate_observed_cursor(cursor(108, 5)).is_err());
736    }
737
738    #[test]
739    fn coordinated_batch_requires_unique_canonical_attempt_participants() {
740        use laminar_core::checkpoint::checkpoint_manifest::PipelineIdentity;
741        use laminar_core::checkpoint::CheckpointAttempt;
742
743        let namespace = CoordinatedCommitNamespace::try_new(
744            PipelineIdentity::empty(),
745            "018f0000-0000-7000-8000-000000000001",
746            "orders",
747        )
748        .unwrap();
749        let target = CheckpointAttempt::canonical(102);
750        let batch = |entries| CoordinatedCommitBatch {
751            namespace: namespace.clone(),
752            expected_predecessor: CoordinatedCommitCursor {
753                checkpoint_id: 100,
754                fencing_token: 1,
755            },
756            fencing_token: 2,
757            target,
758            entries,
759        };
760        let payload = |attempt, participant_id| CoordinatedCommitPayload {
761            attempt,
762            participant_id,
763            payload: None,
764        };
765
766        let duplicate = batch(vec![payload(target, 1), payload(target, 1)]);
767        assert!(duplicate
768            .validate_shape()
769            .unwrap_err()
770            .contains("duplicate"));
771
772        let out_of_order = batch(vec![payload(target, 2), payload(target, 1)]);
773        assert!(out_of_order
774            .validate_shape()
775            .unwrap_err()
776            .contains("out-of-order"));
777
778        let noncanonical = batch(vec![
779            payload(CheckpointAttempt::new(3, 101), 1),
780            payload(target, 2),
781        ]);
782        assert!(noncanonical
783            .validate_shape()
784            .unwrap_err()
785            .contains("canonical checkpoint ID"));
786    }
787
788    #[test]
789    fn coordinated_batch_entry_limit_accepts_max_and_rejects_max_plus_one() {
790        use laminar_core::checkpoint::checkpoint_manifest::PipelineIdentity;
791        use laminar_core::checkpoint::CheckpointAttempt;
792
793        let namespace = CoordinatedCommitNamespace::try_new(
794            PipelineIdentity::empty(),
795            "018f0000-0000-7000-8000-000000000001",
796            "orders",
797        )
798        .unwrap();
799        let target = CheckpointAttempt::canonical(101);
800        let make_batch = |count: usize| CoordinatedCommitBatch {
801            namespace: namespace.clone(),
802            expected_predecessor: CoordinatedCommitCursor {
803                checkpoint_id: 0,
804                fencing_token: 0,
805            },
806            fencing_token: 1,
807            target,
808            entries: (1..=count)
809                .map(|participant_id| CoordinatedCommitPayload {
810                    attempt: target,
811                    participant_id: participant_id as u64,
812                    payload: None,
813                })
814                .collect(),
815        };
816
817        assert!(make_batch(MAX_COORDINATED_COMMIT_BATCH_ENTRIES - 1)
818            .validate_shape()
819            .is_ok());
820        assert!(make_batch(MAX_COORDINATED_COMMIT_BATCH_ENTRIES)
821            .validate_shape()
822            .is_ok());
823        assert!(make_batch(MAX_COORDINATED_COMMIT_BATCH_ENTRIES + 1)
824            .validate_shape()
825            .is_err());
826    }
827
828    struct DefaultPreCommitSink {
829        coordinated: bool,
830    }
831
832    #[async_trait]
833    impl SinkConnector for DefaultPreCommitSink {
834        async fn open(&mut self, _config: &ConnectorConfig) -> Result<(), ConnectorError> {
835            Ok(())
836        }
837        async fn write_batch(
838            &mut self,
839            _batch: &RecordBatch,
840        ) -> Result<WriteResult, ConnectorError> {
841            Ok(WriteResult::new(0, 0))
842        }
843        fn schema(&self) -> SchemaRef {
844            test_schema()
845        }
846        fn suggested_write_timeout(&self) -> std::time::Duration {
847            std::time::Duration::from_secs(5)
848        }
849        fn as_coordinated_committer(&self) -> Option<&dyn CoordinatedCommitter> {
850            self.coordinated
851                .then_some(self as &dyn CoordinatedCommitter)
852        }
853        async fn close(&mut self) -> Result<(), ConnectorError> {
854            Ok(())
855        }
856    }
857
858    #[async_trait]
859    impl CoordinatedCommitter for DefaultPreCommitSink {
860        async fn commit_aggregated(
861            &self,
862            _batch: CoordinatedCommitBatch,
863            _context: CoordinatedCommitContext,
864        ) -> Result<(), ConnectorError> {
865            Ok(())
866        }
867
868        async fn committed_cursor(
869            &self,
870            _namespace: &CoordinatedCommitNamespace,
871        ) -> Result<Option<CoordinatedCommitCursor>, ConnectorError> {
872            Ok(None)
873        }
874    }
875
876    #[tokio::test]
877    async fn default_pre_commit_rejects_coordinated_sink() {
878        let mut sink = DefaultPreCommitSink { coordinated: true };
879        assert!(matches!(
880            sink.pre_commit(1).await,
881            Err(ConnectorError::ConfigurationError(_))
882        ));
883    }
884
885    #[tokio::test]
886    async fn default_pre_commit_ok_for_non_coordinated_sink() {
887        let mut sink = DefaultPreCommitSink { coordinated: false };
888        assert!(matches!(sink.pre_commit(1).await, Ok(None)));
889    }
890}