arrow_sql_server/
connection.rs

1//! SQL Server connection helpers.
2
3use std::fmt;
4
5use arrow_array::RecordBatch;
6use tokio::net::TcpStream;
7use tokio_util::compat::{Compat, TokioAsyncWriteCompatExt};
8
9use crate::{BulkWriter, Error, PlannedSchema, Result, TableName, WriteOptions, WriteStats};
10
11type CompatibleMssqlTransport = Compat<TcpStream>;
12
13/// Opaque SQL Server client constructed with this crate's compatible Tiberius dependency.
14///
15/// Use [`connect_mssql_client_from_ado_string`] to create this type. Its
16/// concrete Tiberius client and async transport types are intentionally hidden
17/// so downstream crates do not have to name or match `tiberius-raw-bulk`
18/// directly.
19pub struct ConnectedMssqlClient {
20    client: tiberius::Client<CompatibleMssqlTransport>,
21}
22
23/// Bulk writer created from a [`ConnectedMssqlClient`].
24///
25/// This wrapper keeps the compatible Tiberius client and transport types out of
26/// downstream signatures while exposing the same write and finish operations as
27/// [`BulkWriter`].
28pub struct ConnectedBulkWriter<'client> {
29    writer: BulkWriter<'client, CompatibleMssqlTransport>,
30}
31
32impl fmt::Debug for ConnectedMssqlClient {
33    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
34        formatter
35            .debug_struct("ConnectedMssqlClient")
36            .finish_non_exhaustive()
37    }
38}
39
40impl fmt::Debug for ConnectedBulkWriter<'_> {
41    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
42        formatter
43            .debug_struct("ConnectedBulkWriter")
44            .finish_non_exhaustive()
45    }
46}
47
48/// Metadata returned after executing SQL through a connected client.
49///
50/// This type is part of the narrow lifecycle SQL API. Statement execution is
51/// added separately from connection construction so connection setup can remain
52/// independently reviewable.
53#[derive(Clone, Debug, Default, Eq, PartialEq)]
54pub struct SqlExecutionOutcome {
55    /// Row counts reported by SQL Server DONE tokens, in server result order.
56    pub rows_affected: Vec<u64>,
57}
58
59impl SqlExecutionOutcome {
60    /// Returns the sum of all reported affected-row counts.
61    pub fn total_rows_affected(&self) -> u64 {
62        self.rows_affected.iter().copied().sum()
63    }
64}
65
66impl ConnectedMssqlClient {
67    /// Returns whether the target table exists in SQL Server metadata.
68    ///
69    /// This is a narrow metadata probe, not a generic query API. For
70    /// schema-qualified names it checks the exact schema and table. For
71    /// unqualified names it checks whether any table with that name exists in
72    /// the current database.
73    pub async fn table_exists(&mut self, table: &TableName) -> Result<bool> {
74        let query = table_exists_query(table);
75        let row = self
76            .client
77            .simple_query(query)
78            .await
79            .map_err(|source| Error::TableExistsQuery { source })?
80            .into_row()
81            .await
82            .map_err(|source| Error::TableExistsQuery { source })?
83            .ok_or_else(|| Error::TableExistsUnexpectedResult {
84                reason: "metadata query returned no rows".to_owned(),
85            })?;
86
87        row.try_get("exists")
88            .map_err(|source| Error::TableExistsQuery { source })?
89            .ok_or_else(|| Error::TableExistsUnexpectedResult {
90                reason: "metadata query returned NULL".to_owned(),
91            })
92    }
93
94    /// Returns `COUNT_BIG(*)` for a target table.
95    ///
96    /// The query uses this crate's bracket-quoted [`TableName`] rendering and
97    /// returns only a checked `u64` count. It does not expose raw SQL text,
98    /// result rows, or the underlying Tiberius client type.
99    pub async fn target_row_count(&mut self, table: &TableName) -> Result<u64> {
100        let query = target_row_count_query(table);
101        let row = self
102            .client
103            .simple_query(query)
104            .await
105            .map_err(|source| Error::TargetRowCountQuery { source })?
106            .into_row()
107            .await
108            .map_err(|source| Error::TargetRowCountQuery { source })?
109            .ok_or_else(|| Error::TargetRowCountUnexpectedResult {
110                reason: "target row count query returned no rows".to_owned(),
111            })?;
112        let count = row
113            .try_get::<i64, _>("row_count")
114            .map_err(|source| Error::TargetRowCountQuery { source })?
115            .ok_or_else(|| Error::TargetRowCountUnexpectedResult {
116                reason: "target row count query returned NULL".to_owned(),
117            })?;
118
119        count_big_i64_to_u64(count)
120    }
121
122    /// Executes a prepared lifecycle SQL statement.
123    ///
124    /// This method accepts statement text but intentionally returns only
125    /// affected-row metadata. It does not expose a generic result-row mapping
126    /// API.
127    pub async fn execute_statement(&mut self, sql: &str) -> Result<SqlExecutionOutcome> {
128        let result = self
129            .client
130            .execute(sql, &[])
131            .await
132            .map_err(|source| Error::SqlExecution { source })?;
133
134        Ok(SqlExecutionOutcome {
135            rows_affected: result.rows_affected().to_vec(),
136        })
137    }
138
139    /// Enables or disables SQL Server's persistent `table lock on bulk load`
140    /// option for a table.
141    ///
142    /// Enabling requires `ALTER` permission. Callers may ignore a known
143    /// nonfatal enable failure and continue without this optimization. Other
144    /// failures should be propagated.
145    ///
146    /// If enabling succeeds for a temporary load table, disabling must succeed
147    /// before the table is published. A disable failure should prevent
148    /// publication and trigger cleanup of the temporary table.
149    pub async fn set_bulk_load_table_lock(
150        &mut self,
151        table: &TableName,
152        enabled: bool,
153    ) -> Result<()> {
154        self.execute_statement(&bulk_load_table_lock_sql(table, enabled))
155            .await?;
156        Ok(())
157    }
158
159    /// Starts a bulk writer on this same SQL Server connection.
160    ///
161    /// The returned writer borrows the connected client, so lifecycle SQL and
162    /// bulk loading cannot accidentally use two different connections through
163    /// this API.
164    pub async fn bulk_writer(
165        &mut self,
166        table: TableName,
167        planned_schema: PlannedSchema,
168        options: WriteOptions,
169    ) -> Result<ConnectedBulkWriter<'_>> {
170        let writer = BulkWriter::new(&mut self.client, table, planned_schema, options).await?;
171
172        Ok(ConnectedBulkWriter { writer })
173    }
174}
175
176impl ConnectedBulkWriter<'_> {
177    /// Writes one Arrow record batch.
178    pub async fn write_batch(&mut self, batch: &RecordBatch) -> Result<WriteStats> {
179        self.writer.write_batch(batch).await
180    }
181
182    /// Finalizes the bulk writer and returns cumulative write statistics.
183    pub async fn finish(self) -> Result<WriteStats> {
184        self.writer.finish().await
185    }
186}
187
188/// Connects to SQL Server from an ADO-style connection string.
189///
190/// The connection uses this crate's `tiberius-raw-bulk` dependency identity and
191/// Tokio TCP transport internally. The returned wrapper hides those concrete
192/// types from downstream crates.
193///
194/// The raw connection string is not stored in the returned client or in errors.
195pub async fn connect_mssql_client_from_ado_string(
196    connection_string: &str,
197) -> Result<ConnectedMssqlClient> {
198    let config = tiberius::Config::from_ado_string(connection_string)
199        .map_err(|_source| Error::InvalidConnectionString)?;
200    let tcp = TcpStream::connect(config.get_addr())
201        .await
202        .map_err(|source| Error::ConnectionTcpConnect { source })?;
203    tcp.set_nodelay(true)
204        .map_err(|source| Error::ConnectionTcpConnect { source })?;
205
206    let client = tiberius::Client::connect(config, tcp.compat_write())
207        .await
208        .map_err(|source| Error::ConnectionClientSetup { source })?;
209
210    Ok(ConnectedMssqlClient { client })
211}
212
213fn table_exists_query(table: &TableName) -> String {
214    let mut conditions = vec![format!(
215        "t.name = {}",
216        sql_string_literal(table.table().as_str())
217    )];
218    if let Some(schema) = table.schema() {
219        conditions.push(format!("s.name = {}", sql_string_literal(schema.as_str())));
220    }
221
222    format!(
223        "SELECT CASE WHEN EXISTS (SELECT 1 FROM sys.tables AS t \
224         INNER JOIN sys.schemas AS s ON s.schema_id = t.schema_id \
225         WHERE {}) THEN CAST(1 AS bit) ELSE CAST(0 AS bit) END AS [exists]",
226        conditions.join(" AND ")
227    )
228}
229
230fn target_row_count_query(table: &TableName) -> String {
231    format!(
232        "SELECT COUNT_BIG(*) AS [row_count] FROM {}",
233        table.quoted_sql()
234    )
235}
236
237fn bulk_load_table_lock_sql(table: &TableName, enabled: bool) -> String {
238    let value = if enabled { "ON" } else { "OFF" };
239
240    format!(
241        "EXEC sys.sp_tableoption {}, 'table lock on bulk load', '{value}';",
242        sql_string_literal(&table.quoted_sql())
243    )
244}
245
246fn count_big_i64_to_u64(count: i64) -> Result<u64> {
247    u64::try_from(count).map_err(|_| Error::TargetRowCountUnexpectedResult {
248        reason: "target row count was outside the supported range".to_owned(),
249    })
250}
251
252fn sql_string_literal(value: &str) -> String {
253    format!("N'{}'", value.replace('\'', "''"))
254}
255
256#[cfg(test)]
257mod tests {
258    use crate::{Error, connect_mssql_client_from_ado_string};
259
260    #[test]
261    fn sql_execution_outcome_records_rows_affected_in_order() {
262        let outcome = crate::SqlExecutionOutcome {
263            rows_affected: vec![2, 3, 5],
264        };
265
266        assert_eq!(outcome.rows_affected, vec![2, 3, 5]);
267        assert_eq!(outcome.total_rows_affected(), 10);
268    }
269
270    #[test]
271    fn table_exists_query_filters_schema_and_table() -> crate::Result<()> {
272        let table = crate::TableName::new("tenant", "people")?;
273        let query = super::table_exists_query(&table);
274
275        assert!(query.contains("FROM sys.tables AS t"));
276        assert!(query.contains("INNER JOIN sys.schemas AS s"));
277        assert!(query.contains("t.name = N'people'"));
278        assert!(query.contains("s.name = N'tenant'"));
279        Ok(())
280    }
281
282    #[test]
283    fn table_exists_query_escapes_string_literals() -> crate::Result<()> {
284        let table = crate::TableName::new("tenant's", "people's")?;
285        let query = super::table_exists_query(&table);
286
287        assert!(query.contains("t.name = N'people''s'"));
288        assert!(query.contains("s.name = N'tenant''s'"));
289        Ok(())
290    }
291
292    #[test]
293    fn unqualified_table_exists_query_filters_only_table_name() -> crate::Result<()> {
294        let table = crate::TableName::unqualified("people")?;
295        let query = super::table_exists_query(&table);
296
297        assert!(query.contains("t.name = N'people'"));
298        assert!(!query.contains("s.name ="));
299        Ok(())
300    }
301
302    #[test]
303    fn target_row_count_query_uses_quoted_table_name() -> crate::Result<()> {
304        let table = crate::TableName::new("tenant.schema", "people]2026")?;
305        let query = super::target_row_count_query(&table);
306
307        assert_eq!(
308            query,
309            "SELECT COUNT_BIG(*) AS [row_count] FROM [tenant.schema].[people]]2026]"
310        );
311        Ok(())
312    }
313
314    #[test]
315    fn bulk_load_table_lock_sql_uses_quoted_table_name_and_requested_state() -> crate::Result<()> {
316        let table = crate::TableName::new("tenant's", "people's")?;
317
318        assert_eq!(
319            super::bulk_load_table_lock_sql(&table, true),
320            "EXEC sys.sp_tableoption N'[tenant''s].[people''s]', 'table lock on bulk load', 'ON';"
321        );
322        assert_eq!(
323            super::bulk_load_table_lock_sql(&table, false),
324            "EXEC sys.sp_tableoption N'[tenant''s].[people''s]', 'table lock on bulk load', 'OFF';"
325        );
326        Ok(())
327    }
328
329    #[test]
330    fn count_big_conversion_rejects_negative_values_without_panicking() {
331        let error = super::count_big_i64_to_u64(-1).err().unwrap_or_else(|| {
332            Error::TargetRowCountUnexpectedResult {
333                reason: "expected negative count to fail".to_owned(),
334            }
335        });
336
337        assert!(matches!(
338            error,
339            Error::TargetRowCountUnexpectedResult { .. }
340        ));
341    }
342
343    #[test]
344    fn connected_client_type_is_public_without_raw_client_signature() {
345        let type_name = std::any::type_name::<crate::ConnectedMssqlClient>();
346
347        assert!(type_name.contains("ConnectedMssqlClient"));
348        assert!(!type_name.contains("tiberius::Client"));
349    }
350
351    #[test]
352    fn connected_writer_type_is_public_without_raw_transport_signature() {
353        let type_name = std::any::type_name::<crate::ConnectedBulkWriter<'static>>();
354
355        assert!(type_name.contains("ConnectedBulkWriter"));
356        assert!(!type_name.contains("tiberius::Client"));
357        assert!(!type_name.contains("tokio::net::TcpStream"));
358    }
359
360    #[tokio::test]
361    async fn invalid_connection_string_error_is_redacted() -> crate::Result<()> {
362        let connection_string =
363            "Server=tcp:localhost,notaport;Password=secret-token-123;Access Token=token-456";
364        let result = connect_mssql_client_from_ado_string(connection_string).await;
365        let Err(error) = result else {
366            return Err(Error::InvalidConnectionString);
367        };
368
369        assert!(matches!(error, Error::InvalidConnectionString));
370        let display = error.to_string();
371        let debug = format!("{error:?}");
372
373        for secret in ["secret-token-123", "token-456", connection_string] {
374            assert!(!display.contains(secret));
375            assert!(!debug.contains(secret));
376        }
377
378        Ok(())
379    }
380
381    #[tokio::test]
382    async fn tcp_connect_error_is_structured_and_redacted() -> crate::Result<()> {
383        let connection_string =
384            "Server=tcp:127.0.0.1,1;User Id=sa;Password=secret-token-123;Encrypt=false";
385        let result = connect_mssql_client_from_ado_string(connection_string).await;
386        let Err(error) = result else {
387            return Err(Error::InvalidConnectionString);
388        };
389
390        assert!(matches!(error, Error::ConnectionTcpConnect { .. }));
391        let display = error.to_string();
392        let debug = format!("{error:?}");
393
394        for secret in ["secret-token-123", connection_string] {
395            assert!(!display.contains(secret));
396            assert!(!debug.contains(secret));
397        }
398
399        Ok(())
400    }
401}
arrow_sql_server/connection.rs

arrow_sql_server/
connection.rs
