Skip to main content

faucet_sink_mssql/
sink.rs

1//! The MSSQL [`Sink`] implementation — connection pool, transaction-wrapped
2//! multi-row `INSERT` with 2100-parameter auto-splitting, and row-isolation
3//! partial-failure handling for DLQ routing.
4
5use std::sync::Mutex;
6use std::time::Duration;
7
8use async_trait::async_trait;
9use faucet_core::check::{CheckContext, CheckReport, Probe};
10use faucet_core::{FaucetError, RowOutcome, Sink};
11use serde_json::Value;
12use tiberius::ToSql;
13
14use faucet_common_mssql::{MssqlPool, MssqlPooledConnection, build_pool, quote_ident_mssql};
15
16/// Width of the watermark table's `scope` PRIMARY KEY column. SQL Server's index
17/// key budget is 900 bytes = 450 UTF-16 chars, so the column cannot be
18/// `NVARCHAR(MAX)`; a longer scope is shortened to a digest form rather than
19/// erroring or truncating onto another row's watermark (#456 L1).
20const SCOPE_COL_WIDTH: usize = 450;
21
22/// Fit a pipeline scope into [`SCOPE_COL_WIDTH`].
23fn scope_key(scope: &str) -> String {
24    faucet_core::idempotency::scope_key(scope, SCOPE_COL_WIDTH)
25}
26
27use crate::config::{MssqlColumnMapping, MssqlSinkConfig};
28use crate::encode::{
29    BoundParam, auto_row_params, build_insert_sql, build_merge, build_merge_delete,
30    max_rows_per_insert, resolve_insert_columns,
31};
32
33/// Microsoft SQL Server sink.
34pub struct MssqlSink {
35    config: MssqlSinkConfig,
36    pool: MssqlPool,
37    table_quoted: String,
38    /// Cached writable (non-IDENTITY) columns for `auto_columns` mode.
39    columns_cache: Mutex<Option<Vec<String>>>,
40}
41
42impl MssqlSink {
43    /// Connect, validate, build the pool, and (in `json_column` + `create_table`
44    /// mode) create the table if it doesn't exist.
45    pub async fn new(config: MssqlSinkConfig) -> Result<Self, FaucetError> {
46        config.validate()?;
47        config.write.validate()?;
48        if !matches!(config.write.write_mode, faucet_core::WriteMode::Append)
49            && !matches!(
50                config.column_mapping,
51                MssqlColumnMapping::AutoColumns { .. }
52            )
53        {
54            return Err(FaucetError::Config(
55                "mssql sink: write_mode upsert/delete requires column_mapping: auto_columns \
56                 (key columns must be real columns, not inside a JSON column)"
57                    .into(),
58            ));
59        }
60        let table_quoted = quote_table(&config.table)?;
61        let pool = build_pool(&config.connection, config.max_connections).await?;
62
63        let sink = Self {
64            config,
65            pool,
66            table_quoted,
67            columns_cache: Mutex::new(None),
68        };
69        sink.maybe_create_table().await?;
70        Ok(sink)
71    }
72
73    fn timeout(&self) -> Option<Duration> {
74        match self.config.statement_timeout_secs {
75            0 => None,
76            secs => Some(Duration::from_secs(secs)),
77        }
78    }
79
80    async fn maybe_create_table(&self) -> Result<(), FaucetError> {
81        if !self.config.create_table {
82            return Ok(());
83        }
84        let MssqlColumnMapping::JsonColumn { column } = &self.config.column_mapping else {
85            return Ok(()); // validated: create_table only with json_column
86        };
87        let col = quote_ident_mssql(column)?;
88        let sql = format!(
89            "IF OBJECT_ID(N'{}', N'U') IS NULL \
90             CREATE TABLE {} (id BIGINT IDENTITY(1,1) PRIMARY KEY, {} NVARCHAR(MAX))",
91            self.config.table.replace('\'', "''"),
92            self.table_quoted,
93            col
94        );
95        let mut conn = self.checkout().await?;
96        conn.simple_query(sql.as_str())
97            .await
98            .map_err(|e| FaucetError::Sink(format!("MSSQL create_table failed: {e}")))?
99            .into_results()
100            .await
101            .map_err(|e| FaucetError::Sink(format!("MSSQL create_table failed: {e}")))?;
102        Ok(())
103    }
104
105    async fn checkout(&self) -> Result<MssqlPooledConnection<'_>, FaucetError> {
106        self.pool
107            .get()
108            .await
109            .map_err(|e| FaucetError::Sink(format!("MSSQL pool checkout failed: {e}")))
110    }
111
112    /// Writable (non-IDENTITY) table columns, discovered once and cached.
113    async fn insertable_columns(&self) -> Result<Vec<String>, FaucetError> {
114        if let Some(cols) = self.columns_cache.lock().expect("columns mutex").clone() {
115            return Ok(cols);
116        }
117        let cols = self.discover_columns().await?;
118        *self.columns_cache.lock().expect("columns mutex") = Some(cols.clone());
119        Ok(cols)
120    }
121
122    async fn discover_columns(&self) -> Result<Vec<String>, FaucetError> {
123        let mut conn = self.checkout().await?;
124        let table: &str = &self.config.table;
125        let rows = conn
126            .query(
127                "SELECT c.name AS name FROM sys.columns c \
128                 WHERE c.object_id = OBJECT_ID(@P1) AND c.is_identity = 0 \
129                 ORDER BY c.column_id",
130                &[&table],
131            )
132            .await
133            .map_err(|e| FaucetError::Sink(format!("MSSQL column discovery failed: {e}")))?
134            .into_first_result()
135            .await
136            .map_err(|e| FaucetError::Sink(format!("MSSQL column discovery failed: {e}")))?;
137
138        let mut cols = Vec::with_capacity(rows.len());
139        for row in &rows {
140            if let Some(name) = row.get::<&str, _>("name") {
141                cols.push(name.to_string());
142            }
143        }
144        if cols.is_empty() {
145            return Err(FaucetError::Sink(format!(
146                "MSSQL table '{}' has no writable columns or does not exist",
147                self.config.table
148            )));
149        }
150        Ok(cols)
151    }
152
153    /// Discover each column's name, system type name, and nullability for the
154    /// target relation. Returns `(name, type_name, is_nullable)` in column order,
155    /// or an empty vec when the table does not exist / has no columns.
156    async fn discover_column_types(&self) -> Result<Vec<(String, String, bool)>, FaucetError> {
157        let mut conn = self.checkout().await?;
158        let table: &str = &self.config.table;
159        let rows = conn
160            .query(
161                "SELECT c.name AS name, ty.name AS type_name, c.is_nullable AS is_nullable \
162                 FROM sys.columns c \
163                 JOIN sys.types ty ON ty.user_type_id = c.user_type_id \
164                 WHERE c.object_id = OBJECT_ID(@P1) \
165                 ORDER BY c.column_id",
166                &[&table],
167            )
168            .await
169            .map_err(|e| FaucetError::Sink(format!("MSSQL schema query failed: {e}")))?
170            .into_first_result()
171            .await
172            .map_err(|e| FaucetError::Sink(format!("MSSQL schema query failed: {e}")))?;
173
174        let mut cols = Vec::with_capacity(rows.len());
175        for row in &rows {
176            let name = row.get::<&str, _>("name");
177            let type_name = row.get::<&str, _>("type_name");
178            // `is_nullable` is a SQL Server `bit` — tiberius decodes it as a bool.
179            let is_nullable = row.get::<bool, _>("is_nullable").unwrap_or(true);
180            if let (Some(name), Some(type_name)) = (name, type_name) {
181                cols.push((name.to_string(), type_name.to_string(), is_nullable));
182            }
183        }
184        Ok(cols)
185    }
186
187    /// Resolve the column list + per-row owned params for one chunk.
188    /// Returns `None` when there is nothing to insert (e.g. auto_columns with no
189    /// matching keys).
190    async fn prepare_chunk(
191        &self,
192        chunk: &[Value],
193    ) -> Result<Option<(Vec<String>, Vec<Vec<BoundParam>>)>, FaucetError> {
194        match &self.config.column_mapping {
195            MssqlColumnMapping::JsonColumn { column } => {
196                let cols = vec![column.clone()];
197                let rows: Vec<Vec<BoundParam>> = chunk
198                    .iter()
199                    .map(|r| {
200                        serde_json::to_string(r)
201                            .map(|s| vec![BoundParam::Str(s)])
202                            .map_err(|e| {
203                                FaucetError::Sink(format!(
204                                    "MSSQL json_column: failed to serialize record to JSON: {e}"
205                                ))
206                            })
207                    })
208                    .collect::<Result<_, _>>()?;
209                Ok(Some((cols, rows)))
210            }
211            MssqlColumnMapping::AutoColumns { on_unknown_field } => {
212                let insertable = self.insertable_columns().await?;
213                let cols = resolve_insert_columns(&insertable, chunk, *on_unknown_field)?;
214                if cols.is_empty() {
215                    return Ok(None);
216                }
217                let rows: Vec<Vec<BoundParam>> =
218                    chunk.iter().map(|r| auto_row_params(r, &cols)).collect();
219                Ok(Some((cols, rows)))
220            }
221        }
222    }
223
224    /// Insert rows within an **already-open** transaction — caller owns
225    /// `BEGIN TRAN`/`COMMIT TRAN`/`ROLLBACK TRAN`.  Splits into ≤2100-param
226    /// sub-INSERTs but does NOT issue any transaction-control statements.
227    ///
228    /// Used by `write_batch_idempotent` so the data INSERTs and the commit-token
229    /// MERGE share one externally-managed transaction.
230    ///
231    /// Returns `Err((error, timed_out))`. When `timed_out` is `true` the `exec`
232    /// future was dropped mid-TDS, leaving the connection desynced — the caller
233    /// must NOT issue ROLLBACK on it (mirrors `insert_chunk`).
234    async fn insert_rows_no_txn(
235        &self,
236        conn: &mut MssqlPooledConnection<'_>,
237        cols: &[String],
238        rows: &[Vec<BoundParam>],
239    ) -> Result<usize, (FaucetError, bool)> {
240        if rows.is_empty() {
241            return Ok(0);
242        }
243        let cols_quoted: Vec<String> = cols
244            .iter()
245            .map(|c| quote_ident_mssql(c))
246            .collect::<Result<_, _>>()
247            .map_err(|e| (e, false))?;
248        let per_insert = max_rows_per_insert(cols_quoted.len());
249        for sub in rows.chunks(per_insert) {
250            let sql = build_insert_sql(&self.table_quoted, &cols_quoted, sub.len());
251            let owned: Vec<&BoundParam> = sub.iter().flatten().collect();
252            let refs: Vec<&dyn ToSql> = owned.iter().map(|p| p.as_tosql()).collect();
253            let exec = async {
254                conn.execute(sql.as_str(), &refs)
255                    .await
256                    .map(|_| ())
257                    .map_err(|e| FaucetError::Sink(format!("MSSQL insert failed: {e}")))
258            };
259            // On timeout the `exec` future is dropped mid-TDS, desyncing the
260            // connection — the caller must NOT issue ROLLBACK on it (mirrors
261            // `insert_chunk`).
262            let (result, timed_out) = match self.timeout() {
263                Some(t) => match tokio::time::timeout(t, exec).await {
264                    Ok(inner) => (inner, false),
265                    Err(_) => (
266                        Err(FaucetError::Sink("MSSQL insert timed out".into())),
267                        true,
268                    ),
269                },
270                None => (exec.await, false),
271            };
272            if let Err(e) = result {
273                return Err((e, timed_out));
274            }
275        }
276        Ok(rows.len())
277    }
278
279    /// Ensure the per-sink commit-token watermark table exists.
280    /// Uses `IF OBJECT_ID … IS NULL CREATE TABLE` so it is idempotent.
281    async fn ensure_commit_table(
282        &self,
283        conn: &mut MssqlPooledConnection<'_>,
284    ) -> Result<(), FaucetError> {
285        // NVARCHAR(450) is the maximum SQL Server index key byte budget (900 B /
286        // 2 bytes-per-char = 450 chars) — fine for the PRIMARY KEY `scope`. The
287        // `token` column is `NVARCHAR(MAX)` because a `#291` commit token embeds
288        // the page's resume bookmark (`{20-digit seq}#{bookmark-json}`) and
289        // exceeds the old `NVARCHAR(32)`, which raised "String or binary data
290        // would be truncated" and broke exactly-once delivery (audit #321 C4).
291        // The table / column names are the fixed constants — no user-controlled
292        // input in this string.
293        let sql = format!(
294            "IF OBJECT_ID(N'{tbl}', N'U') IS NULL \
295             CREATE TABLE [{tbl}] ([scope] NVARCHAR({w}) PRIMARY KEY, \
296             [token] NVARCHAR(MAX) NOT NULL, \
297             [updated_at] DATETIME2 DEFAULT SYSUTCDATETIME())",
298            tbl = faucet_core::idempotency::COMMIT_TOKEN_TABLE,
299            w = SCOPE_COL_WIDTH,
300        );
301        control(conn, &sql).await
302    }
303
304    /// Insert one chunk, splitting into ≤2100-parameter sub-INSERTs wrapped in a
305    /// single transaction (when `transaction_per_batch`). Returns rows inserted.
306    async fn insert_chunk(
307        &self,
308        conn: &mut MssqlPooledConnection<'_>,
309        cols: &[String],
310        rows: &[Vec<BoundParam>],
311    ) -> Result<usize, FaucetError> {
312        if rows.is_empty() {
313            return Ok(0);
314        }
315        let cols_quoted: Vec<String> = cols
316            .iter()
317            .map(|c| quote_ident_mssql(c))
318            .collect::<Result<_, _>>()?;
319        let per_insert = max_rows_per_insert(cols_quoted.len());
320
321        // Wrap the chunk in a transaction when configured, OR whenever it spans
322        // more than one ≤2100-param sub-INSERT. Under autocommit a multi-sub
323        // chunk commits each sub-INSERT independently, so a later failure leaves
324        // earlier sub-INSERTs committed — and both the batch-level transient
325        // retry (`write_batch`) and the per-row isolation (`write_batch_partial`)
326        // re-run the whole chunk, duplicating those committed rows (audit #146
327        // H6). Forcing a transaction makes the chunk atomic so re-running is safe.
328        let txn = self.config.transaction_per_batch || rows.len() > per_insert;
329        if txn {
330            control(conn, "BEGIN TRAN").await?;
331        }
332
333        for sub in rows.chunks(per_insert) {
334            let sql = build_insert_sql(&self.table_quoted, &cols_quoted, sub.len());
335            let owned: Vec<&BoundParam> = sub.iter().flatten().collect();
336            let refs: Vec<&dyn ToSql> = owned.iter().map(|p| p.as_tosql()).collect();
337
338            let exec = async {
339                conn.execute(sql.as_str(), &refs)
340                    .await
341                    .map(|_| ())
342                    .map_err(|e| FaucetError::Sink(format!("MSSQL insert failed: {e}")))
343            };
344            // Track whether the failure was a *timeout* specifically. On timeout
345            // the `exec` future is dropped mid-TDS, leaving an unread response on
346            // the wire — the connection is desynced and must NOT be reused (the
347            // pool helper's contract is "drop it"). Issuing ROLLBACK on it would
348            // run on a corrupt stream. A *normal* error leaves the connection in
349            // sync, so ROLLBACK is safe and releases the transaction promptly.
350            let (result, timed_out) = match self.timeout() {
351                Some(t) => match tokio::time::timeout(t, exec).await {
352                    Ok(inner) => (inner, false),
353                    Err(_) => (
354                        Err(FaucetError::Sink("MSSQL insert timed out".into())),
355                        true,
356                    ),
357                },
358                None => (exec.await, false),
359            };
360            if let Err(e) = result {
361                if txn && !timed_out {
362                    let _ = control(conn, "ROLLBACK TRAN").await;
363                }
364                return Err(e);
365            }
366        }
367
368        if txn {
369            control(conn, "COMMIT TRAN").await?;
370        }
371        Ok(rows.len())
372    }
373
374    /// Run a single `MERGE`-upsert / `MERGE`-delete statement on an
375    /// already-open connection, honouring the per-statement timeout. Returns
376    /// `Err((error, timed_out))`; on timeout the connection is desynced and the
377    /// caller must NOT issue ROLLBACK on it (mirrors `insert_rows_no_txn`).
378    async fn exec_merge(
379        &self,
380        conn: &mut MssqlPooledConnection<'_>,
381        sql: &str,
382        refs: &[&dyn ToSql],
383    ) -> Result<(), (FaucetError, bool)> {
384        let exec = async {
385            conn.execute(sql, refs)
386                .await
387                .map(|_| ())
388                .map_err(|e| FaucetError::Sink(format!("MSSQL merge failed: {e}")))
389        };
390        match self.timeout() {
391            Some(t) => match tokio::time::timeout(t, exec).await {
392                Ok(Ok(())) => Ok(()),
393                Ok(Err(e)) => Err((e, false)),
394                Err(_) => Err((FaucetError::Sink("MSSQL merge timed out".into()), true)),
395            },
396            None => exec.await.map_err(|e| (e, false)),
397        }
398    }
399
400    /// Upsert `upserts` into the table via `MERGE`, on an already-open
401    /// connection (the caller owns the transaction). Resolves the column set
402    /// via `resolve_insert_columns`, chunks by `max_rows_per_insert`, and binds
403    /// each row's params row-major exactly as `build_merge` numbers them.
404    async fn upsert_rows_no_txn(
405        &self,
406        conn: &mut MssqlPooledConnection<'_>,
407        upserts: &[Value],
408    ) -> Result<usize, (FaucetError, bool)> {
409        if upserts.is_empty() {
410            return Ok(0);
411        }
412        let MssqlColumnMapping::AutoColumns { on_unknown_field } = &self.config.column_mapping
413        else {
414            // Validated in `new()` — upsert requires auto_columns.
415            return Err((
416                FaucetError::Sink("MSSQL upsert requires column_mapping: auto_columns".into()),
417                false,
418            ));
419        };
420        let insertable = self.insertable_columns().await.map_err(|e| (e, false))?;
421        let cols = resolve_insert_columns(&insertable, upserts, *on_unknown_field)
422            .map_err(|e| (e, false))?;
423        if cols.is_empty() {
424            return Ok(0);
425        }
426        let per_insert = max_rows_per_insert(cols.len());
427        for sub in upserts.chunks(per_insert) {
428            let sql = build_merge(&self.table_quoted, &self.config.write.key, &cols, sub.len())
429                .map_err(|e| (e, false))?;
430            // Bind every row's params concatenated row-major — matches the @PN
431            // numbering build_merge emits.
432            let owned: Vec<BoundParam> =
433                sub.iter().flat_map(|r| auto_row_params(r, &cols)).collect();
434            let refs: Vec<&dyn ToSql> = owned.iter().map(|p| p.as_tosql()).collect();
435            self.exec_merge(conn, &sql, &refs).await?;
436        }
437        Ok(upserts.len())
438    }
439
440    /// Delete the `deletes` key tuples via `MERGE … WHEN MATCHED THEN DELETE`,
441    /// on an already-open connection. Chunks by `max_rows_per_insert(key.len())`
442    /// and binds each key tuple's values in `key` order, row-major.
443    async fn delete_keys_no_txn(
444        &self,
445        conn: &mut MssqlPooledConnection<'_>,
446        deletes: &[faucet_core::KeyTuple],
447    ) -> Result<usize, (FaucetError, bool)> {
448        if deletes.is_empty() {
449            return Ok(0);
450        }
451        let key = &self.config.write.key;
452        let per = max_rows_per_insert(key.len());
453        for chunk in deletes.chunks(per) {
454            let sql =
455                build_merge_delete(&self.table_quoted, key, chunk.len()).map_err(|e| (e, false))?;
456            // Bind each tuple's values in key order, row-major.
457            let owned: Vec<BoundParam> = chunk
458                .iter()
459                .flat_map(|kt| kt.0.iter().map(|(_, v)| BoundParam::from_value(v)))
460                .collect();
461            let refs: Vec<&dyn ToSql> = owned.iter().map(|p| p.as_tosql()).collect();
462            self.exec_merge(conn, &sql, &refs).await?;
463        }
464        Ok(deletes.len())
465    }
466
467    /// Apply a planned upsert/delete batch atomically: upserts then deletes,
468    /// wrapped in a single `BEGIN TRAN` / `COMMIT TRAN` so they commit together
469    /// (last-write-wins dedup already collapsed conflicting ops in the planner).
470    async fn apply_plan(&self, plan: &faucet_core::WritePlan) -> Result<usize, FaucetError> {
471        let mut conn = self.checkout().await?;
472        control(&mut conn, "BEGIN TRAN").await?;
473
474        let mut affected = 0usize;
475        match self.upsert_rows_no_txn(&mut conn, &plan.upserts).await {
476            Ok(n) => affected += n,
477            Err((e, timed_out)) => {
478                if !timed_out {
479                    let _ = control(&mut conn, "ROLLBACK TRAN").await;
480                }
481                return Err(e);
482            }
483        }
484        match self.delete_keys_no_txn(&mut conn, &plan.deletes).await {
485            Ok(n) => affected += n,
486            Err((e, timed_out)) => {
487                if !timed_out {
488                    let _ = control(&mut conn, "ROLLBACK TRAN").await;
489                }
490                return Err(e);
491            }
492        }
493
494        control(&mut conn, "COMMIT TRAN").await?;
495        Ok(affected)
496    }
497}
498
499/// Run a transaction-control statement and drain its (empty) result.
500async fn control(conn: &mut MssqlPooledConnection<'_>, stmt: &str) -> Result<(), FaucetError> {
501    conn.simple_query(stmt)
502        .await
503        .map_err(|e| FaucetError::Sink(format!("MSSQL {stmt} failed: {e}")))?
504        .into_results()
505        .await
506        .map_err(|e| FaucetError::Sink(format!("MSSQL {stmt} failed: {e}")))?;
507    Ok(())
508}
509
510/// Map a [`SqlBaseType`](faucet_core::SqlBaseType) to the MSSQL type keyword used when adding/widening a
511/// column during schema evolution (issue #194). Integers widen to `BIGINT` and
512/// floats to `FLOAT` so a later, wider value never overflows a narrower column;
513/// text/json land in `NVARCHAR(MAX)`.
514fn mssql_keyword(t: faucet_core::SqlBaseType) -> &'static str {
515    use faucet_core::SqlBaseType::*;
516    match t {
517        Integer => "BIGINT",
518        Double => "FLOAT",
519        Boolean => "BIT",
520        Text => "NVARCHAR(MAX)",
521        Json => "NVARCHAR(MAX)",
522    }
523}
524
525/// Build an idempotent `ADD COLUMN`. T-SQL has no `ADD COLUMN IF NOT EXISTS`, so
526/// guard with `IF NOT EXISTS (SELECT 1 FROM sys.columns …)`.
527///
528/// `table_quoted` is the already-bracket-quoted relation (`[dbo].[events]`).
529/// `table_literal` is the bare (un-quoted) table name used inside the
530/// `OBJECT_ID(N'…')` lookup — the caller passes `self.config.table` so it
531/// resolves the same relation the DDL targets. Both `table_literal` and `col`
532/// are single-quote-escaped for their `N'…'` string literals; `col` is also
533/// bracket-quoted for the `ADD` clause via [`quote_ident_mssql`].
534fn build_add_column_sql(
535    table_quoted: &str,
536    table_literal: &str,
537    col: &str,
538    t: faucet_core::SqlBaseType,
539) -> Result<String, FaucetError> {
540    let qcol = quote_ident_mssql(col)?;
541    Ok(format!(
542        "IF NOT EXISTS (SELECT 1 FROM sys.columns \
543         WHERE object_id = OBJECT_ID(N'{}') AND name = N'{}') \
544         ALTER TABLE {table_quoted} ADD {qcol} {}",
545        table_literal.replace('\'', "''"),
546        col.replace('\'', "''"),
547        mssql_keyword(t),
548    ))
549}
550
551/// `ALTER TABLE <ref> ALTER COLUMN <col> <kw>` — widen an existing column's
552/// type. Naturally idempotent (re-running the same type change is a no-op).
553fn build_alter_type_sql(
554    table_quoted: &str,
555    col: &str,
556    t: faucet_core::SqlBaseType,
557) -> Result<String, FaucetError> {
558    let qcol = quote_ident_mssql(col)?;
559    Ok(format!(
560        "ALTER TABLE {table_quoted} ALTER COLUMN {qcol} {}",
561        mssql_keyword(t),
562    ))
563}
564
565/// `ALTER TABLE <ref> ALTER COLUMN <col> <kw> NULL` — relax a NOT NULL
566/// constraint. MSSQL requires re-stating the column's current type when toggling
567/// nullability, so `kw` must be the column's existing type keyword. Naturally
568/// idempotent.
569fn build_alter_null_sql(table_quoted: &str, col: &str, kw: &str) -> Result<String, FaucetError> {
570    let qcol = quote_ident_mssql(col)?;
571    Ok(format!(
572        "ALTER TABLE {table_quoted} ALTER COLUMN {qcol} {kw} NULL",
573    ))
574}
575
576/// Map an MSSQL system type name (`sys.types.name`, e.g. `bigint`, `float`,
577/// `bit`, `nvarchar`) back to a JSON-Schema type fragment so
578/// [`MssqlSink::current_schema`] round-trips with [`faucet_core::diff_schema`].
579/// `nullable` reflects `sys.columns.is_nullable`.
580fn mssql_type_to_json_schema(type_name: &str, nullable: bool) -> Value {
581    let base = match type_name.to_ascii_lowercase().as_str() {
582        "bigint" | "int" | "smallint" | "tinyint" => "integer",
583        "float" | "real" | "decimal" | "numeric" | "money" | "smallmoney" => "number",
584        "bit" => "boolean",
585        _ => "string",
586    };
587    if nullable {
588        serde_json::json!({ "type": [base, "null"] })
589    } else {
590        serde_json::json!({ "type": base })
591    }
592}
593
594/// Quote a (possibly schema-qualified) table name: `dbo.events` → `[dbo].[events]`.
595fn quote_table(table: &str) -> Result<String, FaucetError> {
596    let parts: Vec<String> = table
597        .split('.')
598        .map(quote_ident_mssql)
599        .collect::<Result<_, _>>()?;
600    Ok(parts.join("."))
601}
602
603/// Heuristic for transient errors that warrant a batch-level retry / outer-Err
604/// propagation rather than per-row DLQ isolation.
605fn is_transient_error(msg: &str) -> bool {
606    let m = msg.to_ascii_lowercase();
607    m.contains("deadlock")
608        || m.contains("timed out")
609        || m.contains("timeout")
610        || m.contains("connection")
611        || m.contains("transport")
612        || m.contains("link failure")
613        || m.contains("1205")
614}
615
616const TRANSIENT_RETRIES: usize = 3;
617
618#[async_trait]
619impl Sink for MssqlSink {
620    async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
621        if records.is_empty() {
622            return Ok(0);
623        }
624
625        // Non-append modes: plan the writes and apply upserts + deletes
626        // atomically. (NOTE: write_batch_partial upsert routing is handled in
627        // the DLQ task; write_batch_idempotent in the exactly-once task.)
628        if !matches!(self.config.write.write_mode, faucet_core::WriteMode::Append) {
629            let plan = faucet_core::plan_writes(records, &self.config.write);
630            if let Some((idx, msg)) = plan.failed.first() {
631                return Err(FaucetError::Sink(format!(
632                    "mssql {}: row {idx}: {msg}",
633                    self.config.write.write_mode.as_str()
634                )));
635            }
636            let total = self.apply_plan(&plan).await?;
637            tracing::info!(
638                table = %self.config.table,
639                mode = self.config.write.write_mode.as_str(),
640                rows = total,
641                "MSSQL write complete"
642            );
643            return Ok(total);
644        }
645
646        let chunks: Vec<&[Value]> = if self.config.batch_size == 0 {
647            vec![records]
648        } else {
649            records.chunks(self.config.batch_size).collect()
650        };
651
652        let mut total = 0usize;
653        for chunk in chunks {
654            let Some((cols, rows)) = self.prepare_chunk(chunk).await? else {
655                continue;
656            };
657            // Bounded retry on transient (deadlock / lock-timeout) errors.
658            let mut attempt = 0;
659            loop {
660                let mut conn = self.checkout().await?;
661                match self.insert_chunk(&mut conn, &cols, &rows).await {
662                    Ok(n) => {
663                        total += n;
664                        break;
665                    }
666                    Err(e) if is_transient_error(&e.to_string()) && attempt < TRANSIENT_RETRIES => {
667                        attempt += 1;
668                        let backoff = Duration::from_millis(50 * (1 << attempt));
669                        tracing::warn!(attempt, error = %e, "MSSQL transient error; retrying batch");
670                        tokio::time::sleep(backoff).await;
671                    }
672                    Err(e) => return Err(e),
673                }
674            }
675        }
676        tracing::info!(table = %self.config.table, rows = total, "MSSQL write complete");
677        Ok(total)
678    }
679
680    async fn write_batch_partial(&self, records: &[Value]) -> Result<Vec<RowOutcome>, FaucetError> {
681        if records.is_empty() {
682            return Ok(Vec::new());
683        }
684
685        // Upsert/delete: apply the good rows (upserts + deletes) and route only
686        // the rows whose key could not be extracted (missing / null key) to the
687        // DLQ per-row. The append path below keeps its row-isolation behaviour.
688        if !matches!(self.config.write.write_mode, faucet_core::WriteMode::Append) {
689            let plan = faucet_core::plan_writes(records, &self.config.write);
690            self.apply_plan(&plan).await?;
691
692            let mut outcomes: Vec<RowOutcome> = records.iter().map(|_| Ok(())).collect();
693            for (idx, msg) in &plan.failed {
694                outcomes[*idx] = Err(FaucetError::Sink(format!(
695                    "mssql {}: {msg}",
696                    self.config.write.write_mode.as_str()
697                )));
698            }
699            return Ok(outcomes);
700        }
701
702        let chunks: Vec<&[Value]> = if self.config.batch_size == 0 {
703            vec![records]
704        } else {
705            records.chunks(self.config.batch_size).collect()
706        };
707
708        let mut outcomes: Vec<RowOutcome> = Vec::with_capacity(records.len());
709        for chunk in chunks {
710            let Some((cols, rows)) = self.prepare_chunk(chunk).await? else {
711                // Nothing to insert for this chunk (no matching columns): the
712                // rows were effectively dropped per on_unknown_field; report Ok.
713                outcomes.extend(chunk.iter().map(|_| Ok(())));
714                continue;
715            };
716
717            let mut conn = self.checkout().await?;
718            match self.insert_chunk(&mut conn, &cols, &rows).await {
719                Ok(_) => outcomes.extend(chunk.iter().map(|_| Ok(()))),
720                Err(e) if is_transient_error(&e.to_string()) => {
721                    // Infra/transient — not row-specific. Propagate so the
722                    // pipeline's on_batch_error policy decides.
723                    return Err(e);
724                }
725                Err(_) if !self.config.isolate_row_failures => {
726                    // One bad row fails the whole batch (caller's choice).
727                    return Err(FaucetError::Sink(
728                        "MSSQL batch insert failed and isolate_row_failures is disabled".into(),
729                    ));
730                }
731                Err(_) => {
732                    // Row-isolate: retry each row alone to find the offender.
733                    for (i, row) in rows.iter().enumerate() {
734                        let single = std::slice::from_ref(row);
735                        let single_cols = cols.clone();
736                        match self.insert_chunk(&mut conn, &single_cols, single).await {
737                            Ok(_) => outcomes.push(Ok(())),
738                            Err(e) if is_transient_error(&e.to_string()) => return Err(e),
739                            Err(e) => {
740                                tracing::warn!(row = i, error = %e, "MSSQL row rejected; routing to DLQ");
741                                outcomes.push(Err(e));
742                            }
743                        }
744                    }
745                }
746            }
747        }
748        Ok(outcomes)
749    }
750
751    async fn flush(&self) -> Result<(), FaucetError> {
752        Ok(())
753    }
754
755    fn supports_idempotent_writes(&self) -> bool {
756        true
757    }
758
759    fn supported_write_modes(&self) -> &'static [faucet_core::WriteMode] {
760        &[
761            faucet_core::WriteMode::Append,
762            faucet_core::WriteMode::Upsert,
763            faucet_core::WriteMode::Delete,
764        ]
765    }
766
767    fn dedups_by_key(&self) -> bool {
768        self.config.write.dedups_by_key()
769    }
770
771    fn supports_schema_evolution(&self) -> bool {
772        true
773    }
774
775    /// Read the live destination schema from `sys.columns` as an
776    /// `infer_schema`-shaped object (`{"type":"object","properties":{…}}`), or
777    /// `None` when the target table does not exist yet (issue #194).
778    async fn current_schema(&self) -> Result<Option<Value>, FaucetError> {
779        let cols = self.discover_column_types().await?;
780        if cols.is_empty() {
781            return Ok(None); // table does not exist yet
782        }
783        let mut props = serde_json::Map::new();
784        for (name, type_name, nullable) in cols {
785            props.insert(name, mssql_type_to_json_schema(&type_name, nullable));
786        }
787        Ok(Some(
788            serde_json::json!({ "type": "object", "properties": props }),
789        ))
790    }
791
792    /// Apply an additive schema evolution (new columns, lossless widenings,
793    /// nullability relaxations) to the destination table. Idempotent — the
794    /// `ADD` is guarded with `IF NOT EXISTS (SELECT 1 FROM sys.columns …)`, and
795    /// re-running the same `ALTER COLUMN` type / `… NULL` is a no-op (issue #194).
796    async fn evolve_schema(
797        &self,
798        evolution: &faucet_core::SchemaEvolution,
799    ) -> Result<(), FaucetError> {
800        let mut conn = self.checkout().await?;
801
802        for c in &evolution.additions {
803            let t =
804                faucet_core::json_schema_base_type(&c.to).unwrap_or(faucet_core::SqlBaseType::Text);
805            let sql = build_add_column_sql(&self.table_quoted, &self.config.table, &c.name, t)?;
806            control(&mut conn, &sql).await.map_err(|e| {
807                FaucetError::Sink(format!("MSSQL ADD COLUMN {} failed: {e}", c.name))
808            })?;
809        }
810        for c in &evolution.widenings {
811            let t =
812                faucet_core::json_schema_base_type(&c.to).unwrap_or(faucet_core::SqlBaseType::Text);
813            let sql = build_alter_type_sql(&self.table_quoted, &c.name, t)?;
814            control(&mut conn, &sql).await.map_err(|e| {
815                FaucetError::Sink(format!("MSSQL ALTER COLUMN {} failed: {e}", c.name))
816            })?;
817        }
818        if !evolution.relax_nullability.is_empty() {
819            // Re-emitting the column as NULL requires its CURRENT type keyword —
820            // derive it from the live schema.
821            let current: std::collections::HashMap<String, &'static str> = self
822                .discover_column_types()
823                .await?
824                .into_iter()
825                .map(|(name, type_name, _)| {
826                    let base = faucet_core::json_schema_base_type(&mssql_type_to_json_schema(
827                        &type_name, false,
828                    ))
829                    .unwrap_or(faucet_core::SqlBaseType::Text);
830                    (name, mssql_keyword(base))
831                })
832                .collect();
833            for col in &evolution.relax_nullability {
834                let Some(kw) = current.get(col) else {
835                    // Column not found in the live schema — nothing to relax.
836                    continue;
837                };
838                let sql = build_alter_null_sql(&self.table_quoted, col, kw)?;
839                control(&mut conn, &sql).await.map_err(|e| {
840                    FaucetError::Sink(format!("MSSQL relax NULL {col} failed: {e}"))
841                })?;
842            }
843        }
844
845        // Columns changed — drop the cached AutoColumns set so the next write
846        // re-discovers them (a newly-added column must be picked up).
847        *self.columns_cache.lock().expect("columns mutex") = None;
848        Ok(())
849    }
850
851    async fn last_committed_token(&self, scope: &str) -> Result<Option<String>, FaucetError> {
852        let mut conn = self.checkout().await?;
853        self.ensure_commit_table(&mut conn).await?;
854        let scope_owned = scope_key(scope);
855        let rows = conn
856            .query(
857                &format!(
858                    "SELECT [token] FROM [{}] WHERE [scope] = @P1",
859                    faucet_core::idempotency::COMMIT_TOKEN_TABLE
860                ),
861                &[&scope_owned],
862            )
863            .await
864            .map_err(|e| FaucetError::Sink(format!("MSSQL token read failed: {e}")))?
865            .into_first_result()
866            .await
867            .map_err(|e| FaucetError::Sink(format!("MSSQL token read failed: {e}")))?;
868        Ok(rows
869            .first()
870            .and_then(|r| r.get::<&str, _>("token"))
871            .map(str::to_string))
872    }
873
874    async fn write_batch_idempotent(
875        &self,
876        records: &[Value],
877        scope: &str,
878        token: &str,
879    ) -> Result<usize, FaucetError> {
880        // For upsert/delete modes, plan the page before opening the transaction
881        // so a key-extraction failure aborts without leaving an open tx.
882        let plan = if matches!(self.config.write.write_mode, faucet_core::WriteMode::Append) {
883            None
884        } else {
885            let plan = faucet_core::plan_writes(records, &self.config.write);
886            if let Some((idx, msg)) = plan.failed.first() {
887                return Err(FaucetError::Sink(format!(
888                    "mssql {}: row {idx}: {msg}",
889                    self.config.write.write_mode.as_str()
890                )));
891            }
892            Some(plan)
893        };
894
895        let mut conn = self.checkout().await?;
896        self.ensure_commit_table(&mut conn).await?;
897        control(&mut conn, "BEGIN TRAN").await?;
898
899        // Data write and the commit-token MERGE share ONE transaction so the
900        // page is committed atomically with its watermark. For upsert/delete the
901        // planned upserts/deletes run via the same no-txn MERGE helpers used by
902        // the append path's INSERTs, inside this same BEGIN TRAN.
903        let written = match &plan {
904            Some(plan) => {
905                let mut affected = 0usize;
906                match self.upsert_rows_no_txn(&mut conn, &plan.upserts).await {
907                    Ok(n) => affected += n,
908                    Err((e, timed_out)) => {
909                        if !timed_out {
910                            let _ = control(&mut conn, "ROLLBACK TRAN").await;
911                        }
912                        return Err(e);
913                    }
914                }
915                match self.delete_keys_no_txn(&mut conn, &plan.deletes).await {
916                    Ok(n) => affected += n,
917                    Err((e, timed_out)) => {
918                        if !timed_out {
919                            let _ = control(&mut conn, "ROLLBACK TRAN").await;
920                        }
921                        return Err(e);
922                    }
923                }
924                affected
925            }
926            None => match self.prepare_chunk(records).await {
927                Ok(Some((cols, rows))) => {
928                    match self.insert_rows_no_txn(&mut conn, &cols, &rows).await {
929                        Ok(n) => n,
930                        Err((e, timed_out)) => {
931                            // Desynced connection on timeout — ROLLBACK would run on a
932                            // corrupt stream (mirrors insert_chunk). Drop the conn instead.
933                            if !timed_out {
934                                let _ = control(&mut conn, "ROLLBACK TRAN").await;
935                            }
936                            return Err(e);
937                        }
938                    }
939                }
940                Ok(None) => 0,
941                Err(e) => {
942                    let _ = control(&mut conn, "ROLLBACK TRAN").await;
943                    return Err(e);
944                }
945            },
946        };
947
948        // UPSERT the commit token atomically with the data rows.
949        let merge = format!(
950            "MERGE [{tbl}] AS t \
951             USING (SELECT @P1 AS [scope], @P2 AS [token]) AS s \
952             ON t.[scope] = s.[scope] \
953             WHEN MATCHED THEN UPDATE SET t.[token] = s.[token], t.[updated_at] = SYSUTCDATETIME() \
954             WHEN NOT MATCHED THEN INSERT ([scope], [token]) VALUES (s.[scope], s.[token]);",
955            tbl = faucet_core::idempotency::COMMIT_TOKEN_TABLE,
956        );
957        let (scope_owned, token_owned) = (scope_key(scope), token.to_string());
958        let refs: Vec<&dyn ToSql> = vec![&scope_owned, &token_owned];
959        if let Err(e) = conn.execute(merge.as_str(), &refs).await {
960            let _ = control(&mut conn, "ROLLBACK TRAN").await;
961            return Err(FaucetError::Sink(format!("MSSQL token merge failed: {e}")));
962        }
963
964        control(&mut conn, "COMMIT TRAN").await?;
965        Ok(written)
966    }
967
968    fn config_schema(&self) -> Value {
969        serde_json::to_value(faucet_core::schema_for!(MssqlSinkConfig))
970            .expect("schema serialization")
971    }
972
973    fn connector_name(&self) -> &'static str {
974        "mssql"
975    }
976
977    fn dataset_uri(&self) -> String {
978        let conn = self
979            .config
980            .connection
981            .connection_url
982            .as_deref()
983            .or(self.config.connection.connection_string.as_deref())
984            .unwrap_or("");
985        format!(
986            "{}?table={}",
987            faucet_core::redact_uri_credentials(conn),
988            self.config.table
989        )
990    }
991
992    async fn check(&self, ctx: &CheckContext) -> Result<CheckReport, FaucetError> {
993        let started = std::time::Instant::now();
994        let probe = match tokio::time::timeout(ctx.timeout, self.pool.get()).await {
995            Ok(Ok(_conn)) => Probe::pass("connect", started.elapsed()),
996            Ok(Err(e)) => Probe::fail_hint(
997                "connect",
998                started.elapsed(),
999                e.to_string(),
1000                "check connection_url / credentials / TLS / that the server is reachable",
1001            ),
1002            Err(_) => Probe::fail_hint(
1003                "connect",
1004                started.elapsed(),
1005                "timed out",
1006                "check connection_url / credentials / TLS / that the server is reachable",
1007            ),
1008        };
1009        Ok(CheckReport::single(probe))
1010    }
1011}
1012
1013#[cfg(test)]
1014mod tests {
1015    use super::*;
1016
1017    // dataset_uri test is skipped: MssqlSink::new() requires a live pool
1018    // (connects to SQL Server in new()), and no offline constructor exists.
1019
1020    #[test]
1021    fn quote_table_handles_schema_qualified() {
1022        assert_eq!(quote_table("dbo.events").unwrap(), "[dbo].[events]");
1023        assert_eq!(quote_table("events").unwrap(), "[events]");
1024        assert_eq!(
1025            quote_table("my.sales.events").unwrap(),
1026            "[my].[sales].[events]"
1027        );
1028    }
1029
1030    #[test]
1031    fn idempotency_constant_names() {
1032        // The commit table and column constants used in ensure_commit_table,
1033        // last_committed_token, and write_batch_idempotent must match the
1034        // canonical values from faucet_core::idempotency.
1035        assert_eq!(
1036            faucet_core::idempotency::COMMIT_TOKEN_TABLE,
1037            "_faucet_commit_token",
1038            "COMMIT_TOKEN_TABLE name changed — update DDL and queries"
1039        );
1040        assert_eq!(
1041            faucet_core::idempotency::COMMIT_TOKEN_SCOPE_COL,
1042            "scope",
1043            "COMMIT_TOKEN_SCOPE_COL name changed — update DDL and queries"
1044        );
1045        assert_eq!(
1046            faucet_core::idempotency::COMMIT_TOKEN_TOKEN_COL,
1047            "token",
1048            "COMMIT_TOKEN_TOKEN_COL name changed — update DDL and queries"
1049        );
1050    }
1051
1052    #[test]
1053    fn mssql_add_column_ddl() {
1054        let sql = build_add_column_sql(
1055            "[dbo].[events]",
1056            "dbo.events",
1057            "email",
1058            faucet_core::SqlBaseType::Text,
1059        )
1060        .unwrap();
1061        assert!(sql.starts_with("IF NOT EXISTS"), "{sql}");
1062        assert!(
1063            sql.contains("ALTER TABLE [dbo].[events] ADD [email] NVARCHAR(MAX)"),
1064            "{sql}"
1065        );
1066        // The OBJECT_ID guard targets the bare table literal.
1067        assert!(sql.contains("OBJECT_ID(N'dbo.events')"), "{sql}");
1068        assert!(sql.contains("name = N'email'"), "{sql}");
1069    }
1070
1071    #[test]
1072    fn mssql_add_column_keyword_per_base_type() {
1073        use faucet_core::SqlBaseType::*;
1074        for (t, kw) in [
1075            (Integer, "BIGINT"),
1076            (Double, "FLOAT"),
1077            (Boolean, "BIT"),
1078            (Text, "NVARCHAR(MAX)"),
1079            (Json, "NVARCHAR(MAX)"),
1080        ] {
1081            let sql = build_add_column_sql("[t]", "t", "c", t).unwrap();
1082            assert!(sql.ends_with(&format!("ADD [c] {kw}")), "{t:?}: {sql}");
1083        }
1084    }
1085
1086    #[test]
1087    fn mssql_add_column_escapes_literals() {
1088        // Single quotes in the table/column literal are doubled for N'…';
1089        // brackets in the identifier are doubled by quote_ident_mssql.
1090        let sql = build_add_column_sql("[d].[t]", "d.o'x", "c'l", faucet_core::SqlBaseType::Text)
1091            .unwrap();
1092        assert!(sql.contains("OBJECT_ID(N'd.o''x')"), "{sql}");
1093        assert!(sql.contains("name = N'c''l'"), "{sql}");
1094        assert!(sql.contains("ADD [c'l] NVARCHAR(MAX)"), "{sql}");
1095    }
1096
1097    #[test]
1098    fn mssql_widen_column_ddl() {
1099        let sql =
1100            build_alter_type_sql("[dbo].[t]", "score", faucet_core::SqlBaseType::Double).unwrap();
1101        assert_eq!(sql, "ALTER TABLE [dbo].[t] ALTER COLUMN [score] FLOAT");
1102    }
1103
1104    #[test]
1105    fn mssql_relax_null_ddl_re_emits_current_type() {
1106        let sql = build_alter_null_sql("[t]", "created_at", "DATETIME2").unwrap();
1107        assert_eq!(
1108            sql,
1109            "ALTER TABLE [t] ALTER COLUMN [created_at] DATETIME2 NULL"
1110        );
1111    }
1112
1113    #[test]
1114    fn mssql_type_round_trips_to_json_schema() {
1115        use serde_json::json;
1116        assert_eq!(
1117            mssql_type_to_json_schema("bigint", false),
1118            json!({"type":"integer"})
1119        );
1120        assert_eq!(
1121            mssql_type_to_json_schema("int", false),
1122            json!({"type":"integer"})
1123        );
1124        assert_eq!(
1125            mssql_type_to_json_schema("float", false),
1126            json!({"type":"number"})
1127        );
1128        assert_eq!(
1129            mssql_type_to_json_schema("decimal", false),
1130            json!({"type":"number"})
1131        );
1132        assert_eq!(
1133            mssql_type_to_json_schema("bit", false),
1134            json!({"type":"boolean"})
1135        );
1136        assert_eq!(
1137            mssql_type_to_json_schema("nvarchar", false),
1138            json!({"type":"string"})
1139        );
1140        // Case-insensitive; nullable widens to a type array.
1141        assert_eq!(
1142            mssql_type_to_json_schema("NVARCHAR", true),
1143            json!({"type":["string","null"]})
1144        );
1145    }
1146
1147    #[test]
1148    fn transient_classifier() {
1149        assert!(is_transient_error(
1150            "Transaction (Process ID 55) was deadlocked"
1151        ));
1152        assert!(is_transient_error(
1153            "Lock request time out period exceeded (1205)"
1154        ));
1155        assert!(is_transient_error("connection reset by peer"));
1156        assert!(!is_transient_error("Violation of PRIMARY KEY constraint"));
1157        assert!(!is_transient_error(
1158            "Conversion failed when converting date"
1159        ));
1160    }
1161}