arrow_tiberius/
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    /// Starts a bulk writer on this same SQL Server connection.
140    ///
141    /// The returned writer borrows the connected client, so lifecycle SQL and
142    /// bulk loading cannot accidentally use two different connections through
143    /// this API.
144    pub async fn bulk_writer(
145        &mut self,
146        table: TableName,
147        planned_schema: PlannedSchema,
148        options: WriteOptions,
149    ) -> Result<ConnectedBulkWriter<'_>> {
150        let writer = BulkWriter::new(&mut self.client, table, planned_schema, options).await?;
151
152        Ok(ConnectedBulkWriter { writer })
153    }
154}
155
156impl ConnectedBulkWriter<'_> {
157    /// Writes one Arrow record batch.
158    pub async fn write_batch(&mut self, batch: &RecordBatch) -> Result<WriteStats> {
159        self.writer.write_batch(batch).await
160    }
161
162    /// Finalizes the bulk writer and returns cumulative write statistics.
163    pub async fn finish(self) -> Result<WriteStats> {
164        self.writer.finish().await
165    }
166}
167
168/// Connects to SQL Server from an ADO-style connection string.
169///
170/// The connection uses this crate's `tiberius-raw-bulk` dependency identity and
171/// Tokio TCP transport internally. The returned wrapper hides those concrete
172/// types from downstream crates.
173///
174/// The raw connection string is not stored in the returned client or in errors.
175pub async fn connect_mssql_client_from_ado_string(
176    connection_string: &str,
177) -> Result<ConnectedMssqlClient> {
178    let config = tiberius::Config::from_ado_string(connection_string)
179        .map_err(|_source| Error::InvalidConnectionString)?;
180    let tcp = TcpStream::connect(config.get_addr())
181        .await
182        .map_err(|source| Error::ConnectionTcpConnect { source })?;
183    tcp.set_nodelay(true)
184        .map_err(|source| Error::ConnectionTcpConnect { source })?;
185
186    let client = tiberius::Client::connect(config, tcp.compat_write())
187        .await
188        .map_err(|source| Error::ConnectionClientSetup { source })?;
189
190    Ok(ConnectedMssqlClient { client })
191}
192
193fn table_exists_query(table: &TableName) -> String {
194    let mut conditions = vec![format!(
195        "t.name = {}",
196        sql_string_literal(table.table().as_str())
197    )];
198    if let Some(schema) = table.schema() {
199        conditions.push(format!("s.name = {}", sql_string_literal(schema.as_str())));
200    }
201
202    format!(
203        "SELECT CASE WHEN EXISTS (SELECT 1 FROM sys.tables AS t \
204         INNER JOIN sys.schemas AS s ON s.schema_id = t.schema_id \
205         WHERE {}) THEN CAST(1 AS bit) ELSE CAST(0 AS bit) END AS [exists]",
206        conditions.join(" AND ")
207    )
208}
209
210fn target_row_count_query(table: &TableName) -> String {
211    format!(
212        "SELECT COUNT_BIG(*) AS [row_count] FROM {}",
213        table.quoted_sql()
214    )
215}
216
217fn count_big_i64_to_u64(count: i64) -> Result<u64> {
218    u64::try_from(count).map_err(|_| Error::TargetRowCountUnexpectedResult {
219        reason: "target row count was outside the supported range".to_owned(),
220    })
221}
222
223fn sql_string_literal(value: &str) -> String {
224    format!("N'{}'", value.replace('\'', "''"))
225}
226
227#[cfg(test)]
228mod tests {
229    use crate::{Error, connect_mssql_client_from_ado_string};
230
231    #[test]
232    fn sql_execution_outcome_records_rows_affected_in_order() {
233        let outcome = crate::SqlExecutionOutcome {
234            rows_affected: vec![2, 3, 5],
235        };
236
237        assert_eq!(outcome.rows_affected, vec![2, 3, 5]);
238        assert_eq!(outcome.total_rows_affected(), 10);
239    }
240
241    #[test]
242    fn table_exists_query_filters_schema_and_table() -> crate::Result<()> {
243        let table = crate::TableName::new("tenant", "people")?;
244        let query = super::table_exists_query(&table);
245
246        assert!(query.contains("FROM sys.tables AS t"));
247        assert!(query.contains("INNER JOIN sys.schemas AS s"));
248        assert!(query.contains("t.name = N'people'"));
249        assert!(query.contains("s.name = N'tenant'"));
250        Ok(())
251    }
252
253    #[test]
254    fn table_exists_query_escapes_string_literals() -> crate::Result<()> {
255        let table = crate::TableName::new("tenant's", "people's")?;
256        let query = super::table_exists_query(&table);
257
258        assert!(query.contains("t.name = N'people''s'"));
259        assert!(query.contains("s.name = N'tenant''s'"));
260        Ok(())
261    }
262
263    #[test]
264    fn unqualified_table_exists_query_filters_only_table_name() -> crate::Result<()> {
265        let table = crate::TableName::unqualified("people")?;
266        let query = super::table_exists_query(&table);
267
268        assert!(query.contains("t.name = N'people'"));
269        assert!(!query.contains("s.name ="));
270        Ok(())
271    }
272
273    #[test]
274    fn target_row_count_query_uses_quoted_table_name() -> crate::Result<()> {
275        let table = crate::TableName::new("tenant.schema", "people]2026")?;
276        let query = super::target_row_count_query(&table);
277
278        assert_eq!(
279            query,
280            "SELECT COUNT_BIG(*) AS [row_count] FROM [tenant.schema].[people]]2026]"
281        );
282        Ok(())
283    }
284
285    #[test]
286    fn count_big_conversion_rejects_negative_values_without_panicking() {
287        let error = super::count_big_i64_to_u64(-1).err().unwrap_or_else(|| {
288            Error::TargetRowCountUnexpectedResult {
289                reason: "expected negative count to fail".to_owned(),
290            }
291        });
292
293        assert!(matches!(
294            error,
295            Error::TargetRowCountUnexpectedResult { .. }
296        ));
297    }
298
299    #[test]
300    fn connected_client_type_is_public_without_raw_client_signature() {
301        let type_name = std::any::type_name::<crate::ConnectedMssqlClient>();
302
303        assert!(type_name.contains("ConnectedMssqlClient"));
304        assert!(!type_name.contains("tiberius::Client"));
305    }
306
307    #[test]
308    fn connected_writer_type_is_public_without_raw_transport_signature() {
309        let type_name = std::any::type_name::<crate::ConnectedBulkWriter<'static>>();
310
311        assert!(type_name.contains("ConnectedBulkWriter"));
312        assert!(!type_name.contains("tiberius::Client"));
313        assert!(!type_name.contains("tokio::net::TcpStream"));
314    }
315
316    #[tokio::test]
317    async fn invalid_connection_string_error_is_redacted() -> crate::Result<()> {
318        let connection_string =
319            "Server=tcp:localhost,notaport;Password=secret-token-123;Access Token=token-456";
320        let result = connect_mssql_client_from_ado_string(connection_string).await;
321        let Err(error) = result else {
322            return Err(Error::InvalidConnectionString);
323        };
324
325        assert!(matches!(error, Error::InvalidConnectionString));
326        let display = error.to_string();
327        let debug = format!("{error:?}");
328
329        for secret in ["secret-token-123", "token-456", connection_string] {
330            assert!(!display.contains(secret));
331            assert!(!debug.contains(secret));
332        }
333
334        Ok(())
335    }
336
337    #[tokio::test]
338    async fn tcp_connect_error_is_structured_and_redacted() -> crate::Result<()> {
339        let connection_string =
340            "Server=tcp:127.0.0.1,1;User Id=sa;Password=secret-token-123;Encrypt=false";
341        let result = connect_mssql_client_from_ado_string(connection_string).await;
342        let Err(error) = result else {
343            return Err(Error::InvalidConnectionString);
344        };
345
346        assert!(matches!(error, Error::ConnectionTcpConnect { .. }));
347        let display = error.to_string();
348        let debug = format!("{error:?}");
349
350        for secret in ["secret-token-123", connection_string] {
351            assert!(!display.contains(secret));
352            assert!(!debug.contains(secret));
353        }
354
355        Ok(())
356    }
357}
arrow_tiberius/connection.rs

arrow_tiberius/
connection.rs
