1use std::{
6 fmt,
7 time::{Duration, Instant},
8};
9
10use arrow_schema::Schema;
11pub use arrow_tiberius::WriteOptions as MssqlWriteOptions;
12use async_trait::async_trait;
13use datafusion::arrow::record_batch::RecordBatch;
14use futures_util::{
15 Stream, StreamExt,
16 io::{AsyncRead, AsyncWrite},
17 pin_mut,
18};
19
20use crate::{
21 DeltaFunnelError, PhaseTimingReport, ReportReasonCode, RowCount, ValidationStatus,
22 observability,
23 report::sql_server::{
24 MssqlBatchShapingReport, MssqlOutputBatchValidationReport, MssqlTargetCleanupStatus,
25 MssqlWriteFailureContext, MssqlWriteReport, MssqlWriteReportMetrics,
26 write::MssqlWriteDiagnostic,
27 },
28};
29
30const POLL_BATCH_STREAM_PHASE: &str = "poll_batch_stream";
31const VALIDATE_BATCH_SCHEMA_PHASE: &str = "validate_batch_schema";
32const WRITE_BATCH_PHASE: &str = "write_batch";
33const FINALIZE_PHASE: &str = "finalize";
34const VALIDATION_PHASE: &str = "validation";
35
36use super::{
37 LoadMode, MssqlLifecycleExecutionGuardrail, MssqlPreparedTarget, MssqlPreparedTargetAction,
38 MssqlTargetOutputPlan,
39};
40
41#[async_trait]
43pub(crate) trait MssqlBulkLoadWriter: Sized + Send {
44 async fn write_batch(
46 &mut self,
47 batch: &RecordBatch,
48 ) -> Result<arrow_tiberius::WriteStats, arrow_tiberius::Error>;
49
50 async fn finish(self) -> Result<arrow_tiberius::WriteStats, arrow_tiberius::Error>;
52}
53
54#[async_trait]
55impl<'client, S> MssqlBulkLoadWriter for arrow_tiberius::BulkWriter<'client, S>
56where
57 S: AsyncRead + AsyncWrite + Unpin + Send,
58{
59 async fn write_batch(
60 &mut self,
61 batch: &RecordBatch,
62 ) -> Result<arrow_tiberius::WriteStats, arrow_tiberius::Error> {
63 arrow_tiberius::BulkWriter::write_batch(self, batch).await
64 }
65
66 async fn finish(self) -> Result<arrow_tiberius::WriteStats, arrow_tiberius::Error> {
67 arrow_tiberius::BulkWriter::finish(self).await
68 }
69}
70
71#[async_trait]
72impl MssqlBulkLoadWriter for arrow_tiberius::ConnectedBulkWriter<'_> {
73 async fn write_batch(
74 &mut self,
75 batch: &RecordBatch,
76 ) -> Result<arrow_tiberius::WriteStats, arrow_tiberius::Error> {
77 arrow_tiberius::ConnectedBulkWriter::write_batch(self, batch).await
78 }
79
80 async fn finish(self) -> Result<arrow_tiberius::WriteStats, arrow_tiberius::Error> {
81 arrow_tiberius::ConnectedBulkWriter::finish(self).await
82 }
83}
84
85#[allow(dead_code)]
87#[async_trait]
88pub(crate) trait MssqlBulkWriterFactory: Send {
89 type Writer: MssqlBulkLoadWriter;
91
92 async fn initialize(
94 self,
95 request: MssqlBulkWriterInitializationRequest,
96 ) -> Result<Self::Writer, arrow_tiberius::Error>;
97}
98
99#[allow(dead_code)]
101pub(crate) struct MssqlConnectedBulkWriterFactory<'client> {
102 client: &'client mut arrow_tiberius::ConnectedMssqlClient,
103}
104
105impl<'client> MssqlConnectedBulkWriterFactory<'client> {
106 #[must_use]
108 pub(crate) const fn new(client: &'client mut arrow_tiberius::ConnectedMssqlClient) -> Self {
109 Self { client }
110 }
111}
112
113#[async_trait]
114impl<'client> MssqlBulkWriterFactory for MssqlConnectedBulkWriterFactory<'client> {
115 type Writer = arrow_tiberius::ConnectedBulkWriter<'client>;
116
117 async fn initialize(
118 self,
119 request: MssqlBulkWriterInitializationRequest,
120 ) -> Result<Self::Writer, arrow_tiberius::Error> {
121 let MssqlBulkWriterInitializationRequest {
122 table,
123 mappings,
124 options,
125 ..
126 } = request;
127
128 self.client.bulk_writer(table, mappings, options).await
129 }
130}
131
132#[derive(Debug, Clone, PartialEq, Eq)]
134pub(crate) struct MssqlBulkWriterInitializationRequest {
135 output_name: String,
136 table: arrow_tiberius::TableName,
137 mappings: Vec<arrow_tiberius::SchemaMapping>,
138 options: MssqlWriteOptions,
139 prepared_action: MssqlPreparedTargetAction,
140 cleanup: MssqlTargetCleanupStatus,
141}
142
143#[allow(dead_code)]
144impl MssqlBulkWriterInitializationRequest {
145 pub(crate) fn from_prepared_target(
147 output_plan: &MssqlTargetOutputPlan,
148 prepared_target: &MssqlPreparedTarget,
149 options: MssqlWriteOptions,
150 ) -> Result<Self, DeltaFunnelError> {
151 ensure_prepared_target_matches_output_plan(output_plan, prepared_target)?;
152
153 Ok(Self {
154 output_name: output_plan.output_name().to_owned(),
155 table: prepared_target.table_name().clone(),
156 mappings: output_plan.schema_mappings().to_vec(),
157 options,
158 prepared_action: prepared_target.report().action(),
159 cleanup: prepared_target.report().cleanup(),
160 })
161 }
162
163 #[must_use]
165 pub(crate) fn output_name(&self) -> &str {
166 &self.output_name
167 }
168
169 #[must_use]
171 pub(crate) const fn table(&self) -> &arrow_tiberius::TableName {
172 &self.table
173 }
174
175 #[must_use]
177 pub(crate) fn mappings(&self) -> &[arrow_tiberius::SchemaMapping] {
178 &self.mappings
179 }
180
181 #[must_use]
183 pub(crate) const fn options(&self) -> MssqlWriteOptions {
184 self.options
185 }
186
187 #[must_use]
189 pub(crate) const fn prepared_action(&self) -> MssqlPreparedTargetAction {
190 self.prepared_action
191 }
192
193 #[must_use]
195 pub(crate) const fn cleanup(&self) -> MssqlTargetCleanupStatus {
196 self.cleanup
197 }
198}
199
200#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
202pub enum MssqlWritePhase {
203 Connect,
205 PrepareTargetLifecycle,
207 InitializeWriter,
209 PollBatchStream,
211 ValidateBatchSchema,
213 WriteBatch,
215 Finalize,
217 Validation,
219 SwapTarget,
221 Cleanup,
223}
224
225impl fmt::Display for MssqlWritePhase {
226 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
227 formatter.write_str(match self {
228 Self::Connect => "connect",
229 Self::PrepareTargetLifecycle => "prepare target lifecycle",
230 Self::InitializeWriter => "initialize writer",
231 Self::PollBatchStream => "poll batch stream",
232 Self::ValidateBatchSchema => "validate batch schema",
233 Self::WriteBatch => "write batch",
234 Self::Finalize => "finalize",
235 Self::Validation => "validation",
236 Self::SwapTarget => "swap target",
237 Self::Cleanup => "cleanup",
238 })
239 }
240}
241
242#[must_use]
244pub fn default_mssql_write_options() -> MssqlWriteOptions {
245 MssqlWriteOptions {
246 backend: arrow_tiberius::WriteBackend::DirectRawBulk,
247 ..MssqlWriteOptions::default()
248 }
249}
250
251#[must_use]
253pub fn mssql_write_options_for_output_plan(
254 output_plan: &MssqlTargetOutputPlan,
255) -> MssqlWriteOptions {
256 MssqlWriteOptions {
257 plan_options: output_plan.schema_plan_options(),
258 ..default_mssql_write_options()
259 }
260}
261
262#[allow(dead_code)]
264pub(crate) async fn initialize_mssql_bulk_writer<'client>(
265 client: &'client mut arrow_tiberius::ConnectedMssqlClient,
266 output_plan: &MssqlTargetOutputPlan,
267 prepared_target: &MssqlPreparedTarget,
268 options: MssqlWriteOptions,
269) -> Result<arrow_tiberius::ConnectedBulkWriter<'client>, DeltaFunnelError> {
270 initialize_mssql_bulk_writer_with_factory(
271 output_plan,
272 prepared_target,
273 options,
274 MssqlConnectedBulkWriterFactory::new(client),
275 )
276 .await
277}
278
279#[allow(dead_code)]
281pub(crate) async fn initialize_mssql_bulk_writer_with_factory<F>(
282 output_plan: &MssqlTargetOutputPlan,
283 prepared_target: &MssqlPreparedTarget,
284 options: MssqlWriteOptions,
285 factory: F,
286) -> Result<F::Writer, DeltaFunnelError>
287where
288 F: MssqlBulkWriterFactory,
289{
290 let request = MssqlBulkWriterInitializationRequest::from_prepared_target(
291 output_plan,
292 prepared_target,
293 options,
294 )?;
295 let cleanup = request.cleanup();
296 let prepared_action = request.prepared_action();
297
298 factory.initialize(request).await.map_err(|source| {
299 mssql_writer_initialization_error(output_plan, prepared_action, cleanup, source.to_string())
300 })
301}
302
303pub fn validate_mssql_output_schema(
308 output_plan: &MssqlTargetOutputPlan,
309 schema: &Schema,
310) -> Result<MssqlOutputBatchValidationReport, DeltaFunnelError> {
311 arrow_tiberius::validate_arrow_schema_against_mappings(schema, output_plan.schema_mappings())
312 .map_err(|source| {
313 mssql_batch_schema_validation_error(
314 output_plan,
315 source,
316 write_report_metrics(
317 RowCount::unavailable(),
318 MssqlBatchShapingReport::not_started(ReportReasonCode::NotExecuted),
319 MssqlWriteProgress::zero(),
320 false,
321 MssqlTargetCleanupStatus::NotApplicable,
322 ),
323 )
324 })?;
325
326 Ok(MssqlOutputBatchValidationReport::from_output_plan(
327 output_plan,
328 ))
329}
330
331pub fn validate_mssql_output_record_batch(
336 output_plan: &MssqlTargetOutputPlan,
337 batch: &RecordBatch,
338) -> Result<MssqlOutputBatchValidationReport, DeltaFunnelError> {
339 arrow_tiberius::validate_record_batch_schema_against_mappings(
340 batch,
341 output_plan.schema_mappings(),
342 )
343 .map_err(|source| {
344 mssql_batch_schema_validation_error(
345 output_plan,
346 source,
347 write_report_metrics(
348 RowCount::unavailable(),
349 MssqlBatchShapingReport::not_started(ReportReasonCode::NotExecuted),
350 MssqlWriteProgress::zero(),
351 false,
352 MssqlTargetCleanupStatus::NotApplicable,
353 ),
354 )
355 })?;
356
357 Ok(MssqlOutputBatchValidationReport::from_output_plan(
358 output_plan,
359 ))
360}
361
362fn mssql_batch_schema_validation_error(
363 output_plan: &MssqlTargetOutputPlan,
364 source: arrow_tiberius::Error,
365 metrics: MssqlWriteReportMetrics,
366) -> DeltaFunnelError {
367 let diagnostics = batch_validation_diagnostics(&source);
368 DeltaFunnelError::MssqlBatchSchemaValidation {
369 context: Box::new(
370 MssqlWriteFailureContext::from_output_plan_with_metrics(
371 output_plan,
372 MssqlWritePhase::ValidateBatchSchema,
373 metrics,
374 )
375 .with_diagnostics(diagnostics),
376 ),
377 source,
378 }
379}
380
381fn batch_validation_diagnostics(source: &arrow_tiberius::Error) -> Vec<MssqlWriteDiagnostic> {
382 match source {
383 arrow_tiberius::Error::ValueConversion { diagnostics } => diagnostics
384 .all()
385 .iter()
386 .map(MssqlWriteDiagnostic::from_arrow_tiberius)
387 .collect(),
388 _ => Vec::new(),
389 }
390}
391
392fn write_report_metrics(
393 output_row_count: RowCount,
394 batch_shaping: MssqlBatchShapingReport,
395 progress: MssqlWriteProgress,
396 partial_write_possible: bool,
397 cleanup: MssqlTargetCleanupStatus,
398) -> MssqlWriteReportMetrics {
399 MssqlWriteReportMetrics::new(
400 output_row_count,
401 batch_shaping,
402 progress.rows_written,
403 progress.batches_written,
404 progress.elapsed_ms,
405 partial_write_possible,
406 cleanup,
407 )
408}
409
410fn write_failure_report_metrics(
411 output_row_count: RowCount,
412 batch_shaping: MssqlBatchShapingReport,
413 progress: MssqlWriteProgress,
414 partial_write_possible: bool,
415 cleanup: MssqlTargetCleanupStatus,
416) -> MssqlWriteReportMetrics {
417 write_report_metrics(
418 output_row_count,
419 batch_shaping,
420 progress,
421 partial_write_possible,
422 cleanup,
423 )
424 .with_target_validation(
425 RowCount::unavailable(),
426 ValidationStatus::skipped(ReportReasonCode::FailureBeforeValidation),
427 )
428}
429
430#[derive(Default)]
431struct MssqlWriteLoopPhaseTimings {
432 poll_batch_stream: Duration,
433 validate_batch_schema: Duration,
434 write_batch: Duration,
435 validation_attempted: bool,
436 write_attempted: bool,
437}
438
439impl MssqlWriteLoopPhaseTimings {
440 fn add_poll_batch_stream(&mut self, elapsed: Duration) {
441 self.poll_batch_stream = self.poll_batch_stream.saturating_add(elapsed);
442 }
443
444 fn add_validate_batch_schema(&mut self, elapsed: Duration) {
445 self.validation_attempted = true;
446 self.validate_batch_schema = self.validate_batch_schema.saturating_add(elapsed);
447 }
448
449 fn add_write_batch(&mut self, elapsed: Duration) {
450 self.write_attempted = true;
451 self.write_batch = self.write_batch.saturating_add(elapsed);
452 }
453
454 fn completed(&self, finalize_elapsed: Duration) -> Vec<PhaseTimingReport> {
455 vec![
456 PhaseTimingReport::completed(POLL_BATCH_STREAM_PHASE, self.poll_batch_stream),
457 self.validate_batch_schema_completed_or_not_started(),
458 self.write_batch_completed_or_not_started(),
459 PhaseTimingReport::completed(FINALIZE_PHASE, finalize_elapsed),
460 PhaseTimingReport::not_started(VALIDATION_PHASE, ReportReasonCode::NotExecuted),
461 ]
462 }
463
464 fn poll_batch_stream_failed(&self) -> Vec<PhaseTimingReport> {
465 vec![
466 PhaseTimingReport::failed(POLL_BATCH_STREAM_PHASE, self.poll_batch_stream),
467 self.validate_batch_schema_completed_or_not_started(),
468 self.write_batch_completed_or_not_started(),
469 PhaseTimingReport::not_started(FINALIZE_PHASE, ReportReasonCode::NotExecuted),
470 PhaseTimingReport::not_started(
471 VALIDATION_PHASE,
472 ReportReasonCode::FailureBeforeValidation,
473 ),
474 ]
475 }
476
477 fn validate_batch_schema_failed(&self) -> Vec<PhaseTimingReport> {
478 vec![
479 PhaseTimingReport::completed(POLL_BATCH_STREAM_PHASE, self.poll_batch_stream),
480 PhaseTimingReport::failed(VALIDATE_BATCH_SCHEMA_PHASE, self.validate_batch_schema),
481 self.write_batch_completed_or_not_started(),
482 PhaseTimingReport::not_started(FINALIZE_PHASE, ReportReasonCode::NotExecuted),
483 PhaseTimingReport::not_started(
484 VALIDATION_PHASE,
485 ReportReasonCode::FailureBeforeValidation,
486 ),
487 ]
488 }
489
490 fn write_batch_failed(&self) -> Vec<PhaseTimingReport> {
491 vec![
492 PhaseTimingReport::completed(POLL_BATCH_STREAM_PHASE, self.poll_batch_stream),
493 self.validate_batch_schema_completed_or_not_started(),
494 PhaseTimingReport::failed(WRITE_BATCH_PHASE, self.write_batch),
495 PhaseTimingReport::not_started(FINALIZE_PHASE, ReportReasonCode::NotExecuted),
496 PhaseTimingReport::not_started(
497 VALIDATION_PHASE,
498 ReportReasonCode::FailureBeforeValidation,
499 ),
500 ]
501 }
502
503 fn finalize_failed(&self, finalize_elapsed: Duration) -> Vec<PhaseTimingReport> {
504 vec![
505 PhaseTimingReport::completed(POLL_BATCH_STREAM_PHASE, self.poll_batch_stream),
506 self.validate_batch_schema_completed_or_not_started(),
507 self.write_batch_completed_or_not_started(),
508 PhaseTimingReport::failed(FINALIZE_PHASE, finalize_elapsed),
509 PhaseTimingReport::not_started(
510 VALIDATION_PHASE,
511 ReportReasonCode::FailureBeforeValidation,
512 ),
513 ]
514 }
515
516 fn validate_batch_schema_completed_or_not_started(&self) -> PhaseTimingReport {
517 if self.validation_attempted {
518 PhaseTimingReport::completed(VALIDATE_BATCH_SCHEMA_PHASE, self.validate_batch_schema)
519 } else {
520 PhaseTimingReport::not_started(
521 VALIDATE_BATCH_SCHEMA_PHASE,
522 ReportReasonCode::NotExecuted,
523 )
524 }
525 }
526
527 fn write_batch_completed_or_not_started(&self) -> PhaseTimingReport {
528 if self.write_attempted {
529 PhaseTimingReport::completed(WRITE_BATCH_PHASE, self.write_batch)
530 } else {
531 PhaseTimingReport::not_started(WRITE_BATCH_PHASE, ReportReasonCode::NotExecuted)
532 }
533 }
534}
535
536#[allow(dead_code)]
538pub(crate) async fn write_mssql_batches_with_writer<W, S>(
539 output_plan: &MssqlTargetOutputPlan,
540 batches: S,
541 mut writer: W,
542 _options: MssqlWriteOptions,
543) -> Result<MssqlWriteReport, DeltaFunnelError>
544where
545 W: MssqlBulkLoadWriter,
546 S: Stream<Item = Result<RecordBatch, DeltaFunnelError>>,
547{
548 let mut rows_written = 0_u64;
549 let mut batches_written = 0_u64;
550 let mut input_rows = 0_u64;
551 let mut input_batches = 0_u64;
552 let mut shaped_rows = 0_u64;
553 let mut shaped_batches = 0_u64;
554 let cleanup = MssqlTargetCleanupStatus::NotApplicable;
555 let started_at = Instant::now();
556 let mut phase_timings = MssqlWriteLoopPhaseTimings::default();
557 pin_mut!(batches);
558 observability::datafusion_batch_stream_started(
559 output_plan.output_name(),
560 output_plan.target_table(),
561 output_plan.load_mode(),
562 );
563
564 loop {
565 let poll_started_at = Instant::now();
566 let Some(batch) = batches.next().await else {
567 phase_timings.add_poll_batch_stream(poll_started_at.elapsed());
568 break;
569 };
570 phase_timings.add_poll_batch_stream(poll_started_at.elapsed());
571
572 let batch = batch.map_err(|source| {
573 let elapsed_ms = elapsed_ms_since(started_at);
574 let batch_shaping = trace_datafusion_batch_stream_finished(
575 output_plan,
576 MssqlBatchShapingReport::failed(
577 input_batches,
578 input_rows,
579 shaped_batches,
580 shaped_rows,
581 ),
582 );
583 mssql_write_phase_error(
584 output_plan,
585 MssqlWritePhase::PollBatchStream,
586 write_failure_report_metrics(
587 RowCount::partial(input_rows),
588 batch_shaping,
589 MssqlWriteProgress::new(rows_written, batches_written, elapsed_ms),
590 partial_write_possible(output_plan, rows_written, batches_written),
591 cleanup,
592 )
593 .with_phase_timings(phase_timings.poll_batch_stream_failed()),
594 source.to_string(),
595 )
596 })?;
597
598 let row_count = batch_row_count(batch.num_rows());
599 input_rows = input_rows.saturating_add(row_count);
600 input_batches = input_batches.saturating_add(1);
601
602 let validation_started_at = Instant::now();
603 let validation_result = arrow_tiberius::validate_record_batch_schema_against_mappings(
604 &batch,
605 output_plan.schema_mappings(),
606 );
607 phase_timings.add_validate_batch_schema(validation_started_at.elapsed());
608 validation_result.map_err(|source| {
609 let elapsed_ms = elapsed_ms_since(started_at);
610 let batch_shaping = trace_datafusion_batch_stream_finished(
611 output_plan,
612 MssqlBatchShapingReport::failed(
613 input_batches,
614 input_rows,
615 shaped_batches,
616 shaped_rows,
617 ),
618 );
619 mssql_batch_schema_validation_error(
620 output_plan,
621 source,
622 write_failure_report_metrics(
623 RowCount::partial(input_rows),
624 batch_shaping,
625 MssqlWriteProgress::new(rows_written, batches_written, elapsed_ms),
626 partial_write_possible(output_plan, rows_written, batches_written),
627 cleanup,
628 )
629 .with_phase_timings(phase_timings.validate_batch_schema_failed()),
630 )
631 })?;
632
633 let write_batch_started_at = Instant::now();
634 let write_batch_result = MssqlBulkLoadWriter::write_batch(&mut writer, &batch).await;
635 phase_timings.add_write_batch(write_batch_started_at.elapsed());
636 write_batch_result.map_err(|source| {
637 let elapsed_ms = elapsed_ms_since(started_at);
638 let batch_shaping = trace_datafusion_batch_stream_finished(
639 output_plan,
640 MssqlBatchShapingReport::failed(
641 input_batches,
642 input_rows,
643 shaped_batches,
644 shaped_rows,
645 ),
646 );
647 mssql_write_phase_error(
648 output_plan,
649 MssqlWritePhase::WriteBatch,
650 write_failure_report_metrics(
651 RowCount::partial(input_rows),
652 batch_shaping,
653 MssqlWriteProgress::new(rows_written, batches_written, elapsed_ms),
654 partial_write_possible(output_plan, rows_written, batches_written),
655 cleanup,
656 )
657 .with_phase_timings(phase_timings.write_batch_failed()),
658 source.to_string(),
659 )
660 })?;
661
662 rows_written = rows_written.saturating_add(row_count);
663 batches_written = batches_written.saturating_add(1);
664 shaped_rows = shaped_rows.saturating_add(row_count);
665 shaped_batches = shaped_batches.saturating_add(1);
666 }
667
668 let finalize_started_at = Instant::now();
669 let finish_result = MssqlBulkLoadWriter::finish(writer).await;
670 let finalize_elapsed = finalize_started_at.elapsed();
671 finish_result.map_err(|source| {
672 let elapsed_ms = elapsed_ms_since(started_at);
673 let batch_shaping = trace_datafusion_batch_stream_finished(
674 output_plan,
675 MssqlBatchShapingReport::completed(
676 input_batches,
677 input_rows,
678 shaped_batches,
679 shaped_rows,
680 ),
681 );
682 mssql_write_phase_error(
683 output_plan,
684 MssqlWritePhase::Finalize,
685 write_failure_report_metrics(
686 RowCount::exact(input_rows),
687 batch_shaping,
688 MssqlWriteProgress::new(rows_written, batches_written, elapsed_ms),
689 partial_write_possible(output_plan, rows_written, batches_written),
690 cleanup,
691 )
692 .with_phase_timings(phase_timings.finalize_failed(finalize_elapsed)),
693 source.to_string(),
694 )
695 })?;
696
697 Ok(MssqlWriteReport::from_output_plan_with_metrics(
698 output_plan,
699 write_report_metrics(
700 RowCount::exact(input_rows),
701 trace_datafusion_batch_stream_finished(
702 output_plan,
703 MssqlBatchShapingReport::completed(
704 input_batches,
705 input_rows,
706 shaped_batches,
707 shaped_rows,
708 ),
709 ),
710 MssqlWriteProgress::new(rows_written, batches_written, elapsed_ms_since(started_at)),
711 false,
712 cleanup,
713 )
714 .with_phase_timings(phase_timings.completed(finalize_elapsed)),
715 ))
716}
717
718fn trace_datafusion_batch_stream_finished(
719 output_plan: &MssqlTargetOutputPlan,
720 report: MssqlBatchShapingReport,
721) -> MssqlBatchShapingReport {
722 observability::datafusion_batch_stream_finished(
723 output_plan.output_name(),
724 output_plan.target_table(),
725 output_plan.load_mode(),
726 report,
727 );
728 report
729}
730
731fn mssql_write_phase_error(
732 output_plan: &MssqlTargetOutputPlan,
733 phase: MssqlWritePhase,
734 metrics: MssqlWriteReportMetrics,
735 message: String,
736) -> DeltaFunnelError {
737 DeltaFunnelError::MssqlWritePhase {
738 context: Box::new(MssqlWriteFailureContext::from_output_plan_with_metrics(
739 output_plan,
740 phase,
741 metrics,
742 )),
743 message,
744 }
745}
746
747fn mssql_writer_initialization_error(
748 output_plan: &MssqlTargetOutputPlan,
749 prepared_action: MssqlPreparedTargetAction,
750 cleanup: MssqlTargetCleanupStatus,
751 message: String,
752) -> DeltaFunnelError {
753 mssql_write_phase_error(
754 output_plan,
755 MssqlWritePhase::InitializeWriter,
756 write_report_metrics(
757 RowCount::unavailable(),
758 MssqlBatchShapingReport::not_started(ReportReasonCode::NotExecuted),
759 MssqlWriteProgress::zero(),
760 false,
761 cleanup,
762 ),
763 format!("prepared target action {prepared_action:?}: {message}"),
764 )
765}
766
767fn ensure_prepared_target_matches_output_plan(
768 output_plan: &MssqlTargetOutputPlan,
769 prepared_target: &MssqlPreparedTarget,
770) -> Result<(), DeltaFunnelError> {
771 let report = prepared_target.report();
772 let matches_plan = report.output_name() == output_plan.output_name()
773 && report.target_table() == output_plan.target_table()
774 && report.load_mode() == output_plan.load_mode()
775 && report.connection_source() == output_plan.connection_source()
776 && report.connection() == output_plan.connection()
777 && report.expected_target_state() == output_plan.lifecycle_plan().expected_target_state();
778
779 if !matches_plan {
780 return Err(mssql_writer_initialization_error(
781 output_plan,
782 report.action(),
783 report.cleanup(),
784 "prepared target does not match output plan".to_owned(),
785 ));
786 }
787
788 if !output_plan
789 .lifecycle_plan()
790 .execution_guardrails()
791 .contains(&MssqlLifecycleExecutionGuardrail::BulkWriterConstruction)
792 {
793 return Err(mssql_writer_initialization_error(
794 output_plan,
795 report.action(),
796 report.cleanup(),
797 "target lifecycle plan does not allow BulkWriterConstruction".to_owned(),
798 ));
799 }
800
801 Ok(())
802}
803
804fn partial_write_possible(
805 output_plan: &MssqlTargetOutputPlan,
806 rows_written: u64,
807 batches_written: u64,
808) -> bool {
809 output_plan.load_mode() == LoadMode::AppendExisting && (rows_written > 0 || batches_written > 0)
810}
811
812#[derive(Clone, Copy)]
813struct MssqlWriteProgress {
814 rows_written: u64,
815 batches_written: u64,
816 elapsed_ms: u64,
817}
818
819impl MssqlWriteProgress {
820 const fn new(rows_written: u64, batches_written: u64, elapsed_ms: u64) -> Self {
821 Self {
822 rows_written,
823 batches_written,
824 elapsed_ms,
825 }
826 }
827
828 const fn zero() -> Self {
829 Self::new(0, 0, 0)
830 }
831}
832
833fn elapsed_ms_since(started_at: Instant) -> u64 {
834 u64::try_from(started_at.elapsed().as_millis()).unwrap_or(u64::MAX)
835}
836
837fn batch_row_count(row_count: usize) -> u64 {
838 u64::try_from(row_count).unwrap_or(u64::MAX)
839}
840
841#[cfg(test)]
842mod tests {
843 use std::{
844 sync::{Arc, Mutex, MutexGuard},
845 time::Duration,
846 };
847
848 use arrow_schema::{DataType, Field, Schema};
849 use arrow_tiberius::{
850 DiagnosticCode, DiagnosticSeverity, PlanOptions, SchemaCheck, StringPolicy, WriteBackend,
851 WriteOptions,
852 };
853 use datafusion::arrow::{
854 array::{Int64Array, StringArray, StringViewArray},
855 record_batch::RecordBatch,
856 };
857 use futures_util::stream;
858
859 use super::*;
860 use crate::{
861 DeltaFunnelError, MssqlBatchShapingReport, MssqlConnectionConfig, MssqlConnectionSource,
862 MssqlOutputFieldReport, MssqlTargetConfig, MssqlTargetTable, MssqlWriteStats, PhaseStatus,
863 PhaseTimingReport, ReportReasonCode, ValidationStatus, plan_mssql_target_for_output,
864 };
865
866 fn secret_connection() -> Result<MssqlConnectionConfig, DeltaFunnelError> {
867 Ok(MssqlConnectionConfig::new(
868 "server=tcp:sql.example.com;database=warehouse;user=admin;password=secret-token",
869 )?
870 .with_display_label("warehouse-primary"))
871 }
872
873 fn orders_schema() -> Schema {
874 Schema::new(vec![
875 Field::new("order_id", DataType::Int64, false),
876 Field::new("status", DataType::Utf8, true),
877 ])
878 }
879
880 #[derive(Debug, Default)]
881 struct FakeBulkLoadWriter {
882 accepted_rows: u64,
883 accepted_batches: u64,
884 batch_rows: Vec<usize>,
885 log: Option<Arc<Mutex<FakeBulkLoadWriterLog>>>,
886 fail_on_write_batch: Option<u64>,
887 fail_on_finish: bool,
888 write_delay: Option<Duration>,
889 }
890
891 #[derive(Debug, Default, PartialEq, Eq)]
892 struct FakeBulkLoadWriterLog {
893 batch_rows: Vec<usize>,
894 finish_count: u64,
895 }
896
897 #[derive(Clone, Default)]
898 struct RecordingBulkWriterFactory {
899 request: Arc<Mutex<Option<MssqlBulkWriterInitializationRequest>>>,
900 error: Option<String>,
901 }
902
903 impl FakeBulkLoadWriter {
904 fn with_log(log: Arc<Mutex<FakeBulkLoadWriterLog>>) -> Self {
905 Self {
906 log: Some(log),
907 ..Self::default()
908 }
909 }
910
911 fn fail_on_write_batch(mut self, write_batch_call: u64) -> Self {
912 self.fail_on_write_batch = Some(write_batch_call);
913 self
914 }
915
916 fn delay_writes_by(mut self, delay: Duration) -> Self {
917 self.write_delay = Some(delay);
918 self
919 }
920
921 fn fail_on_finish(mut self) -> Self {
922 self.fail_on_finish = true;
923 self
924 }
925 }
926
927 impl RecordingBulkWriterFactory {
928 fn with_request_log(
929 request: Arc<Mutex<Option<MssqlBulkWriterInitializationRequest>>>,
930 ) -> Self {
931 Self {
932 request,
933 error: None,
934 }
935 }
936
937 fn fail_with(mut self, error: impl Into<String>) -> Self {
938 self.error = Some(error.into());
939 self
940 }
941 }
942
943 #[async_trait]
944 impl MssqlBulkLoadWriter for FakeBulkLoadWriter {
945 async fn write_batch(
946 &mut self,
947 batch: &RecordBatch,
948 ) -> Result<arrow_tiberius::WriteStats, arrow_tiberius::Error> {
949 let row_count = batch.num_rows();
950 self.batch_rows.push(row_count);
951 if let Some(log) = &self.log {
952 let Ok(mut log) = log.lock() else {
953 return Err(arrow_tiberius::Error::BackendUnavailable {
954 backend: arrow_tiberius::WriteBackend::DirectRawBulk,
955 reason: "fake writer log mutex was poisoned".to_owned(),
956 });
957 };
958 log.batch_rows.push(row_count);
959 }
960 if let Some(delay) = self.write_delay {
961 tokio::time::sleep(delay).await;
962 }
963 let write_batch_call = self.accepted_batches.saturating_add(1);
964 if self.fail_on_write_batch == Some(write_batch_call) {
965 return Err(arrow_tiberius::Error::BackendUnavailable {
966 backend: arrow_tiberius::WriteBackend::DirectRawBulk,
967 reason: format!("fake writer failed on batch {write_batch_call}"),
968 });
969 }
970 self.accepted_rows = self.accepted_rows.saturating_add(row_count as u64);
971 self.accepted_batches = self.accepted_batches.saturating_add(1);
972
973 Ok(arrow_tiberius::WriteStats {
974 rows_written: self.accepted_rows,
975 batches_written: self.accepted_batches,
976 })
977 }
978
979 async fn finish(self) -> Result<arrow_tiberius::WriteStats, arrow_tiberius::Error> {
980 if let Some(log) = &self.log {
981 let Ok(mut log) = log.lock() else {
982 return Err(arrow_tiberius::Error::BackendUnavailable {
983 backend: arrow_tiberius::WriteBackend::DirectRawBulk,
984 reason: "fake writer log mutex was poisoned".to_owned(),
985 });
986 };
987 log.finish_count = log.finish_count.saturating_add(1);
988 }
989 if self.fail_on_finish {
990 return Err(arrow_tiberius::Error::BackendUnavailable {
991 backend: arrow_tiberius::WriteBackend::DirectRawBulk,
992 reason: "fake writer failed on finish".to_owned(),
993 });
994 }
995
996 Ok(arrow_tiberius::WriteStats {
997 rows_written: self.accepted_rows,
998 batches_written: self.accepted_batches,
999 })
1000 }
1001 }
1002
1003 #[async_trait]
1004 impl MssqlBulkWriterFactory for RecordingBulkWriterFactory {
1005 type Writer = FakeBulkLoadWriter;
1006
1007 async fn initialize(
1008 self,
1009 request: MssqlBulkWriterInitializationRequest,
1010 ) -> Result<Self::Writer, arrow_tiberius::Error> {
1011 let Ok(mut request_log) = self.request.lock() else {
1012 return Err(arrow_tiberius::Error::BackendUnavailable {
1013 backend: arrow_tiberius::WriteBackend::DirectRawBulk,
1014 reason: "fake initializer request log mutex was poisoned".to_owned(),
1015 });
1016 };
1017 *request_log = Some(request);
1018
1019 if let Some(reason) = self.error {
1020 return Err(arrow_tiberius::Error::BackendUnavailable {
1021 backend: arrow_tiberius::WriteBackend::DirectRawBulk,
1022 reason,
1023 });
1024 }
1025
1026 Ok(FakeBulkLoadWriter::default())
1027 }
1028 }
1029
1030 fn orders_batch(
1031 order_ids: Vec<i64>,
1032 statuses: Vec<Option<&str>>,
1033 ) -> Result<RecordBatch, DeltaFunnelError> {
1034 RecordBatch::try_new(
1035 Arc::new(orders_schema()),
1036 vec![
1037 Arc::new(Int64Array::from(order_ids)),
1038 Arc::new(StringArray::from(statuses)),
1039 ],
1040 )
1041 .map_err(|error| DeltaFunnelError::Config {
1042 message: error.to_string(),
1043 })
1044 }
1045
1046 fn orders_batch_with_int32_order_id() -> Result<RecordBatch, DeltaFunnelError> {
1047 let schema = Schema::new(vec![
1048 Field::new("order_id", DataType::Int32, false),
1049 Field::new("status", DataType::Utf8, true),
1050 ]);
1051
1052 RecordBatch::try_new(
1053 Arc::new(schema),
1054 vec![
1055 Arc::new(datafusion::arrow::array::Int32Array::from(vec![1_i32])),
1056 Arc::new(StringArray::from(vec![Some("open")])),
1057 ],
1058 )
1059 .map_err(|error| DeltaFunnelError::Config {
1060 message: error.to_string(),
1061 })
1062 }
1063
1064 fn orders_batch_with_utf8_view_status() -> Result<RecordBatch, DeltaFunnelError> {
1065 let schema = Schema::new(vec![
1066 Field::new("order_id", DataType::Int64, false),
1067 Field::new("status", DataType::Utf8View, true),
1068 ]);
1069
1070 RecordBatch::try_new(
1071 Arc::new(schema),
1072 vec![
1073 Arc::new(Int64Array::from(vec![1_i64, 2])),
1074 Arc::new(StringViewArray::from(vec![Some("open"), None])),
1075 ],
1076 )
1077 .map_err(|error| DeltaFunnelError::Config {
1078 message: error.to_string(),
1079 })
1080 }
1081
1082 fn lock_fake_writer_log(
1083 log: &Arc<Mutex<FakeBulkLoadWriterLog>>,
1084 ) -> Result<MutexGuard<'_, FakeBulkLoadWriterLog>, DeltaFunnelError> {
1085 log.lock().map_err(|_| DeltaFunnelError::Config {
1086 message: "fake writer log mutex was poisoned".to_owned(),
1087 })
1088 }
1089
1090 fn take_initialization_request(
1091 request: &Arc<Mutex<Option<MssqlBulkWriterInitializationRequest>>>,
1092 ) -> Result<MssqlBulkWriterInitializationRequest, DeltaFunnelError> {
1093 let mut request = request.lock().map_err(|_| DeltaFunnelError::Config {
1094 message: "fake initializer request log mutex was poisoned".to_owned(),
1095 })?;
1096
1097 request.take().ok_or_else(|| DeltaFunnelError::Config {
1098 message: "expected fake initializer request".to_owned(),
1099 })
1100 }
1101
1102 fn output_plan_for_orders_schema() -> Result<MssqlTargetOutputPlan, DeltaFunnelError> {
1103 output_plan_for_orders_schema_with_load_mode(LoadMode::AppendExisting)
1104 }
1105
1106 fn output_plan_for_orders_schema_with_load_mode(
1107 load_mode: LoadMode,
1108 ) -> Result<MssqlTargetOutputPlan, DeltaFunnelError> {
1109 let connection = secret_connection()?;
1110 let target_config = MssqlTargetConfig::new(MssqlTargetTable::new("dbo", "orders")?)
1111 .with_load_mode(load_mode);
1112
1113 plan_mssql_target_for_output(
1114 orders_schema(),
1115 "orders_output",
1116 &target_config,
1117 Some(&connection),
1118 PlanOptions::default(),
1119 )
1120 }
1121
1122 fn assert_batch_schema_validation_error(
1123 error: DeltaFunnelError,
1124 expected_field: Option<(usize, &str)>,
1125 ) -> Result<Box<MssqlWriteFailureContext>, DeltaFunnelError> {
1126 let DeltaFunnelError::MssqlBatchSchemaValidation { context, source } = error else {
1127 return Err(DeltaFunnelError::Config {
1128 message: "expected MSSQL batch schema validation error".to_owned(),
1129 });
1130 };
1131
1132 assert_eq!(context.phase(), MssqlWritePhase::ValidateBatchSchema);
1133 assert_eq!(context.output_name(), "orders_output");
1134 assert_eq!(context.target_table().schema(), Some("dbo"));
1135 assert_eq!(context.target_table().table(), "orders");
1136 assert_eq!(context.load_mode(), LoadMode::AppendExisting);
1137 assert!(!context.partial_write_possible());
1138
1139 let arrow_tiberius::Error::ValueConversion { diagnostics } = source else {
1140 return Err(DeltaFunnelError::Config {
1141 message: "expected arrow-tiberius value conversion error".to_owned(),
1142 });
1143 };
1144 assert_eq!(diagnostics.len(), 1);
1145 let diagnostic = &diagnostics.all()[0];
1146 assert_eq!(diagnostic.code(), DiagnosticCode::SchemaMismatch);
1147 assert_eq!(
1148 diagnostic
1149 .field()
1150 .map(|field| (field.index(), field.name())),
1151 expected_field
1152 );
1153 assert_eq!(context.diagnostics().len(), 1);
1154 let context_diagnostic = &context.diagnostics()[0];
1155 assert_eq!(context_diagnostic.severity(), DiagnosticSeverity::Error);
1156 assert_eq!(context_diagnostic.code(), DiagnosticCode::SchemaMismatch);
1157 assert_eq!(
1158 context_diagnostic
1159 .field()
1160 .map(|field| (field.index(), field.name())),
1161 expected_field.map(|(index, name)| (u64::try_from(index).unwrap_or(u64::MAX), name))
1162 );
1163 let value = context.to_json_value();
1164 assert_eq!(value["diagnostics"][0]["severity"], "error");
1165 assert_eq!(value["diagnostics"][0]["code"], "schema_mismatch");
1166 assert_eq!(value["cleanup_error"], serde_json::Value::Null);
1167 Ok(context)
1168 }
1169
1170 fn assert_write_phase_error(
1171 error: DeltaFunnelError,
1172 expected_phase: MssqlWritePhase,
1173 expected_rows_written: u64,
1174 expected_batches_written: u64,
1175 expected_partial_write_possible: bool,
1176 ) -> Result<Box<MssqlWriteFailureContext>, DeltaFunnelError> {
1177 let DeltaFunnelError::MssqlWritePhase { context, .. } = error else {
1178 return Err(DeltaFunnelError::Config {
1179 message: "expected MSSQL write phase error".to_owned(),
1180 });
1181 };
1182
1183 assert_eq!(context.phase(), expected_phase);
1184 assert_eq!(context.output_name(), "orders_output");
1185 assert_eq!(context.target_table().schema(), Some("dbo"));
1186 assert_eq!(context.target_table().table(), "orders");
1187 assert_eq!(context.load_mode(), LoadMode::AppendExisting);
1188 assert_eq!(context.stats().rows_written(), expected_rows_written);
1189 assert_eq!(context.stats().batches_written(), expected_batches_written);
1190 assert_eq!(
1191 context.partial_write_possible(),
1192 expected_partial_write_possible
1193 );
1194 assert_eq!(context.cleanup(), MssqlTargetCleanupStatus::NotApplicable);
1195 Ok(context)
1196 }
1197
1198 fn assert_output_schema(fields: &[MssqlOutputFieldReport]) {
1199 assert_eq!(fields.len(), 2);
1200 assert_eq!(fields[0].index(), 0);
1201 assert_eq!(fields[0].name(), "order_id");
1202 assert_eq!(fields[0].arrow_type(), "Int64");
1203 assert!(!fields[0].nullable());
1204 assert_eq!(fields[1].index(), 1);
1205 assert_eq!(fields[1].name(), "status");
1206 assert_eq!(fields[1].arrow_type(), "Utf8");
1207 assert!(fields[1].nullable());
1208 }
1209
1210 fn assert_batch_shaping(
1211 report: MssqlBatchShapingReport,
1212 expected_status: PhaseStatus,
1213 expected_input_batches: u64,
1214 expected_input_rows: u64,
1215 expected_output_batches: u64,
1216 expected_output_rows: u64,
1217 ) {
1218 assert_eq!(report.status(), expected_status);
1219 assert_eq!(report.input_batches(), expected_input_batches);
1220 assert_eq!(report.input_rows(), expected_input_rows);
1221 assert_eq!(report.output_batches(), expected_output_batches);
1222 assert_eq!(report.output_rows(), expected_output_rows);
1223 }
1224
1225 fn assert_initialize_writer_error(
1226 error: DeltaFunnelError,
1227 expected_message: &str,
1228 expected_cleanup: MssqlTargetCleanupStatus,
1229 ) -> Result<String, DeltaFunnelError> {
1230 let display = error.to_string();
1231 assert!(display.contains("orders_output"));
1232 assert!(display.contains("initialize writer"));
1233 assert!(display.contains(expected_message));
1234 assert!(!display.contains("secret-token"));
1235 assert!(!display.contains("password"));
1236 assert!(!display.contains("server=tcp"));
1237 let DeltaFunnelError::MssqlWritePhase { context, message } = error else {
1238 return Err(DeltaFunnelError::Config {
1239 message: "expected MSSQL write phase error".to_owned(),
1240 });
1241 };
1242
1243 assert_eq!(context.phase(), MssqlWritePhase::InitializeWriter);
1244 assert_eq!(context.output_name(), "orders_output");
1245 assert_eq!(context.stats().rows_written(), 0);
1246 assert_eq!(context.stats().batches_written(), 0);
1247 assert!(!context.partial_write_possible());
1248 assert_eq!(context.cleanup(), expected_cleanup);
1249 Ok(message)
1250 }
1251
1252 #[test]
1253 fn write_stats_preserve_output_counts_and_elapsed_time() {
1254 let stats = MssqlWriteStats::new("orders", 42, 3, 125);
1255
1256 assert_eq!(stats.output_name(), "orders");
1257 assert_eq!(stats.rows_written(), 42);
1258 assert_eq!(stats.batches_written(), 3);
1259 assert_eq!(stats.elapsed_ms(), 125);
1260 }
1261
1262 #[test]
1263 fn connected_bulk_writer_adapts_to_bulk_load_writer_trait() {
1264 fn assert_bulk_load_writer<W: MssqlBulkLoadWriter>() {}
1265
1266 assert_bulk_load_writer::<arrow_tiberius::ConnectedBulkWriter<'static>>();
1267 }
1268
1269 #[tokio::test]
1270 async fn bulk_writer_initialization_passes_prepared_identity_mappings_and_options()
1271 -> Result<(), DeltaFunnelError> {
1272 let output_plan = output_plan_for_orders_schema()?;
1273 let prepared_target = MssqlPreparedTarget::from_output_plan(
1274 &output_plan,
1275 MssqlPreparedTargetAction::VerifiedExisting,
1276 )?;
1277 let options = WriteOptions {
1278 backend: WriteBackend::BaselineTokenRow,
1279 schema_check: SchemaCheck::Strict,
1280 plan_options: PlanOptions {
1281 string_policy: StringPolicy::NVarChar(128),
1282 ..PlanOptions::default()
1283 },
1284 };
1285 let request_log = Arc::new(Mutex::new(None));
1286 let factory = RecordingBulkWriterFactory::with_request_log(Arc::clone(&request_log));
1287
1288 let writer = initialize_mssql_bulk_writer_with_factory(
1289 &output_plan,
1290 &prepared_target,
1291 options,
1292 factory,
1293 )
1294 .await?;
1295 let final_stats = MssqlBulkLoadWriter::finish(writer)
1296 .await
1297 .map_err(|source| DeltaFunnelError::MssqlWrite { source })?;
1298 let request = take_initialization_request(&request_log)?;
1299
1300 assert_eq!(final_stats.rows_written, 0);
1301 assert_eq!(final_stats.batches_written, 0);
1302 assert_eq!(request.output_name(), "orders_output");
1303 assert_eq!(request.table().quoted_sql(), "[dbo].[orders]");
1304 assert_eq!(request.mappings(), output_plan.schema_mappings());
1305 assert_eq!(request.options(), options);
1306 assert_eq!(
1307 request.prepared_action(),
1308 MssqlPreparedTargetAction::VerifiedExisting
1309 );
1310 assert_eq!(request.cleanup(), MssqlTargetCleanupStatus::NotApplicable);
1311 Ok(())
1312 }
1313
1314 #[tokio::test]
1315 async fn bulk_writer_initialization_errors_map_to_initialize_writer_phase()
1316 -> Result<(), DeltaFunnelError> {
1317 let output_plan = output_plan_for_orders_schema_with_load_mode(LoadMode::CreateAndLoad)?;
1318 let prepared_target = MssqlPreparedTarget::from_output_plan(
1319 &output_plan,
1320 MssqlPreparedTargetAction::CreatedTable,
1321 )?;
1322 let request_log = Arc::new(Mutex::new(None));
1323 let factory = RecordingBulkWriterFactory::with_request_log(Arc::clone(&request_log))
1324 .fail_with("target metadata failed\nfor test");
1325
1326 let error = initialize_mssql_bulk_writer_with_factory(
1327 &output_plan,
1328 &prepared_target,
1329 default_mssql_write_options(),
1330 factory,
1331 )
1332 .await
1333 .err()
1334 .ok_or_else(|| DeltaFunnelError::Config {
1335 message: "expected bulk writer initialization error".to_owned(),
1336 })?;
1337 let message = assert_initialize_writer_error(
1338 error,
1339 r"target metadata failed\nfor test",
1340 MssqlTargetCleanupStatus::NotAttempted,
1341 )?;
1342 let request = take_initialization_request(&request_log)?;
1343
1344 assert!(message.contains("CreatedTable"));
1345 assert_eq!(request.cleanup(), MssqlTargetCleanupStatus::NotAttempted);
1346 assert_eq!(
1347 request.prepared_action(),
1348 MssqlPreparedTargetAction::CreatedTable
1349 );
1350 Ok(())
1351 }
1352
1353 #[tokio::test]
1354 async fn bulk_writer_initialization_rejects_mismatched_prepared_target_before_factory()
1355 -> Result<(), DeltaFunnelError> {
1356 let output_plan = output_plan_for_orders_schema()?;
1357 let connection = secret_connection()?;
1358 let other_target_config =
1359 MssqlTargetConfig::new(MssqlTargetTable::new("dbo", "other_orders")?);
1360 let other_output_plan = plan_mssql_target_for_output(
1361 orders_schema(),
1362 "other_output",
1363 &other_target_config,
1364 Some(&connection),
1365 PlanOptions::default(),
1366 )?;
1367 let prepared_target = MssqlPreparedTarget::from_output_plan(
1368 &other_output_plan,
1369 MssqlPreparedTargetAction::VerifiedExisting,
1370 )?;
1371 let request_log = Arc::new(Mutex::new(None));
1372 let factory = RecordingBulkWriterFactory::with_request_log(Arc::clone(&request_log));
1373
1374 let error = initialize_mssql_bulk_writer_with_factory(
1375 &output_plan,
1376 &prepared_target,
1377 default_mssql_write_options(),
1378 factory,
1379 )
1380 .await
1381 .err()
1382 .ok_or_else(|| DeltaFunnelError::Config {
1383 message: "expected mismatched prepared target error".to_owned(),
1384 })?;
1385
1386 assert_initialize_writer_error(
1387 error,
1388 "prepared target does not match output plan",
1389 MssqlTargetCleanupStatus::NotApplicable,
1390 )?;
1391 assert!(take_initialization_request(&request_log).is_err());
1392 Ok(())
1393 }
1394
1395 #[tokio::test]
1396 async fn bulk_load_writer_trait_consumes_writer_on_finish() -> Result<(), DeltaFunnelError> {
1397 let mut writer = FakeBulkLoadWriter::default();
1398 let first = orders_batch(vec![1, 2], vec![Some("open"), Some("closed")])?;
1399 let second = orders_batch(vec![3], vec![None])?;
1400
1401 let first_stats = MssqlBulkLoadWriter::write_batch(&mut writer, &first)
1402 .await
1403 .map_err(|source| DeltaFunnelError::MssqlWrite { source })?;
1404 let second_stats = MssqlBulkLoadWriter::write_batch(&mut writer, &second)
1405 .await
1406 .map_err(|source| DeltaFunnelError::MssqlWrite { source })?;
1407 assert_eq!(writer.batch_rows, vec![2, 1]);
1408
1409 let final_stats = MssqlBulkLoadWriter::finish(writer)
1410 .await
1411 .map_err(|source| DeltaFunnelError::MssqlWrite { source })?;
1412
1413 assert_eq!(first_stats.rows_written, 2);
1414 assert_eq!(first_stats.batches_written, 1);
1415 assert_eq!(second_stats.rows_written, 3);
1416 assert_eq!(second_stats.batches_written, 2);
1417 assert_eq!(final_stats.rows_written, 3);
1418 assert_eq!(final_stats.batches_written, 2);
1419 Ok(())
1420 }
1421
1422 #[tokio::test]
1423 async fn write_loop_writes_batches_in_order_counts_accepted_and_finishes()
1424 -> Result<(), DeltaFunnelError> {
1425 let output_plan = output_plan_for_orders_schema()?;
1426 let log = Arc::new(Mutex::new(FakeBulkLoadWriterLog::default()));
1427 let writer = FakeBulkLoadWriter::with_log(Arc::clone(&log))
1428 .delay_writes_by(Duration::from_millis(2));
1429 let first = orders_batch(vec![1, 2], vec![Some("open"), Some("closed")])?;
1430 let second = orders_batch(vec![3], vec![None])?;
1431 let batches = stream::iter(vec![Ok(first), Ok(second)]);
1432
1433 let report = write_mssql_batches_with_writer(
1434 &output_plan,
1435 batches,
1436 writer,
1437 default_mssql_write_options(),
1438 )
1439 .await?;
1440
1441 assert_eq!(report.output_name(), "orders_output");
1442 assert_eq!(report.stats().rows_written(), 3);
1443 assert_eq!(report.stats().batches_written(), 2);
1444 assert!(report.stats().elapsed_ms() > 0);
1445 assert!(!report.partial_write_possible());
1446 assert_eq!(report.cleanup(), MssqlTargetCleanupStatus::NotApplicable);
1447 assert_output_schema(report.output_schema());
1448 assert_eq!(report.output_row_count(), RowCount::exact(3));
1449 assert_batch_shaping(report.batch_shaping(), PhaseStatus::completed(), 2, 3, 2, 3);
1450 assert_phase_timing(
1451 report.phase_timings(),
1452 POLL_BATCH_STREAM_PHASE,
1453 PhaseStatus::completed(),
1454 )?;
1455 assert_phase_timing(
1456 report.phase_timings(),
1457 VALIDATE_BATCH_SCHEMA_PHASE,
1458 PhaseStatus::completed(),
1459 )?;
1460 assert_phase_timing(
1461 report.phase_timings(),
1462 WRITE_BATCH_PHASE,
1463 PhaseStatus::completed(),
1464 )?;
1465 assert_phase_timing(
1466 report.phase_timings(),
1467 FINALIZE_PHASE,
1468 PhaseStatus::completed(),
1469 )?;
1470 assert_phase_timing(
1471 report.phase_timings(),
1472 VALIDATION_PHASE,
1473 PhaseStatus::not_started(ReportReasonCode::NotExecuted),
1474 )?;
1475 let debug = format!("{report:?}");
1476 assert!(!debug.contains("open"));
1477 assert!(!debug.contains("closed"));
1478
1479 let log = lock_fake_writer_log(&log)?;
1480 assert_eq!(log.batch_rows, vec![2, 1]);
1481 assert_eq!(log.finish_count, 1);
1482 Ok(())
1483 }
1484
1485 #[tokio::test]
1486 async fn write_loop_accepts_utf8_view_for_planned_utf8() -> Result<(), DeltaFunnelError> {
1487 let output_plan = output_plan_for_orders_schema()?;
1488 let log = Arc::new(Mutex::new(FakeBulkLoadWriterLog::default()));
1489 let writer = FakeBulkLoadWriter::with_log(Arc::clone(&log));
1490 let batches = stream::iter(vec![Ok(orders_batch_with_utf8_view_status()?)]);
1491
1492 let report = write_mssql_batches_with_writer(
1493 &output_plan,
1494 batches,
1495 writer,
1496 default_mssql_write_options(),
1497 )
1498 .await?;
1499
1500 assert_eq!(report.stats().rows_written(), 2);
1501 assert_eq!(report.stats().batches_written(), 1);
1502 let log = lock_fake_writer_log(&log)?;
1503 assert_eq!(log.batch_rows, vec![2]);
1504 assert_eq!(log.finish_count, 1);
1505 Ok(())
1506 }
1507
1508 #[tokio::test]
1509 async fn write_loop_empty_stream_finishes_with_zero_stats() -> Result<(), DeltaFunnelError> {
1510 let output_plan = output_plan_for_orders_schema()?;
1511 let log = Arc::new(Mutex::new(FakeBulkLoadWriterLog::default()));
1512 let writer = FakeBulkLoadWriter::with_log(Arc::clone(&log));
1513 let batches = stream::empty::<Result<RecordBatch, DeltaFunnelError>>();
1514
1515 let report = write_mssql_batches_with_writer(
1516 &output_plan,
1517 batches,
1518 writer,
1519 default_mssql_write_options(),
1520 )
1521 .await?;
1522
1523 assert_eq!(report.output_name(), "orders_output");
1524 assert_eq!(report.stats().rows_written(), 0);
1525 assert_eq!(report.stats().batches_written(), 0);
1526 assert!(!report.partial_write_possible());
1527 assert_eq!(report.output_row_count(), RowCount::exact(0));
1528 assert_batch_shaping(report.batch_shaping(), PhaseStatus::completed(), 0, 0, 0, 0);
1529 let log = lock_fake_writer_log(&log)?;
1530 assert!(log.batch_rows.is_empty());
1531 assert_eq!(log.finish_count, 1);
1532 Ok(())
1533 }
1534
1535 #[tokio::test]
1536 async fn write_loop_stream_error_preserves_stats_and_skips_finish()
1537 -> Result<(), DeltaFunnelError> {
1538 let output_plan = output_plan_for_orders_schema()?;
1539 let log = Arc::new(Mutex::new(FakeBulkLoadWriterLog::default()));
1540 let writer = FakeBulkLoadWriter::with_log(Arc::clone(&log));
1541 let first = orders_batch(vec![1, 2], vec![Some("open"), Some("closed")])?;
1542 let batches = stream::iter(vec![
1543 Ok(first),
1544 Err(DeltaFunnelError::Config {
1545 message: "stream failed after first batch".to_owned(),
1546 }),
1547 ]);
1548
1549 let error = match write_mssql_batches_with_writer(
1550 &output_plan,
1551 batches,
1552 writer,
1553 default_mssql_write_options(),
1554 )
1555 .await
1556 {
1557 Ok(_) => {
1558 return Err(DeltaFunnelError::Config {
1559 message: "expected stream error".to_owned(),
1560 });
1561 }
1562 Err(error) => error,
1563 };
1564
1565 let context =
1566 assert_write_phase_error(error, MssqlWritePhase::PollBatchStream, 2, 1, true)?;
1567 assert_eq!(context.output_row_count(), RowCount::partial(2));
1568 assert_batch_shaping(context.batch_shaping(), PhaseStatus::failed(), 1, 2, 1, 2);
1569 assert_phase_timing(
1570 context.phase_timings(),
1571 POLL_BATCH_STREAM_PHASE,
1572 PhaseStatus::failed(),
1573 )?;
1574 assert_phase_timing(
1575 context.phase_timings(),
1576 FINALIZE_PHASE,
1577 PhaseStatus::not_started(ReportReasonCode::NotExecuted),
1578 )?;
1579 assert_phase_timing(
1580 context.phase_timings(),
1581 VALIDATION_PHASE,
1582 PhaseStatus::not_started(ReportReasonCode::FailureBeforeValidation),
1583 )?;
1584 let log = lock_fake_writer_log(&log)?;
1585 assert_eq!(log.batch_rows, vec![2]);
1586 assert_eq!(log.finish_count, 0);
1587 Ok(())
1588 }
1589
1590 #[tokio::test]
1591 async fn write_loop_schema_mismatch_skips_writer_and_finish() -> Result<(), DeltaFunnelError> {
1592 let output_plan = output_plan_for_orders_schema()?;
1593 let log = Arc::new(Mutex::new(FakeBulkLoadWriterLog::default()));
1594 let writer = FakeBulkLoadWriter::with_log(Arc::clone(&log));
1595 let batches = stream::iter(vec![Ok(orders_batch_with_int32_order_id()?)]);
1596
1597 let error = match write_mssql_batches_with_writer(
1598 &output_plan,
1599 batches,
1600 writer,
1601 default_mssql_write_options(),
1602 )
1603 .await
1604 {
1605 Ok(_) => {
1606 return Err(DeltaFunnelError::Config {
1607 message: "expected batch schema validation error".to_owned(),
1608 });
1609 }
1610 Err(error) => error,
1611 };
1612
1613 let context = assert_batch_schema_validation_error(error, Some((0, "order_id")))?;
1614 assert_eq!(context.output_row_count(), RowCount::partial(1));
1615 assert_batch_shaping(context.batch_shaping(), PhaseStatus::failed(), 1, 1, 0, 0);
1616 assert_phase_timing(
1617 context.phase_timings(),
1618 VALIDATE_BATCH_SCHEMA_PHASE,
1619 PhaseStatus::failed(),
1620 )?;
1621 assert_phase_timing(
1622 context.phase_timings(),
1623 WRITE_BATCH_PHASE,
1624 PhaseStatus::not_started(ReportReasonCode::NotExecuted),
1625 )?;
1626 let log = lock_fake_writer_log(&log)?;
1627 assert!(log.batch_rows.is_empty());
1628 assert_eq!(log.finish_count, 0);
1629 Ok(())
1630 }
1631
1632 #[tokio::test]
1633 async fn write_loop_write_failure_preserves_stats_and_skips_finish()
1634 -> Result<(), DeltaFunnelError> {
1635 let output_plan = output_plan_for_orders_schema()?;
1636 let log = Arc::new(Mutex::new(FakeBulkLoadWriterLog::default()));
1637 let writer = FakeBulkLoadWriter::with_log(Arc::clone(&log))
1638 .fail_on_write_batch(2)
1639 .delay_writes_by(Duration::from_millis(2));
1640 let first = orders_batch(vec![1, 2], vec![Some("open"), Some("closed")])?;
1641 let second = orders_batch(vec![3], vec![None])?;
1642 let batches = stream::iter(vec![Ok(first), Ok(second)]);
1643
1644 let error = match write_mssql_batches_with_writer(
1645 &output_plan,
1646 batches,
1647 writer,
1648 default_mssql_write_options(),
1649 )
1650 .await
1651 {
1652 Ok(_) => {
1653 return Err(DeltaFunnelError::Config {
1654 message: "expected write batch error".to_owned(),
1655 });
1656 }
1657 Err(error) => error,
1658 };
1659
1660 let DeltaFunnelError::MssqlWritePhase { context, .. } = error else {
1661 return Err(DeltaFunnelError::Config {
1662 message: "expected MSSQL write phase error".to_owned(),
1663 });
1664 };
1665 assert_eq!(context.phase(), MssqlWritePhase::WriteBatch);
1666 assert_eq!(context.output_name(), "orders_output");
1667 assert_eq!(context.stats().rows_written(), 2);
1668 assert_eq!(context.stats().batches_written(), 1);
1669 assert_eq!(context.output_row_count(), RowCount::partial(3));
1670 assert_batch_shaping(context.batch_shaping(), PhaseStatus::failed(), 2, 3, 1, 2);
1671 assert!(context.stats().elapsed_ms() > 0);
1672 assert!(context.partial_write_possible());
1673 assert_eq!(context.cleanup(), MssqlTargetCleanupStatus::NotApplicable);
1674 assert_phase_timing(
1675 context.phase_timings(),
1676 WRITE_BATCH_PHASE,
1677 PhaseStatus::failed(),
1678 )?;
1679 assert_phase_timing(
1680 context.phase_timings(),
1681 FINALIZE_PHASE,
1682 PhaseStatus::not_started(ReportReasonCode::NotExecuted),
1683 )?;
1684 assert_phase_timing(
1685 context.phase_timings(),
1686 VALIDATION_PHASE,
1687 PhaseStatus::not_started(ReportReasonCode::FailureBeforeValidation),
1688 )?;
1689 let log = lock_fake_writer_log(&log)?;
1690 assert_eq!(log.batch_rows, vec![2, 1]);
1691 assert_eq!(log.finish_count, 0);
1692 Ok(())
1693 }
1694
1695 #[tokio::test]
1696 async fn write_loop_create_and_load_write_failure_does_not_report_partial_write()
1697 -> Result<(), DeltaFunnelError> {
1698 let output_plan = output_plan_for_orders_schema_with_load_mode(LoadMode::CreateAndLoad)?;
1699 let log = Arc::new(Mutex::new(FakeBulkLoadWriterLog::default()));
1700 let writer = FakeBulkLoadWriter::with_log(Arc::clone(&log)).fail_on_write_batch(2);
1701 let first = orders_batch(vec![1, 2], vec![Some("open"), Some("closed")])?;
1702 let second = orders_batch(vec![3], vec![None])?;
1703 let batches = stream::iter(vec![Ok(first), Ok(second)]);
1704
1705 let error = match write_mssql_batches_with_writer(
1706 &output_plan,
1707 batches,
1708 writer,
1709 default_mssql_write_options(),
1710 )
1711 .await
1712 {
1713 Ok(_) => {
1714 return Err(DeltaFunnelError::Config {
1715 message: "expected create-and-load write batch error".to_owned(),
1716 });
1717 }
1718 Err(error) => error,
1719 };
1720
1721 let DeltaFunnelError::MssqlWritePhase { context, .. } = error else {
1722 return Err(DeltaFunnelError::Config {
1723 message: "expected MSSQL write phase error".to_owned(),
1724 });
1725 };
1726 assert_eq!(context.phase(), MssqlWritePhase::WriteBatch);
1727 assert_eq!(context.load_mode(), LoadMode::CreateAndLoad);
1728 assert_eq!(context.stats().rows_written(), 2);
1729 assert_eq!(context.stats().batches_written(), 1);
1730 assert!(!context.partial_write_possible());
1731 let log = lock_fake_writer_log(&log)?;
1732 assert_eq!(log.batch_rows, vec![2, 1]);
1733 assert_eq!(log.finish_count, 0);
1734 Ok(())
1735 }
1736
1737 #[tokio::test]
1738 async fn write_loop_finalize_failure_preserves_stats() -> Result<(), DeltaFunnelError> {
1739 let output_plan = output_plan_for_orders_schema()?;
1740 let log = Arc::new(Mutex::new(FakeBulkLoadWriterLog::default()));
1741 let writer = FakeBulkLoadWriter::with_log(Arc::clone(&log)).fail_on_finish();
1742 let first = orders_batch(vec![1, 2], vec![Some("open"), Some("closed")])?;
1743 let second = orders_batch(vec![3], vec![None])?;
1744 let batches = stream::iter(vec![Ok(first), Ok(second)]);
1745
1746 let error = match write_mssql_batches_with_writer(
1747 &output_plan,
1748 batches,
1749 writer,
1750 default_mssql_write_options(),
1751 )
1752 .await
1753 {
1754 Ok(_) => {
1755 return Err(DeltaFunnelError::Config {
1756 message: "expected finalize error".to_owned(),
1757 });
1758 }
1759 Err(error) => error,
1760 };
1761
1762 let context = assert_write_phase_error(error, MssqlWritePhase::Finalize, 3, 2, true)?;
1763 assert_eq!(context.output_row_count(), RowCount::exact(3));
1764 assert_batch_shaping(
1765 context.batch_shaping(),
1766 PhaseStatus::completed(),
1767 2,
1768 3,
1769 2,
1770 3,
1771 );
1772 assert_phase_timing(
1773 context.phase_timings(),
1774 FINALIZE_PHASE,
1775 PhaseStatus::failed(),
1776 )?;
1777 assert_phase_timing(
1778 context.phase_timings(),
1779 VALIDATION_PHASE,
1780 PhaseStatus::not_started(ReportReasonCode::FailureBeforeValidation),
1781 )?;
1782 let log = lock_fake_writer_log(&log)?;
1783 assert_eq!(log.batch_rows, vec![2, 1]);
1784 assert_eq!(log.finish_count, 1);
1785 Ok(())
1786 }
1787
1788 #[test]
1789 fn write_phase_display_is_stable() {
1790 let phases = [
1791 (MssqlWritePhase::Connect, "connect"),
1792 (
1793 MssqlWritePhase::PrepareTargetLifecycle,
1794 "prepare target lifecycle",
1795 ),
1796 (MssqlWritePhase::InitializeWriter, "initialize writer"),
1797 (MssqlWritePhase::PollBatchStream, "poll batch stream"),
1798 (
1799 MssqlWritePhase::ValidateBatchSchema,
1800 "validate batch schema",
1801 ),
1802 (MssqlWritePhase::WriteBatch, "write batch"),
1803 (MssqlWritePhase::Finalize, "finalize"),
1804 (MssqlWritePhase::Validation, "validation"),
1805 (MssqlWritePhase::SwapTarget, "swap target"),
1806 (MssqlWritePhase::Cleanup, "cleanup"),
1807 ];
1808
1809 for (phase, expected) in phases {
1810 assert_eq!(phase.to_string(), expected);
1811 assert!(!format!("{phase:?}").contains("password"));
1812 }
1813 }
1814
1815 #[test]
1816 fn cleanup_status_display_is_stable() {
1817 let statuses = [
1818 (MssqlTargetCleanupStatus::NotApplicable, "not applicable"),
1819 (MssqlTargetCleanupStatus::NotAttempted, "not attempted"),
1820 (MssqlTargetCleanupStatus::Succeeded, "succeeded"),
1821 (MssqlTargetCleanupStatus::Failed, "failed"),
1822 ];
1823
1824 for (status, expected) in statuses {
1825 assert_eq!(status.to_string(), expected);
1826 assert!(!format!("{status:?}").contains("password"));
1827 }
1828 }
1829
1830 #[test]
1831 fn write_report_preserves_plan_context_stats_and_cleanup() -> Result<(), DeltaFunnelError> {
1832 let connection = secret_connection()?;
1833 let target_config = MssqlTargetConfig::new(MssqlTargetTable::new("dbo", "orders")?);
1834 let output_plan = plan_mssql_target_for_output(
1835 orders_schema(),
1836 "orders_output",
1837 &target_config,
1838 Some(&connection),
1839 PlanOptions::default(),
1840 )?;
1841
1842 let report = MssqlWriteReport::from_output_plan(
1843 &output_plan,
1844 42,
1845 3,
1846 125,
1847 true,
1848 MssqlTargetCleanupStatus::NotApplicable,
1849 );
1850
1851 assert_eq!(report.output_name(), "orders_output");
1852 assert_eq!(report.target_table().schema(), Some("dbo"));
1853 assert_eq!(report.target_table().table(), "orders");
1854 assert_eq!(report.load_mode(), LoadMode::AppendExisting);
1855 assert_eq!(
1856 report.connection_source(),
1857 MssqlConnectionSource::ContextDefault
1858 );
1859 assert_eq!(
1860 report.connection().display_label(),
1861 Some("warehouse-primary")
1862 );
1863 assert_eq!(report.stats().output_name(), "orders_output");
1864 assert_eq!(report.stats().rows_written(), 42);
1865 assert_eq!(report.stats().batches_written(), 3);
1866 assert_eq!(report.stats().elapsed_ms(), 125);
1867 assert_output_schema(report.output_schema());
1868 assert_eq!(report.output_row_count(), RowCount::exact(42));
1869 assert_eq!(report.target_row_count(), RowCount::unavailable());
1870 assert_eq!(
1871 report.target_row_count_before_write(),
1872 RowCount::unavailable()
1873 );
1874 assert_eq!(
1875 report.target_row_count_after_write(),
1876 RowCount::unavailable()
1877 );
1878 assert_eq!(
1879 report.validation_status(),
1880 ValidationStatus::skipped(ReportReasonCode::NotExecuted)
1881 );
1882 assert_batch_shaping(
1883 report.batch_shaping(),
1884 PhaseStatus::completed(),
1885 3,
1886 42,
1887 3,
1888 42,
1889 );
1890 assert!(report.partial_write_possible());
1891 assert_eq!(report.cleanup(), MssqlTargetCleanupStatus::NotApplicable);
1892 Ok(())
1893 }
1894
1895 #[test]
1896 fn write_report_records_target_validation_outcome() -> Result<(), DeltaFunnelError> {
1897 let connection = secret_connection()?;
1898 let target_config = MssqlTargetConfig::new(MssqlTargetTable::new("dbo", "orders")?);
1899 let output_plan = plan_mssql_target_for_output(
1900 orders_schema(),
1901 "orders_output",
1902 &target_config,
1903 Some(&connection),
1904 PlanOptions::default(),
1905 )?;
1906 let report = MssqlWriteReport::from_output_plan_with_metrics(
1907 &output_plan,
1908 MssqlWriteReportMetrics::new(
1909 RowCount::exact(42),
1910 MssqlBatchShapingReport::completed(3, 42, 3, 42),
1911 42,
1912 3,
1913 125,
1914 false,
1915 MssqlTargetCleanupStatus::NotApplicable,
1916 )
1917 .with_phase_timings(vec![PhaseTimingReport::not_started(
1918 VALIDATION_PHASE,
1919 ReportReasonCode::NotExecuted,
1920 )]),
1921 );
1922
1923 let report = report.with_target_validation(
1924 RowCount::exact(42),
1925 ValidationStatus::passed(),
1926 PhaseTimingReport::completed(VALIDATION_PHASE, Duration::from_micros(7)),
1927 );
1928
1929 assert_eq!(report.target_row_count(), RowCount::exact(42));
1930 assert_eq!(
1931 report.target_row_count_before_write(),
1932 RowCount::unavailable()
1933 );
1934 assert_eq!(report.target_row_count_after_write(), RowCount::exact(42));
1935 assert_eq!(report.validation_status(), ValidationStatus::passed());
1936 assert_eq!(
1937 report
1938 .phase_timings()
1939 .iter()
1940 .filter(|timing| timing.phase_name() == VALIDATION_PHASE)
1941 .count(),
1942 1
1943 );
1944 assert_phase_timing(
1945 report.phase_timings(),
1946 VALIDATION_PHASE,
1947 PhaseStatus::completed(),
1948 )?;
1949 Ok(())
1950 }
1951
1952 #[test]
1953 fn write_report_records_target_delta_validation_counts() -> Result<(), DeltaFunnelError> {
1954 let connection = secret_connection()?;
1955 let target_config = MssqlTargetConfig::new(MssqlTargetTable::new("dbo", "orders")?);
1956 let output_plan = plan_mssql_target_for_output(
1957 orders_schema(),
1958 "orders_output",
1959 &target_config,
1960 Some(&connection),
1961 PlanOptions::default(),
1962 )?;
1963 let report = MssqlWriteReport::from_output_plan_with_metrics(
1964 &output_plan,
1965 MssqlWriteReportMetrics::new(
1966 RowCount::exact(3),
1967 MssqlBatchShapingReport::completed(1, 3, 1, 3),
1968 3,
1969 1,
1970 125,
1971 false,
1972 MssqlTargetCleanupStatus::NotApplicable,
1973 )
1974 .with_phase_timings(vec![PhaseTimingReport::not_started(
1975 VALIDATION_PHASE,
1976 ReportReasonCode::NotExecuted,
1977 )]),
1978 );
1979
1980 let report = report.with_target_delta_validation(
1981 RowCount::exact(10),
1982 RowCount::exact(13),
1983 ValidationStatus::passed(),
1984 PhaseTimingReport::completed(VALIDATION_PHASE, Duration::from_micros(7)),
1985 );
1986
1987 assert_eq!(report.target_row_count(), RowCount::exact(13));
1988 assert_eq!(report.target_row_count_before_write(), RowCount::exact(10));
1989 assert_eq!(report.target_row_count_after_write(), RowCount::exact(13));
1990 assert_eq!(report.validation_status(), ValidationStatus::passed());
1991 assert_phase_timing(
1992 report.phase_timings(),
1993 VALIDATION_PHASE,
1994 PhaseStatus::completed(),
1995 )?;
1996 Ok(())
1997 }
1998
1999 #[test]
2000 fn write_report_debug_redacts_connection_secret() -> Result<(), DeltaFunnelError> {
2001 let connection = secret_connection()?;
2002 let target_config = MssqlTargetConfig::new(MssqlTargetTable::new("dbo", "orders")?);
2003 let output_plan = plan_mssql_target_for_output(
2004 orders_schema(),
2005 "orders_output",
2006 &target_config,
2007 Some(&connection),
2008 PlanOptions::default(),
2009 )?;
2010
2011 let report = MssqlWriteReport::from_output_plan(
2012 &output_plan,
2013 0,
2014 0,
2015 0,
2016 false,
2017 MssqlTargetCleanupStatus::NotApplicable,
2018 );
2019 let debug = format!("{report:?}");
2020
2021 assert!(debug.contains("warehouse-primary"));
2022 assert!(!debug.contains("secret-token"));
2023 assert!(!debug.contains("password"));
2024 assert!(!debug.contains("server=tcp"));
2025 Ok(())
2026 }
2027
2028 #[test]
2029 fn write_failure_context_preserves_phase_report_and_accepted_stats()
2030 -> Result<(), DeltaFunnelError> {
2031 let connection = secret_connection()?;
2032 let target_config = MssqlTargetConfig::new(MssqlTargetTable::new("dbo", "orders")?);
2033 let output_plan = plan_mssql_target_for_output(
2034 orders_schema(),
2035 "orders_output",
2036 &target_config,
2037 Some(&connection),
2038 PlanOptions::default(),
2039 )?;
2040
2041 let context = MssqlWriteFailureContext::from_output_plan(
2042 &output_plan,
2043 MssqlWritePhase::WriteBatch,
2044 42,
2045 3,
2046 125,
2047 true,
2048 MssqlTargetCleanupStatus::NotApplicable,
2049 );
2050
2051 assert_eq!(context.phase(), MssqlWritePhase::WriteBatch);
2052 assert_eq!(context.output_name(), "orders_output");
2053 assert_eq!(context.target_table().table(), "orders");
2054 assert_eq!(context.load_mode(), LoadMode::AppendExisting);
2055 assert_eq!(
2056 context.connection_source(),
2057 MssqlConnectionSource::ContextDefault
2058 );
2059 assert_eq!(
2060 context.connection().display_label(),
2061 Some("warehouse-primary")
2062 );
2063 assert_eq!(context.stats().rows_written(), 42);
2064 assert_eq!(context.stats().batches_written(), 3);
2065 assert_eq!(context.stats().elapsed_ms(), 125);
2066 assert_eq!(context.output_row_count(), RowCount::partial(42));
2067 assert_eq!(context.target_row_count(), RowCount::unavailable());
2068 assert_eq!(
2069 context.validation_status(),
2070 ValidationStatus::skipped(ReportReasonCode::NotExecuted)
2071 );
2072 assert_batch_shaping(context.batch_shaping(), PhaseStatus::failed(), 3, 42, 3, 42);
2073 assert!(context.partial_write_possible());
2074 assert_eq!(context.cleanup(), MssqlTargetCleanupStatus::NotApplicable);
2075 assert_eq!(context.report().output_name(), "orders_output");
2076 Ok(())
2077 }
2078
2079 fn assert_phase_timing(
2080 timings: &[PhaseTimingReport],
2081 phase_name: &str,
2082 expected_status: PhaseStatus,
2083 ) -> Result<(), DeltaFunnelError> {
2084 let timing = timings
2085 .iter()
2086 .find(|timing| timing.phase_name() == phase_name)
2087 .ok_or_else(|| DeltaFunnelError::Config {
2088 message: format!("missing phase timing {phase_name}"),
2089 })?;
2090
2091 assert_eq!(timing.status(), expected_status);
2092 if expected_status.is_completed() || expected_status.is_failed() {
2093 assert!(timing.elapsed_micros().is_some());
2094 } else {
2095 assert_eq!(timing.elapsed_micros(), None);
2096 }
2097 Ok(())
2098 }
2099
2100 #[test]
2101 fn write_failure_context_debug_redacts_connection_secret() -> Result<(), DeltaFunnelError> {
2102 let connection = secret_connection()?;
2103 let target_config = MssqlTargetConfig::new(MssqlTargetTable::new("dbo", "orders")?);
2104 let output_plan = plan_mssql_target_for_output(
2105 orders_schema(),
2106 "orders_output",
2107 &target_config,
2108 Some(&connection),
2109 PlanOptions::default(),
2110 )?;
2111
2112 let context = MssqlWriteFailureContext::from_output_plan(
2113 &output_plan,
2114 MssqlWritePhase::InitializeWriter,
2115 0,
2116 0,
2117 0,
2118 false,
2119 MssqlTargetCleanupStatus::NotAttempted,
2120 );
2121 let debug = format!("{context:?}");
2122
2123 assert!(debug.contains("warehouse-primary"));
2124 assert!(!debug.contains("secret-token"));
2125 assert!(!debug.contains("password"));
2126 assert!(!debug.contains("server=tcp"));
2127 Ok(())
2128 }
2129
2130 #[test]
2131 fn default_options_pin_direct_raw_bulk_backend() {
2132 let options = default_mssql_write_options();
2133
2134 assert_eq!(options.backend, WriteBackend::DirectRawBulk);
2135 }
2136
2137 #[test]
2138 fn default_options_preserve_arrow_tiberius_schema_check_default() {
2139 let options = default_mssql_write_options();
2140
2141 assert_eq!(options.schema_check, WriteOptions::default().schema_check);
2142 assert_eq!(options.schema_check, SchemaCheck::Strict);
2143 }
2144
2145 #[test]
2146 fn default_options_preserve_arrow_tiberius_plan_options_default() {
2147 let options = default_mssql_write_options();
2148
2149 assert_eq!(options.plan_options, WriteOptions::default().plan_options);
2150 assert_eq!(options.plan_options, PlanOptions::default());
2151 }
2152
2153 #[test]
2154 fn write_options_for_output_plan_preserve_schema_plan_options() -> Result<(), DeltaFunnelError>
2155 {
2156 let connection = secret_connection()?;
2157 let target_config = MssqlTargetConfig::new(MssqlTargetTable::new("dbo", "orders")?);
2158 let plan_options = PlanOptions {
2159 string_policy: StringPolicy::NVarChar(128),
2160 ..PlanOptions::default()
2161 };
2162 let output_plan = plan_mssql_target_for_output(
2163 orders_schema(),
2164 "orders_output",
2165 &target_config,
2166 Some(&connection),
2167 plan_options,
2168 )?;
2169
2170 let write_options = mssql_write_options_for_output_plan(&output_plan);
2171
2172 assert_eq!(write_options.backend, WriteBackend::DirectRawBulk);
2173 assert_eq!(write_options.schema_check, SchemaCheck::Strict);
2174 assert_eq!(write_options.plan_options, plan_options);
2175 Ok(())
2176 }
2177
2178 #[test]
2179 fn output_record_batch_validation_accepts_matching_planned_schema()
2180 -> Result<(), DeltaFunnelError> {
2181 let schema = Arc::new(orders_schema());
2182 let output_plan = output_plan_for_orders_schema()?;
2183 let batch = RecordBatch::try_new(
2184 schema,
2185 vec![
2186 Arc::new(Int64Array::from(vec![1_i64, 2])),
2187 Arc::new(StringArray::from(vec![Some("open"), None])),
2188 ],
2189 )
2190 .map_err(|error| DeltaFunnelError::Config {
2191 message: error.to_string(),
2192 })?;
2193
2194 let report = validate_mssql_output_record_batch(&output_plan, &batch)?;
2195
2196 assert_eq!(report.output_name(), "orders_output");
2197 assert_eq!(report.target_table().schema(), Some("dbo"));
2198 assert_eq!(report.target_table().table(), "orders");
2199 assert_eq!(report.load_mode(), LoadMode::AppendExisting);
2200 assert_eq!(
2201 report.connection_source(),
2202 MssqlConnectionSource::ContextDefault
2203 );
2204 assert_eq!(
2205 report.connection().display_label(),
2206 Some("warehouse-primary")
2207 );
2208 assert!(!format!("{report:?}").contains("secret-token"));
2209 Ok(())
2210 }
2211
2212 #[test]
2213 fn output_record_batch_validation_accepts_utf8_view_for_planned_utf8()
2214 -> Result<(), DeltaFunnelError> {
2215 let output_plan = output_plan_for_orders_schema()?;
2216 let batch = orders_batch_with_utf8_view_status()?;
2217
2218 let report = validate_mssql_output_record_batch(&output_plan, &batch)?;
2219
2220 assert_eq!(report.output_name(), "orders_output");
2221 assert_eq!(report.target_table().schema(), Some("dbo"));
2222 assert_eq!(report.target_table().table(), "orders");
2223 Ok(())
2224 }
2225
2226 #[test]
2227 fn output_schema_validation_accepts_aliased_output_field_names() -> Result<(), DeltaFunnelError>
2228 {
2229 let connection = secret_connection()?;
2230 let target_config = MssqlTargetConfig::new(MssqlTargetTable::new("dbo", "orders")?);
2231 let schema = Schema::new(vec![
2232 Field::new("gross_total", DataType::Float64, true),
2233 Field::new("order_id", DataType::Int32, false),
2234 ]);
2235 let output_plan = plan_mssql_target_for_output(
2236 schema.clone(),
2237 "orders_output",
2238 &target_config,
2239 Some(&connection),
2240 PlanOptions::default(),
2241 )?;
2242
2243 let report = validate_mssql_output_schema(&output_plan, &schema)?;
2244
2245 assert_eq!(
2246 output_plan.schema_mappings()[0].arrow().name(),
2247 "gross_total"
2248 );
2249 assert_eq!(output_plan.schema_mappings()[1].arrow().name(), "order_id");
2250 assert_eq!(report.output_name(), "orders_output");
2251 assert_eq!(report.target_table().schema(), Some("dbo"));
2252 assert_eq!(report.target_table().table(), "orders");
2253 Ok(())
2254 }
2255
2256 #[test]
2257 fn output_schema_validation_rejects_reordered_fields() -> Result<(), DeltaFunnelError> {
2258 let output_plan = output_plan_for_orders_schema()?;
2259 let schema = Schema::new(vec![
2260 Field::new("status", DataType::Utf8, true),
2261 Field::new("order_id", DataType::Int64, false),
2262 ]);
2263
2264 let error = validate_mssql_output_schema(&output_plan, &schema).unwrap_err();
2265
2266 assert_batch_schema_validation_error(error, Some((0, "order_id")))?;
2267 Ok(())
2268 }
2269
2270 #[test]
2271 fn output_schema_validation_rejects_type_mismatch() -> Result<(), DeltaFunnelError> {
2272 let output_plan = output_plan_for_orders_schema()?;
2273 let schema = Schema::new(vec![
2274 Field::new("order_id", DataType::Int32, false),
2275 Field::new("status", DataType::Utf8, true),
2276 ]);
2277
2278 let error = validate_mssql_output_schema(&output_plan, &schema).unwrap_err();
2279
2280 assert_batch_schema_validation_error(error, Some((0, "order_id")))?;
2281 Ok(())
2282 }
2283
2284 #[test]
2285 fn output_schema_validation_rejects_missing_field() -> Result<(), DeltaFunnelError> {
2286 let output_plan = output_plan_for_orders_schema()?;
2287 let schema = Schema::new(vec![Field::new("order_id", DataType::Int64, false)]);
2288
2289 let error = validate_mssql_output_schema(&output_plan, &schema).unwrap_err();
2290
2291 assert_batch_schema_validation_error(error, Some((1, "status")))?;
2292 Ok(())
2293 }
2294
2295 #[test]
2296 fn output_schema_validation_rejects_extra_field() -> Result<(), DeltaFunnelError> {
2297 let output_plan = output_plan_for_orders_schema()?;
2298 let schema = Schema::new(vec![
2299 Field::new("order_id", DataType::Int64, false),
2300 Field::new("status", DataType::Utf8, true),
2301 Field::new("extra", DataType::Utf8, true),
2302 ]);
2303
2304 let error = validate_mssql_output_schema(&output_plan, &schema).unwrap_err();
2305
2306 assert_batch_schema_validation_error(error, None)?;
2307 Ok(())
2308 }
2309
2310 #[test]
2311 fn output_schema_validation_rejects_nullability_mismatch() -> Result<(), DeltaFunnelError> {
2312 let output_plan = output_plan_for_orders_schema()?;
2313 let schema = Schema::new(vec![
2314 Field::new("order_id", DataType::Int64, true),
2315 Field::new("status", DataType::Utf8, true),
2316 ]);
2317
2318 let error = validate_mssql_output_schema(&output_plan, &schema).unwrap_err();
2319
2320 assert_batch_schema_validation_error(error, Some((0, "order_id")))?;
2321 Ok(())
2322 }
2323}