Skip to main content

delta_funnel/sql_server/execution/
write.rs

1//! SQL Server write options.
2//!
3//! This module owns Delta Funnel write defaults around `arrow-tiberius`.
4
5use 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/// Fakeable bulk-load writer boundary for one planned SQL Server output.
42#[async_trait]
43pub(crate) trait MssqlBulkLoadWriter: Sized + Send {
44    /// Writes one record batch.
45    async fn write_batch(
46        &mut self,
47        batch: &RecordBatch,
48    ) -> Result<arrow_tiberius::WriteStats, arrow_tiberius::Error>;
49
50    /// Finalizes the writer and consumes it, matching `arrow-tiberius`.
51    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/// Fakeable boundary for constructing one SQL Server bulk writer.
86#[allow(dead_code)]
87#[async_trait]
88pub(crate) trait MssqlBulkWriterFactory: Send {
89    /// Writer type returned by this factory.
90    type Writer: MssqlBulkLoadWriter;
91
92    /// Constructs a bulk writer from an already prepared target request.
93    async fn initialize(
94        self,
95        request: MssqlBulkWriterInitializationRequest,
96    ) -> Result<Self::Writer, arrow_tiberius::Error>;
97}
98
99/// Production bulk writer factory for an already connected SQL Server client.
100#[allow(dead_code)]
101pub(crate) struct MssqlConnectedBulkWriterFactory<'client> {
102    client: &'client mut arrow_tiberius::ConnectedMssqlClient,
103}
104
105impl<'client> MssqlConnectedBulkWriterFactory<'client> {
106    /// Wraps the already connected SQL Server client used for lifecycle work.
107    #[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/// Planned inputs for constructing one SQL Server bulk writer.
133#[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    /// Builds writer initialization inputs from a previously prepared target.
146    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    /// Returns the selected output name.
164    #[must_use]
165    pub(crate) fn output_name(&self) -> &str {
166        &self.output_name
167    }
168
169    /// Returns the prepared SQL Server table identity.
170    #[must_use]
171    pub(crate) const fn table(&self) -> &arrow_tiberius::TableName {
172        &self.table
173    }
174
175    /// Returns the planned schema mappings passed to the writer.
176    #[must_use]
177    pub(crate) fn mappings(&self) -> &[arrow_tiberius::SchemaMapping] {
178        &self.mappings
179    }
180
181    /// Returns the write options passed to the writer.
182    #[must_use]
183    pub(crate) const fn options(&self) -> MssqlWriteOptions {
184        self.options
185    }
186
187    /// Returns the prepared lifecycle action that permits writer initialization.
188    #[must_use]
189    pub(crate) const fn prepared_action(&self) -> MssqlPreparedTargetAction {
190        self.prepared_action
191    }
192
193    /// Returns cleanup state if writer initialization fails.
194    #[must_use]
195    pub(crate) const fn cleanup(&self) -> MssqlTargetCleanupStatus {
196        self.cleanup
197    }
198}
199
200/// Phase of one-output SQL Server write execution.
201#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
202pub enum MssqlWritePhase {
203    /// Establish the SQL Server connection.
204    Connect,
205    /// Execute target lifecycle preparation before writer construction.
206    PrepareTargetLifecycle,
207    /// Construct the SQL Server bulk writer and validate target metadata.
208    InitializeWriter,
209    /// Poll the selected output batch stream.
210    PollBatchStream,
211    /// Validate an incoming batch schema against the planned schema.
212    ValidateBatchSchema,
213    /// Write an accepted batch into SQL Server.
214    WriteBatch,
215    /// Finalize the SQL Server bulk writer.
216    Finalize,
217    /// Validate target-side write evidence after finalize succeeds.
218    Validation,
219    /// Atomically swap a staged replace table into the final target name.
220    SwapTarget,
221    /// Clean up a Delta Funnel-created target after a later failure.
222    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/// Returns Delta Funnel's default SQL Server write options.
243#[must_use]
244pub fn default_mssql_write_options() -> MssqlWriteOptions {
245    MssqlWriteOptions {
246        backend: arrow_tiberius::WriteBackend::DirectRawBulk,
247        ..MssqlWriteOptions::default()
248    }
249}
250
251/// Builds write options from a planned SQL Server output target.
252#[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/// Initializes one SQL Server bulk writer after target lifecycle preparation.
263#[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/// Initializes one SQL Server bulk writer through an injected factory.
280#[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
303/// Validates a runtime Arrow schema against a planned SQL Server output.
304///
305/// Delta Funnel owns the output context and redacted report shape, while the
306/// schema contract comparison is delegated to `arrow-tiberius`.
307pub 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
331/// Validates a runtime `RecordBatch` schema against a planned SQL Server output.
332///
333/// This helper validates `batch.schema()` before row writes and does not inspect
334/// row values, connect to SQL Server, or construct a writer.
335pub 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/// Writes one planned SQL Server output through an injected bulk-load writer.
537#[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}