mssql_client/client.rs
1//! SQL Server client implementation.
2//!
3//! ## DDL and statement routing
4//!
5//! [`Client::execute`] routes automatically by parameter count: with no
6//! parameters it sends a SQL batch (which permits DDL such as `CREATE` / `ALTER`
7//! / `DROP`); with parameters it uses `sp_executesql`, whose procedure context
8//! SQL Server forbids DDL in. Run DDL with an empty parameter slice:
9//!
10//! ```rust,no_run
11//! # async fn create_table(config: mssql_client::Config) -> Result<(), mssql_client::Error> {
12//! # let mut client = mssql_client::Client::connect(config).await?;
13//! client.execute("CREATE TABLE dbo.t (id INT)", &[]).await?;
14//! # Ok(())
15//! # }
16//! ```
17//!
18//! Use [`Client::simple_query`] for fire-and-forget batches (including
19//! multi-statement, `;`-separated DDL) when you don't need the affected-row count.
20
21// Allow unwrap/expect for chrono date construction with known-valid constant dates
22// and for regex patterns that are compile-time constants
23#![allow(clippy::unwrap_used, clippy::expect_used, clippy::needless_range_loop)]
24
25mod connect;
26mod params;
27pub(crate) mod response;
28
29use std::marker::PhantomData;
30
31use mssql_codec::connection::Connection;
32#[cfg(feature = "tls")]
33use mssql_tls::TlsStream;
34use tds_protocol::packet::PacketType;
35use tds_protocol::rpc::RpcRequest;
36use tds_protocol::token::{EnvChange, EnvChangeType};
37use tokio::net::TcpStream;
38
39use crate::config::Config;
40use crate::error::{Error, Result};
41#[cfg(feature = "otel")]
42use crate::instrumentation::InstrumentationContext;
43use crate::state::{ConnectionState, InTransaction, Ready};
44use crate::statement_cache::StatementCache;
45use crate::stream::{MultiResultStream, QueryStream};
46use crate::transaction::SavePoint;
47
48/// How long to wait for the server to acknowledge an Attention packet after
49/// a command timeout. SqlClient waits 5 seconds before dooming the
50/// connection; we match it.
51const ATTENTION_ACK_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
52
53/// Run a network future under an optional command deadline.
54///
55/// On timeout this sends an Attention packet via `canceller` and then awaits
56/// the future so its own read loop drains the server's DONE_ATTN
57/// acknowledgement, leaving the connection clean before returning
58/// [`Error::CommandTimeout`]. This is the cancel-safe alternative to dropping
59/// the future (e.g. via `tokio::time::timeout`), which would leave unconsumed
60/// TDS data in the connection buffer and desync the next request.
61///
62/// The drain itself is bounded by [`ATTENTION_ACK_TIMEOUT`] — a hung server
63/// that never acknowledges the attention must not turn the timeout into an
64/// infinite wait. When the bound expires the connection is abandoned
65/// mid-response: `in_flight` stays set, so the pool discards the connection
66/// at check-in instead of reusing it.
67pub(crate) async fn run_with_deadline<F, T>(
68 fut: F,
69 deadline: Option<std::time::Duration>,
70 canceller: crate::cancel::CancelHandle,
71) -> Result<T>
72where
73 F: std::future::Future<Output = Result<T>>,
74{
75 let Some(d) = deadline else {
76 return fut.await;
77 };
78 tokio::pin!(fut);
79 tokio::select! {
80 biased;
81 res = &mut fut => res,
82 () = tokio::time::sleep(d) => {
83 // Signal cancellation, then let the in-flight read consume the
84 // server's attention acknowledgement so the connection stays usable.
85 let drain = async {
86 let _ = canceller.cancel().await;
87 let _ = (&mut fut).await;
88 };
89 if tokio::time::timeout(ATTENTION_ACK_TIMEOUT, drain).await.is_err() {
90 tracing::warn!(
91 timeout = ?ATTENTION_ACK_TIMEOUT,
92 "server did not acknowledge attention; abandoning the connection as dirty"
93 );
94 }
95 Err(Error::CommandTimeout)
96 }
97 }
98}
99
100/// SQL Server client with type-state connection management.
101///
102/// The generic parameter `S` represents the current connection state,
103/// ensuring at compile time that certain operations are only available
104/// in appropriate states.
105pub struct Client<S: ConnectionState> {
106 config: Config,
107 _state: PhantomData<S>,
108 /// The underlying connection (present only when connected)
109 connection: Option<ConnectionHandle>,
110 /// Server version from LoginAck (raw u32 TDS version)
111 server_version: Option<u32>,
112 /// Current database from EnvChange
113 current_database: Option<String>,
114 /// Server's default collation from SqlCollation EnvChange during login.
115 /// Used when `SendStringParametersAsUnicode=false` to encode VARCHAR
116 /// parameters with the correct character encoding and collation bytes.
117 server_collation: Option<tds_protocol::token::Collation>,
118 /// Prepared statement cache for query optimization
119 statement_cache: StatementCache,
120 /// Transaction descriptor from BeginTransaction EnvChange.
121 /// Per MS-TDS spec, this value must be included in ALL_HEADERS for subsequent
122 /// requests within an explicit transaction. 0 indicates auto-commit mode.
123 transaction_descriptor: u64,
124 /// Whether a request has been sent and the response has not yet been fully read.
125 /// Used by the connection pool to detect dirty connections after cancel/timeout.
126 in_flight: bool,
127 /// Whether this connection needs a reset on next use.
128 /// Set by connection pool on checkin, cleared after first query/execute.
129 /// When true, the RESETCONNECTION flag is set on the first TDS packet.
130 needs_reset: bool,
131 /// OpenTelemetry instrumentation context (when otel feature is enabled)
132 #[cfg(feature = "otel")]
133 instrumentation: InstrumentationContext,
134 /// Always Encrypted context for column decryption (when always-encrypted feature is enabled)
135 #[cfg(feature = "always-encrypted")]
136 pub(crate) encryption_context: Option<std::sync::Arc<crate::encryption::EncryptionContext>>,
137}
138
139/// Internal connection handle wrapping the actual connection.
140///
141/// This is an enum to support different connection types:
142/// - TLS (TDS 8.0 strict mode) - requires `tls` feature
143/// - TLS with PreLogin wrapping (TDS 7.x style) - requires `tls` feature
144/// - Plain TCP (for internal networks or when `tls` feature is disabled)
145enum ConnectionHandle {
146 /// TLS connection (TDS 8.0 strict mode - TLS before any TDS traffic)
147 #[cfg(feature = "tls")]
148 Tls(Connection<TlsStream<TcpStream>>),
149 /// TLS connection with PreLogin wrapping (TDS 7.x style)
150 #[cfg(feature = "tls")]
151 TlsPrelogin(Connection<TlsStream<mssql_tls::TlsPreloginWrapper<TcpStream>>>),
152 /// Plain TCP connection (for internal networks or when `tls` feature is disabled)
153 Plain(Connection<TcpStream>),
154}
155
156/// The parameter `TypeInfo` to declare a typed NULL ([`crate::null`]) with, from
157/// its [`crate::ToSql::sql_type`] name. Returns `None` for an untyped NULL
158/// (`Option::None`, type `"NULL"`), which falls back to the default param type.
159#[cfg(feature = "always-encrypted")]
160fn null_param_type_info(sql_type: &str) -> Option<tds_protocol::rpc::TypeInfo> {
161 use tds_protocol::rpc::TypeInfo;
162 Some(match sql_type {
163 "BIT" => TypeInfo::bit(),
164 "TINYINT" => TypeInfo::tinyint(),
165 "SMALLINT" => TypeInfo::smallint(),
166 "INT" => TypeInfo::int(),
167 "BIGINT" => TypeInfo::bigint(),
168 "REAL" => TypeInfo::real(),
169 "FLOAT" => TypeInfo::float(),
170 "NVARCHAR" => TypeInfo::nvarchar(1),
171 "VARBINARY" => TypeInfo::varbinary(1),
172 "UNIQUEIDENTIFIER" => TypeInfo::uuid(),
173 "DATE" => TypeInfo::date(),
174 _ => return None,
175 })
176}
177
178/// Map a typed-parameter wrapper's [`EncryptedParamType`] to the `TypeInfo` the
179/// driver declares it as (for `sp_describe_parameter_encryption` and the
180/// `CryptoMetadata` base type). Unknown future variants error rather than
181/// silently declaring the wrong type.
182#[cfg(feature = "always-encrypted")]
183fn encrypted_param_type_info(
184 ty: mssql_types::EncryptedParamType,
185) -> Result<tds_protocol::rpc::TypeInfo> {
186 use mssql_types::EncryptedParamType as E;
187 use tds_protocol::rpc::TypeInfo;
188 Ok(match ty {
189 E::Decimal { precision, scale } => TypeInfo::decimal(precision, scale),
190 E::Time { scale } => TypeInfo::time(scale),
191 E::DateTime2 { scale } => TypeInfo::datetime2(scale),
192 E::DateTimeOffset { scale } => TypeInfo::datetimeoffset(scale),
193 E::DateTime => TypeInfo::datetime(),
194 E::Char { length } => TypeInfo::char(length),
195 E::NChar { length } => TypeInfo::nchar(length),
196 E::Binary { length } => TypeInfo::binary(length),
197 _ => {
198 return Err(Error::Encryption(
199 "unsupported Always Encrypted parameter type".to_string(),
200 ));
201 }
202 })
203}
204
205// Private helper methods available to all connection states
206impl<S: ConnectionState> Client<S> {
207 /// The default per-command deadline from `command_timeout`.
208 ///
209 /// Returns `None` when `command_timeout` is zero, which means "no limit"
210 /// (matching ADO.NET's `SqlCommand.CommandTimeout = 0`).
211 pub(crate) fn command_deadline(&self) -> Option<std::time::Duration> {
212 let t = self.config.command_timeout;
213 if t.is_zero() { None } else { Some(t) }
214 }
215
216 /// Build a cancel handle for the current connection, regardless of
217 /// connection state. The public, documented surface is
218 /// [`Client::<Ready>::cancel_handle`]; both state-specific methods
219 /// delegate here.
220 pub(crate) fn connection_cancel_handle(&self) -> crate::cancel::CancelHandle {
221 let connection = self
222 .connection
223 .as_ref()
224 .expect("connection should be present");
225 match connection {
226 #[cfg(feature = "tls")]
227 ConnectionHandle::Tls(conn) => {
228 crate::cancel::CancelHandle::from_tls(conn.cancel_handle())
229 }
230 #[cfg(feature = "tls")]
231 ConnectionHandle::TlsPrelogin(conn) => {
232 crate::cancel::CancelHandle::from_tls_prelogin(conn.cancel_handle())
233 }
234 ConnectionHandle::Plain(conn) => {
235 crate::cancel::CancelHandle::from_plain(conn.cancel_handle())
236 }
237 }
238 }
239
240 /// Cancel an in-flight response that was abandoned without being drained —
241 /// e.g. a [`RowStream`](crate::RowStream) dropped or cancelled mid-result.
242 ///
243 /// Sends an Attention and drains to the server's DONE_ATTN acknowledgement so
244 /// the socket is clean and the connection reusable. A no-op when nothing is
245 /// in flight. Bounded by [`ATTENTION_ACK_TIMEOUT`]: if the acknowledgement
246 /// never arrives the connection is left marked in-flight (so the pool
247 /// discards it on return) and an error is returned.
248 pub(crate) async fn cancel_in_flight_response(&mut self) -> Result<()> {
249 if !self.in_flight {
250 return Ok(());
251 }
252 let canceller = self.connection_cancel_handle();
253 let drain = async {
254 canceller.cancel().await?;
255 // With the cancelling flag set, `read_response_message` routes through
256 // the codec's drain-after-cancel path and returns `Err(Cancelled)`
257 // once the DONE_ATTN acknowledgement is consumed (clearing
258 // `in_flight`). Any full messages that arrive before the ack are
259 // discarded.
260 loop {
261 match self.read_response_message().await {
262 Err(Error::Cancelled) => return Ok(()),
263 Ok(_) => continue,
264 Err(e) => return Err(e),
265 }
266 }
267 };
268 match tokio::time::timeout(ATTENTION_ACK_TIMEOUT, drain).await {
269 Ok(result) => result,
270 Err(_) => {
271 tracing::warn!(
272 timeout = ?ATTENTION_ACK_TIMEOUT,
273 "attention acknowledgement not received while cancelling an \
274 abandoned response; connection left dirty"
275 );
276 Err(Error::Cancelled)
277 }
278 }
279 }
280
281 /// Process transaction-related EnvChange tokens.
282 ///
283 /// This handles BeginTransaction, CommitTransaction, and RollbackTransaction
284 /// EnvChange tokens, updating the transaction descriptor accordingly.
285 ///
286 /// This enables executing BEGIN TRANSACTION, COMMIT, and ROLLBACK via raw SQL
287 /// while still having the transaction descriptor tracked correctly.
288 fn process_transaction_env_change(env: &EnvChange, transaction_descriptor: &mut u64) {
289 use tds_protocol::token::EnvChangeValue;
290
291 match env.env_type {
292 EnvChangeType::BeginTransaction => {
293 if let EnvChangeValue::Binary(ref data) = env.new_value {
294 if data.len() >= 8 {
295 let descriptor = u64::from_le_bytes([
296 data[0], data[1], data[2], data[3], data[4], data[5], data[6], data[7],
297 ]);
298 tracing::debug!(descriptor = descriptor, "transaction started via raw SQL");
299 *transaction_descriptor = descriptor;
300 }
301 }
302 }
303 EnvChangeType::CommitTransaction | EnvChangeType::RollbackTransaction => {
304 tracing::debug!(
305 env_type = ?env.env_type,
306 "transaction ended via raw SQL"
307 );
308 *transaction_descriptor = 0;
309 }
310 _ => {}
311 }
312 }
313
314 /// Apply a transaction-related `ENVCHANGE` to this client's descriptor.
315 ///
316 /// Lets the streaming readers (which live in sibling modules) keep the
317 /// transaction descriptor in sync with raw `BEGIN`/`COMMIT`/`ROLLBACK`
318 /// batches seen mid-stream, exactly as the buffered readers do.
319 pub(crate) fn apply_transaction_env_change(&mut self, env: &EnvChange) {
320 Self::process_transaction_env_change(env, &mut self.transaction_descriptor);
321 }
322
323 /// Send a SQL batch to the server.
324 ///
325 /// Uses the client's current transaction descriptor in ALL_HEADERS.
326 /// Per MS-TDS spec, when in an explicit transaction, the descriptor
327 /// returned by BeginTransaction must be included.
328 ///
329 /// If `needs_reset` is set (from pool return), the RESETCONNECTION flag
330 /// is included in the first packet to reset connection state.
331 async fn send_sql_batch(&mut self, sql: &str) -> Result<()> {
332 // If a previous streamed response was abandoned (a RowStream dropped
333 // mid-result), drain it before issuing a new request so the next read
334 // does not pick up the old response's bytes.
335 self.cancel_in_flight_response().await?;
336
337 let payload = tds_protocol::__private::encode_sql_batch_with_transaction(
338 sql,
339 self.transaction_descriptor,
340 );
341 let max_packet = self.config.packet_size as usize;
342
343 // Check if we need to reset the connection on this request
344 let reset = self.needs_reset;
345 if reset {
346 self.needs_reset = false; // Clear flag before sending
347 // RESETCONNECTION invalidates all server-side prepared handles, so
348 // drop the cache (no sp_unprepare needed — the server released them).
349 let _ = self.statement_cache.clear();
350 tracing::debug!("sending SQL batch with RESETCONNECTION flag");
351 }
352
353 self.in_flight = true;
354 let connection = self.connection.as_mut().ok_or(Error::ConnectionClosed)?;
355
356 match connection {
357 #[cfg(feature = "tls")]
358 ConnectionHandle::Tls(conn) => {
359 conn.send_message_with_reset(PacketType::SqlBatch, payload, max_packet, reset)
360 .await?;
361 }
362 #[cfg(feature = "tls")]
363 ConnectionHandle::TlsPrelogin(conn) => {
364 conn.send_message_with_reset(PacketType::SqlBatch, payload, max_packet, reset)
365 .await?;
366 }
367 ConnectionHandle::Plain(conn) => {
368 conn.send_message_with_reset(PacketType::SqlBatch, payload, max_packet, reset)
369 .await?;
370 }
371 }
372
373 Ok(())
374 }
375
376 /// Send an RPC request to the server.
377 ///
378 /// Uses the client's current transaction descriptor in ALL_HEADERS.
379 ///
380 /// If `needs_reset` is set (from pool return), the RESETCONNECTION flag
381 /// is included in the first packet to reset connection state.
382 pub(crate) async fn send_rpc(&mut self, rpc: &RpcRequest) -> Result<()> {
383 // Drain an abandoned streamed response (see `send_sql_batch`) before
384 // issuing this request.
385 self.cancel_in_flight_response().await?;
386
387 let payload = rpc.encode_with_transaction(self.transaction_descriptor);
388 let max_packet = self.config.packet_size as usize;
389
390 // Check if we need to reset the connection on this request
391 let reset = self.needs_reset;
392 if reset {
393 self.needs_reset = false; // Clear flag before sending
394 // RESETCONNECTION invalidates all server-side prepared handles, so
395 // drop the cache (no sp_unprepare needed — the server released them).
396 let _ = self.statement_cache.clear();
397 tracing::debug!("sending RPC with RESETCONNECTION flag");
398 }
399
400 self.in_flight = true;
401 let connection = self.connection.as_mut().ok_or(Error::ConnectionClosed)?;
402
403 match connection {
404 #[cfg(feature = "tls")]
405 ConnectionHandle::Tls(conn) => {
406 conn.send_message_with_reset(PacketType::Rpc, payload, max_packet, reset)
407 .await?;
408 }
409 #[cfg(feature = "tls")]
410 ConnectionHandle::TlsPrelogin(conn) => {
411 conn.send_message_with_reset(PacketType::Rpc, payload, max_packet, reset)
412 .await?;
413 }
414 ConnectionHandle::Plain(conn) => {
415 conn.send_message_with_reset(PacketType::Rpc, payload, max_packet, reset)
416 .await?;
417 }
418 }
419
420 Ok(())
421 }
422
423 /// Start building a stored procedure call with full control over parameters.
424 ///
425 /// Returns a [`crate::procedure::ProcedureBuilder`] that allows adding named input and output
426 /// parameters before executing the call.
427 ///
428 /// The procedure name is validated to prevent SQL injection. It may be
429 /// schema-qualified (e.g., `"dbo.MyProc"`).
430 ///
431 /// # Example
432 ///
433 /// ```rust,no_run
434 /// # async fn ex(client: &mut mssql_client::Client<mssql_client::Ready>) -> Result<(), mssql_client::Error> {
435 /// let result = client.procedure("dbo.CalculateSum")?
436 /// .input("@a", &10i32)
437 /// .input("@b", &20i32)
438 /// .output_int("@result")
439 /// .execute().await?;
440 ///
441 /// let sum = result.get_output("@result").unwrap();
442 /// # let _ = sum;
443 /// # Ok(())
444 /// # }
445 /// ```
446 pub fn procedure(
447 &mut self,
448 proc_name: &str,
449 ) -> Result<crate::procedure::ProcedureBuilder<'_, S>> {
450 crate::validation::validate_qualified_identifier(proc_name)?;
451 Ok(crate::procedure::ProcedureBuilder::new(self, proc_name))
452 }
453
454 /// Execute a stored procedure with positional input parameters.
455 ///
456 /// This is a convenience method for the common case of calling a procedure
457 /// with input-only parameters. For output parameters or named parameters,
458 /// use [`procedure()`](Client::procedure) instead.
459 ///
460 /// # Example
461 ///
462 /// ```rust,no_run
463 /// # async fn ex(client: &mut mssql_client::Client<mssql_client::Ready>) -> Result<(), mssql_client::Error> {
464 /// let result = client.call_procedure("dbo.GetUser", &[&1i32]).await?;
465 /// assert_eq!(result.return_value, 0);
466 ///
467 /// if let Some(rs) = result.first_result_set() {
468 /// println!("columns: {:?}", rs.columns());
469 /// }
470 /// # Ok(())
471 /// # }
472 /// ```
473 pub async fn call_procedure(
474 &mut self,
475 proc_name: &str,
476 params: &[&(dyn crate::ToSql + Sync)],
477 ) -> Result<crate::stream::ProcedureResult> {
478 crate::validation::validate_qualified_identifier(proc_name)?;
479
480 tracing::debug!(
481 proc_name = proc_name,
482 params_count = params.len(),
483 "executing stored procedure"
484 );
485
486 let rpc_params =
487 Self::convert_params_positional(params, self.send_unicode(), self.server_collation())?;
488 let mut rpc = RpcRequest::named(proc_name);
489 for param in rpc_params {
490 rpc = rpc.param(param);
491 }
492
493 #[cfg(feature = "otel")]
494 let instrumentation = self.instrumentation.clone();
495 #[cfg(feature = "otel")]
496 let mut span = instrumentation.procedure_span(proc_name);
497 #[cfg(feature = "otel")]
498 let timer = crate::instrumentation::OperationTimer::start("EXECUTE");
499
500 let deadline = self.command_deadline();
501 let canceller = self.connection_cancel_handle();
502 let result = run_with_deadline(
503 async {
504 self.send_rpc(&rpc).await?;
505 self.read_procedure_result().await
506 },
507 deadline,
508 canceller,
509 )
510 .await;
511
512 #[cfg(feature = "otel")]
513 match &result {
514 Ok(r) => InstrumentationContext::record_success(&mut span, Some(r.rows_affected)),
515 Err(e) => InstrumentationContext::record_error(&mut span, e),
516 }
517 #[cfg(feature = "otel")]
518 timer.finish(instrumentation.metrics(), result.is_ok());
519 #[cfg(feature = "otel")]
520 drop(span);
521
522 result
523 }
524
525 /// Ask the server how each parameter of a statement must be encrypted.
526 ///
527 /// Issues the `sp_describe_parameter_encryption` system RPC for the
528 /// parameterized statement `tsql` with the parameter declaration `params`
529 /// (e.g. `"@id int, @name nvarchar(64)"`), and parses the two result sets
530 /// into a [`ParameterEncryptionInfo`](crate::encryption::ParameterEncryptionInfo): the
531 /// CEK table, plus — for each parameter the server reports as encrypted —
532 /// which CEK and whether deterministic or randomized. Parameters the server
533 /// reports as plaintext are omitted.
534 ///
535 /// This is the first step of Always Encrypted parameter encryption; the
536 /// connection must have negotiated it (`Column Encryption Setting=Enabled`).
537 #[cfg(feature = "always-encrypted")]
538 pub(crate) async fn describe_parameter_encryption(
539 &mut self,
540 tsql: &str,
541 params: &str,
542 ) -> Result<crate::encryption::ParameterEncryptionInfo> {
543 let tsql_arg = tsql.to_string();
544 let params_arg = params.to_string();
545 let mut result = self
546 .call_procedure(
547 "sp_describe_parameter_encryption",
548 &[&tsql_arg, ¶ms_arg],
549 )
550 .await?;
551 crate::encryption::ParameterEncryptionInfo::from_describe_result_sets(
552 &mut result.result_sets,
553 )
554 }
555
556 /// Build the `sp_executesql` request for a parameterized statement.
557 ///
558 /// When the connection has Always Encrypted enabled, parameters the server
559 /// reports as encrypted are encrypted client-side first (an extra
560 /// `sp_describe_parameter_encryption` round-trip). Otherwise this is the
561 /// plain parameter conversion.
562 pub(crate) async fn build_parameterized_rpc(
563 &mut self,
564 sql: &str,
565 params: &[&(dyn crate::ToSql + Sync)],
566 ) -> Result<RpcRequest> {
567 #[cfg(feature = "always-encrypted")]
568 if self.encryption_context.is_some() {
569 return self.build_encrypted_sql_rpc(sql, params).await;
570 }
571 let rpc_params =
572 Self::convert_params(params, self.send_unicode(), self.server_collation())?;
573 Ok(RpcRequest::execute_sql(sql, rpc_params))
574 }
575
576 /// Send a parameterized `query` request, consulting the prepared-statement
577 /// cache when [`Config::statement_cache`](crate::Config::statement_cache)
578 /// is enabled. Leaves the execution response ready for the caller's
579 /// `read_query_response`.
580 ///
581 /// Falls back to the default path (SQL batch for no params, `sp_executesql`
582 /// otherwise) when the cache is disabled, the query has no parameters, or
583 /// Always Encrypted is active (prepared + AE parameter encryption is out of
584 /// scope for this first increment).
585 async fn send_query_request(
586 &mut self,
587 sql: &str,
588 params: &[&(dyn crate::ToSql + Sync)],
589 ) -> Result<()> {
590 #[cfg(feature = "always-encrypted")]
591 let ae_active = self.encryption_context.is_some();
592 #[cfg(not(feature = "always-encrypted"))]
593 let ae_active = false;
594
595 if !self.config.statement_cache || params.is_empty() || ae_active {
596 if params.is_empty() {
597 self.send_sql_batch(sql).await?;
598 } else {
599 let rpc = self.build_parameterized_rpc(sql, params).await?;
600 self.send_rpc(&rpc).await?;
601 }
602 return Ok(());
603 }
604
605 // If a connection reset is pending, the next packet carries
606 // RESETCONNECTION (set in `send_rpc`), which invalidates every
607 // server-side prepared handle. Drop the cache BEFORE the lookup so this
608 // request re-prepares instead of `sp_execute`-ing a handle the reset is
609 // about to invalidate (otherwise the server rejects it with "Could not
610 // find prepared statement with handle N").
611 if self.needs_reset {
612 let _ = self.statement_cache.clear();
613 }
614
615 let rpc_params =
616 Self::convert_params(params, self.send_unicode(), self.server_collation())?;
617 // Key on the parameter declaration + SQL: a cached handle is only valid
618 // for the exact prepared parameter types, so two calls with the same
619 // SQL but different param types must not share a handle.
620 let key = format!(
621 "{}\u{1}{sql}",
622 RpcRequest::build_param_declarations(&rpc_params)
623 );
624
625 if let Some(handle) = self.statement_cache.get(&key) {
626 // Hit: sp_execute the cached handle. Clear any stale pending key
627 // (e.g. from a prior request whose read aborted) so the read path
628 // does not try to capture a handle from this response.
629 self.statement_cache.set_pending(None);
630 let rpc = RpcRequest::execute(handle, rpc_params);
631 self.send_rpc(&rpc).await?;
632 } else {
633 // Miss: sp_prepexec prepares and executes in ONE round-trip. The
634 // caller's read_query_response reads the row response and captures
635 // the `@handle` RETURNVALUE, then stores it under `key` (see
636 // `store_pending_prepared_handle`).
637 self.statement_cache.set_pending(Some(key));
638 let rpc = RpcRequest::prepexec(sql, rpc_params);
639 self.send_rpc(&rpc).await?;
640 }
641 Ok(())
642 }
643
644 /// Store the handle captured from an `sp_prepexec` execution response under
645 /// the pending cache key (set by [`send_query_request`](Self::send_query_request)
646 /// on a cold miss), releasing any LRU-evicted server-side handle.
647 ///
648 /// Called by `read_query_response` after it has read the row response and
649 /// the trailing `@handle` RETURNVALUE. A no-op when no prepexec is pending.
650 /// Eviction `sp_unprepare` is best-effort: a failure leaks one handle until
651 /// connection reset, never corrupts data.
652 pub(super) async fn store_pending_prepared_handle(
653 &mut self,
654 handle: Option<i32>,
655 ) -> Result<()> {
656 let Some(key) = self.statement_cache.take_pending() else {
657 return Ok(());
658 };
659 let Some(handle) = handle else {
660 // No @handle came back (unexpected for sp_prepexec): leave the
661 // statement uncached so the next call simply re-prepares.
662 return Ok(());
663 };
664 if let Some(evicted) = self
665 .statement_cache
666 .insert(crate::statement_cache::PreparedStatement::new(handle, key))
667 {
668 let unprepare = RpcRequest::unprepare(evicted.handle());
669 self.send_rpc(&unprepare).await?;
670 let _ = self.read_procedure_result().await?;
671 }
672 Ok(())
673 }
674
675 /// Encrypt the Always Encrypted parameters of a statement, then build its
676 /// `sp_executesql` request.
677 ///
678 /// Asks the server which parameters are encrypted
679 /// ([`describe_parameter_encryption`](Self::describe_parameter_encryption)),
680 /// then for each one normalizes the value, resolves its column encryption
681 /// key, encrypts, and emits an encrypted RPC parameter. Parameters the
682 /// server reports as plaintext are sent unchanged.
683 #[cfg(feature = "always-encrypted")]
684 async fn build_encrypted_sql_rpc(
685 &mut self,
686 sql: &str,
687 params: &[&(dyn crate::ToSql + Sync)],
688 ) -> Result<RpcRequest> {
689 use tds_protocol::rpc::RpcParam;
690
691 let Some(ctx) = self.encryption_context.clone() else {
692 let rpc_params =
693 Self::convert_params(params, self.send_unicode(), self.server_collation())?;
694 return Ok(RpcRequest::execute_sql(sql, rpc_params));
695 };
696
697 // Resolve each parameter's value once (AE normalization needs the typed
698 // value, not the wire encoding) and build the plaintext RPC params.
699 let send_unicode = self.send_unicode();
700 let collation = self.server_collation().cloned();
701 let mut values: Vec<mssql_types::SqlValue> = Vec::with_capacity(params.len());
702 let mut plaintext: Vec<RpcParam> = Vec::with_capacity(params.len());
703 let mut hints: Vec<Option<mssql_types::EncryptedParamType>> =
704 Vec::with_capacity(params.len());
705 for (i, p) in params.iter().enumerate() {
706 let name = format!("@p{}", i + 1);
707 let value = p.to_sql()?;
708 let hint = p.encrypted_param_type();
709 // A typed NULL (e.g. `null::<i32>()`) is declared by its SQL type so
710 // describe accepts it against the target encrypted column; an untyped
711 // NULL falls back to the default in `sql_value_to_rpc_param`.
712 let rpc_param = match (&value, null_param_type_info(p.sql_type())) {
713 (mssql_types::SqlValue::Null, Some(type_info)) => RpcParam::null(&name, type_info),
714 _ => {
715 let mut param = Self::sql_value_to_rpc_param(
716 &name,
717 &value,
718 send_unicode,
719 collation.as_ref(),
720 )?;
721 // A typed-parameter wrapper (e.g. `numeric(v, p, s)`,
722 // `datetime2(v, scale)`) declares an explicit SQL type so
723 // describe matches the encrypted column exactly — the value
724 // alone cannot convey precision/scale or the legacy-`datetime`
725 // vs `datetime2` distinction.
726 if let Some(ty) = hint {
727 param.type_info = encrypted_param_type_info(ty)?;
728 }
729 param
730 }
731 };
732 plaintext.push(rpc_param);
733 values.push(value);
734 hints.push(hint);
735 }
736
737 if plaintext.is_empty() {
738 return Ok(RpcRequest::execute_sql(sql, plaintext));
739 }
740
741 // Ask the server which parameters need encryption.
742 let declarations = RpcRequest::build_param_declarations(&plaintext);
743 let info = self
744 .describe_parameter_encryption(sql, &declarations)
745 .await?;
746 if info.parameters.is_empty() {
747 return Ok(RpcRequest::execute_sql(sql, plaintext));
748 }
749
750 // Encrypt the flagged parameters; pass the rest through untouched.
751 let mut final_params: Vec<RpcParam> = Vec::with_capacity(plaintext.len());
752 for ((value, param), hint) in values.into_iter().zip(plaintext).zip(hints) {
753 let Some(crypto) = info.get_parameter(¶m.name) else {
754 final_params.push(param);
755 continue;
756 };
757 let entry = info.cek_table.get(crypto.cek_ordinal).ok_or_else(|| {
758 Error::Protocol(format!(
759 "encrypted parameter {} references missing CEK ordinal {}",
760 param.name, crypto.cek_ordinal
761 ))
762 })?;
763 let metadata = tds_protocol::rpc::EncryptedParamMetadata {
764 base_type_info: param.type_info.clone(),
765 algorithm_id: crypto.algorithm_id,
766 encryption_type: crypto.encryption_type,
767 database_id: entry.database_id,
768 cek_id: entry.cek_id,
769 cek_version: entry.cek_version,
770 cek_md_version: entry.cek_md_version,
771 normalization_rule_version: crypto.normalization_rule_version,
772 };
773 // A NULL value bound to an encrypted column is sent as an encrypted
774 // NULL (the server rejects a plaintext parameter for an encrypted
775 // column); there is nothing to encrypt.
776 if matches!(value, mssql_types::SqlValue::Null) {
777 final_params.push(RpcParam::encrypted_null(param.name, metadata));
778 continue;
779 }
780 let normalized = crate::encryption::normalize_for_encryption(&value, hint)?;
781 let ciphertext = ctx
782 .encrypt_value(&normalized, entry, crypto.encryption_type)
783 .await?;
784 final_params.push(RpcParam::encrypted(
785 param.name,
786 bytes::Bytes::from(ciphertext),
787 metadata,
788 ));
789 }
790
791 Ok(RpcRequest::execute_sql(sql, final_params))
792 }
793
794 /// Start a bulk insert operation for the specified table.
795 ///
796 /// Sends the `INSERT BULK` statement to the server and returns a
797 /// [`crate::bulk::BulkWriter`] for streaming rows. The writer holds
798 /// a mutable borrow on the client, preventing other operations while
799 /// the bulk insert is in progress.
800 ///
801 /// # Example
802 ///
803 /// ```rust,no_run
804 /// # async fn ex(client: &mut mssql_client::Client<mssql_client::Ready>) -> Result<(), mssql_client::Error> {
805 /// use mssql_client::{BulkInsertBuilder, BulkColumn, SqlValue};
806 ///
807 /// let builder = BulkInsertBuilder::new("dbo.Users")
808 /// .with_typed_columns(vec![
809 /// BulkColumn::new("id", "INT", 0)?,
810 /// BulkColumn::new("name", "NVARCHAR(100)", 1)?,
811 /// ]);
812 ///
813 /// let mut writer = client.bulk_insert(&builder).await?;
814 /// writer.send_row_values(&[SqlValue::Int(1), SqlValue::String("Alice".into())])?;
815 /// writer.send_row_values(&[SqlValue::Int(2), SqlValue::String("Bob".into())])?;
816 /// let result = writer.finish().await?;
817 /// println!("Inserted {} rows", result.rows_affected);
818 /// # Ok(())
819 /// # }
820 /// ```
821 pub async fn bulk_insert(
822 &mut self,
823 builder: &crate::bulk::BulkInsertBuilder,
824 ) -> Result<crate::bulk::BulkWriter<'_, S>> {
825 use tds_protocol::token::{ColMetaData, Token};
826
827 tracing::debug!(
828 table = builder.table_name(),
829 columns = builder.columns().len(),
830 "starting bulk insert"
831 );
832
833 // Step 1: Query the server for column metadata.
834 // This gives us the exact type encoding the server expects for BulkLoad,
835 // following the pattern established by Tiberius.
836 let meta_query = format!("SELECT TOP 0 * FROM {}", builder.table_name());
837 let deadline = self.command_deadline();
838 let canceller = self.connection_cancel_handle();
839 let message = run_with_deadline(
840 async {
841 self.send_sql_batch(&meta_query).await?;
842 self.read_response_message().await
843 },
844 deadline,
845 canceller,
846 )
847 .await?;
848 self.in_flight = false;
849
850 // Capture both the raw COLMETADATA bytes and parsed column info
851 let raw_payload = message.payload.clone();
852 let mut parser = self.create_parser(message.payload);
853 let mut server_metadata: Option<ColMetaData> = None;
854 let mut meta_start: usize = 0;
855 let mut meta_end: usize = 0;
856
857 loop {
858 let pos_before = raw_payload.len() - parser.remaining();
859 let token = parser.next_token_with_metadata(server_metadata.as_ref())?;
860 let pos_after = raw_payload.len() - parser.remaining();
861 let Some(token) = token else { break };
862
863 match token {
864 Token::ColMetaData(meta) => {
865 meta_start = pos_before;
866 meta_end = pos_after;
867 server_metadata = Some(meta);
868 }
869 Token::Done(_) => break,
870 _ => {}
871 }
872 }
873
874 // Reject deprecated TEXT/NTEXT/IMAGE columns reported by the server.
875 // These types require a legacy TEXTPTR wire format that this driver
876 // does not support — users should migrate the column to VARCHAR(MAX) /
877 // NVARCHAR(MAX) / VARBINARY(MAX).
878 if let Some(ref meta) = server_metadata {
879 use tds_protocol::types::TypeId;
880 for col in meta.columns.iter() {
881 let (rejected, replacement) = match col.type_id {
882 TypeId::Text => (Some("TEXT"), "VARCHAR(MAX)"),
883 TypeId::NText => (Some("NTEXT"), "NVARCHAR(MAX)"),
884 TypeId::Image => (Some("IMAGE"), "VARBINARY(MAX)"),
885 _ => (None, ""),
886 };
887 if let Some(sql_type) = rejected {
888 return Err(Error::from(mssql_types::TypeError::UnsupportedType {
889 sql_type: sql_type.to_string(),
890 reason: format!(
891 "column `{}` in table `{}` is {} — TEXT/NTEXT/IMAGE \
892 are not supported. Alter the column to {} instead \
893 (Microsoft deprecated TEXT/NTEXT/IMAGE in SQL \
894 Server 2005).",
895 col.name,
896 builder.table_name(),
897 sql_type,
898 replacement,
899 ),
900 }));
901 }
902 }
903 }
904
905 // Step 2: Send INSERT BULK statement to put server in bulk load mode
906 let stmt = builder.build_insert_bulk_statement()?;
907 let deadline = self.command_deadline();
908 let canceller = self.connection_cancel_handle();
909 run_with_deadline(
910 async {
911 self.send_sql_batch(&stmt).await?;
912 self.read_execute_result().await
913 },
914 deadline,
915 canceller,
916 )
917 .await?;
918
919 // Step 3: Create bulk writer with server's metadata
920 let raw_meta = if meta_end > meta_start {
921 Some(raw_payload.slice(meta_start..meta_end))
922 } else {
923 None
924 };
925
926 let server_cols = server_metadata.as_ref().map(|m| m.columns.as_slice());
927 let bulk = crate::bulk::BulkInsert::new_with_server_metadata(
928 builder.columns().to_vec(),
929 builder.options().batch_size,
930 raw_meta,
931 server_cols,
932 );
933
934 Ok(crate::bulk::BulkWriter::new(self, bulk))
935 }
936
937 /// Start a bulk insert without querying the server for column metadata.
938 ///
939 /// Unlike [`bulk_insert()`](Self::bulk_insert), this method does not send
940 /// `SELECT TOP 0 * FROM table` to discover column types. Instead, the
941 /// column metadata is constructed from the `BulkColumn` types provided
942 /// on the builder. This saves a round-trip when the schema is known.
943 ///
944 /// # Caveats
945 ///
946 /// The caller must ensure `BulkColumn` entries match the target table's
947 /// column definitions exactly. Mismatched types, lengths, precision/scale,
948 /// or column ordering will cause the server to reject the BulkLoad packet.
949 ///
950 /// For most use cases, prefer [`bulk_insert()`](Self::bulk_insert) — the
951 /// extra round-trip is usually negligible and the server-supplied metadata
952 /// is guaranteed correct.
953 pub async fn bulk_insert_without_schema_discovery(
954 &mut self,
955 builder: &crate::bulk::BulkInsertBuilder,
956 ) -> Result<crate::bulk::BulkWriter<'_, S>> {
957 tracing::debug!(
958 table = builder.table_name(),
959 columns = builder.columns().len(),
960 "starting bulk insert (no schema discovery)"
961 );
962
963 // Send INSERT BULK statement to put server in bulk load mode
964 let stmt = builder.build_insert_bulk_statement()?;
965 let deadline = self.command_deadline();
966 let canceller = self.connection_cancel_handle();
967 run_with_deadline(
968 async {
969 self.send_sql_batch(&stmt).await?;
970 self.read_execute_result().await
971 },
972 deadline,
973 canceller,
974 )
975 .await?;
976
977 // Create bulk writer with hand-crafted metadata
978 let bulk =
979 crate::bulk::BulkInsert::new(builder.columns().to_vec(), builder.options().batch_size);
980
981 Ok(crate::bulk::BulkWriter::new(self, bulk))
982 }
983
984 /// Send bulk load data as a BulkLoad (0x07) message and read the server response.
985 ///
986 /// Used internally by [`crate::bulk::BulkWriter::finish()`] to transmit accumulated
987 /// row data after the `INSERT BULK` statement has been acknowledged.
988 pub(crate) async fn send_and_read_bulk_load(&mut self, payload: bytes::Bytes) -> Result<u64> {
989 let max_packet = self.config.packet_size as usize;
990
991 self.in_flight = true;
992 let connection = self.connection.as_mut().ok_or(Error::ConnectionClosed)?;
993
994 match connection {
995 #[cfg(feature = "tls")]
996 ConnectionHandle::Tls(conn) => {
997 conn.send_message(PacketType::BulkLoad, payload, max_packet)
998 .await?;
999 }
1000 #[cfg(feature = "tls")]
1001 ConnectionHandle::TlsPrelogin(conn) => {
1002 conn.send_message(PacketType::BulkLoad, payload, max_packet)
1003 .await?;
1004 }
1005 ConnectionHandle::Plain(conn) => {
1006 conn.send_message(PacketType::BulkLoad, payload, max_packet)
1007 .await?;
1008 }
1009 }
1010
1011 // Read the server's Done response with row count
1012 self.read_execute_result().await
1013 }
1014
1015 /// Execute a query with named parameters and return a streaming result set.
1016 ///
1017 /// This method accepts [`NamedParam`](crate::to_params::NamedParam) values,
1018 /// making it compatible with the [`ToParams`](crate::to_params::ToParams) trait
1019 /// and the `#[derive(ToParams)]` macro.
1020 ///
1021 /// # Example
1022 ///
1023 /// ```rust,no_run
1024 /// # async fn ex(client: &mut mssql_client::Client<mssql_client::Ready>) -> Result<(), mssql_client::Error> {
1025 /// use mssql_client::{NamedParam, ToParams};
1026 ///
1027 /// // With derive macro:
1028 /// #[derive(mssql_derive::ToParams)]
1029 /// struct UserQuery { name: String }
1030 ///
1031 /// let q = UserQuery { name: "Alice".into() };
1032 /// let rows = client.query_named(
1033 /// "SELECT * FROM users WHERE name = @name",
1034 /// &q.to_params()?,
1035 /// ).await?;
1036 ///
1037 /// // Or manually:
1038 /// let params = vec![NamedParam::from_value("name", &"Alice")?];
1039 /// let rows = client.query_named(
1040 /// "SELECT * FROM users WHERE name = @name",
1041 /// ¶ms,
1042 /// ).await?;
1043 /// # let _ = rows;
1044 /// # Ok(())
1045 /// # }
1046 /// ```
1047 pub async fn query_named<'a>(
1048 &'a mut self,
1049 sql: &str,
1050 params: &[crate::to_params::NamedParam],
1051 ) -> Result<QueryStream<'a>> {
1052 tracing::debug!(
1053 sql = sql,
1054 params_count = params.len(),
1055 "executing query with named parameters"
1056 );
1057
1058 #[cfg(feature = "otel")]
1059 let instrumentation = self.instrumentation.clone();
1060 #[cfg(feature = "otel")]
1061 let mut span = instrumentation.query_span(sql);
1062 #[cfg(feature = "otel")]
1063 let timer = crate::instrumentation::OperationTimer::start(
1064 crate::instrumentation::extract_operation(sql),
1065 );
1066
1067 let result = async {
1068 if params.is_empty() {
1069 self.send_sql_batch(sql).await?;
1070 } else {
1071 let rpc_params = Self::convert_named_params(
1072 params,
1073 self.send_unicode(),
1074 self.server_collation(),
1075 )?;
1076 let rpc = RpcRequest::execute_sql(sql, rpc_params);
1077 self.send_rpc(&rpc).await?;
1078 }
1079
1080 self.read_query_response().await
1081 }
1082 .await;
1083
1084 #[cfg(feature = "otel")]
1085 match &result {
1086 Ok(_) => InstrumentationContext::record_success(&mut span, None),
1087 Err(e) => InstrumentationContext::record_error(&mut span, e),
1088 }
1089 #[cfg(feature = "otel")]
1090 timer.finish(instrumentation.metrics(), result.is_ok());
1091 #[cfg(feature = "otel")]
1092 drop(span);
1093
1094 let resp = result?;
1095 #[cfg(feature = "always-encrypted")]
1096 {
1097 Ok(QueryStream::from_raw(
1098 resp.columns,
1099 resp.pending_rows,
1100 resp.meta,
1101 resp.decryptor,
1102 ))
1103 }
1104 #[cfg(not(feature = "always-encrypted"))]
1105 {
1106 Ok(QueryStream::from_raw(
1107 resp.columns,
1108 resp.pending_rows,
1109 resp.meta,
1110 ))
1111 }
1112 }
1113
1114 /// Execute a statement with named parameters.
1115 ///
1116 /// Returns the number of affected rows. This is the named-parameter
1117 /// counterpart of [`execute()`](Client::execute), compatible with the
1118 /// [`ToParams`](crate::to_params::ToParams) trait.
1119 ///
1120 /// # Example
1121 ///
1122 /// ```rust,no_run
1123 /// # async fn ex(client: &mut mssql_client::Client<mssql_client::Ready>) -> Result<(), mssql_client::Error> {
1124 /// use mssql_client::NamedParam;
1125 ///
1126 /// let params = vec![
1127 /// NamedParam::from_value("name", &"Alice")?,
1128 /// NamedParam::from_value("email", &"alice@example.com")?,
1129 /// ];
1130 /// let rows_affected = client.execute_named(
1131 /// "INSERT INTO users (name, email) VALUES (@name, @email)",
1132 /// ¶ms,
1133 /// ).await?;
1134 /// # let _ = rows_affected;
1135 /// # Ok(())
1136 /// # }
1137 /// ```
1138 pub async fn execute_named(
1139 &mut self,
1140 sql: &str,
1141 params: &[crate::to_params::NamedParam],
1142 ) -> Result<u64> {
1143 tracing::debug!(
1144 sql = sql,
1145 params_count = params.len(),
1146 "executing statement with named parameters"
1147 );
1148
1149 #[cfg(feature = "otel")]
1150 let instrumentation = self.instrumentation.clone();
1151 #[cfg(feature = "otel")]
1152 let mut span = instrumentation.query_span(sql);
1153 #[cfg(feature = "otel")]
1154 let timer = crate::instrumentation::OperationTimer::start(
1155 crate::instrumentation::extract_operation(sql),
1156 );
1157
1158 let deadline = self.command_deadline();
1159 let canceller = self.connection_cancel_handle();
1160 let result = run_with_deadline(
1161 async {
1162 if params.is_empty() {
1163 self.send_sql_batch(sql).await?;
1164 } else {
1165 let rpc_params = Self::convert_named_params(
1166 params,
1167 self.send_unicode(),
1168 self.server_collation(),
1169 )?;
1170 let rpc = RpcRequest::execute_sql(sql, rpc_params);
1171 self.send_rpc(&rpc).await?;
1172 }
1173
1174 self.read_execute_result().await
1175 },
1176 deadline,
1177 canceller,
1178 )
1179 .await;
1180
1181 #[cfg(feature = "otel")]
1182 match &result {
1183 Ok(rows) => InstrumentationContext::record_success(&mut span, Some(*rows)),
1184 Err(e) => InstrumentationContext::record_error(&mut span, e),
1185 }
1186 #[cfg(feature = "otel")]
1187 timer.finish(instrumentation.metrics(), result.is_ok());
1188 #[cfg(feature = "otel")]
1189 drop(span);
1190
1191 result
1192 }
1193
1194 /// The connection's OpenTelemetry instrumentation context.
1195 #[cfg(feature = "otel")]
1196 pub(crate) fn instrumentation(&self) -> &InstrumentationContext {
1197 &self.instrumentation
1198 }
1199
1200 /// Snapshot this connection's prepared-statement cache statistics.
1201 ///
1202 /// Reflects activity since the connection was established (or its last
1203 /// reset). Meaningful only when
1204 /// [`Config::statement_cache`](crate::Config::statement_cache) is enabled;
1205 /// otherwise the cache is never consulted and all counts stay zero.
1206 #[must_use]
1207 pub fn statement_cache_stats(&self) -> crate::StatementCacheStats {
1208 self.statement_cache.stats()
1209 }
1210
1211 /// Whether string parameters are sent as NVARCHAR (Unicode).
1212 pub(crate) fn send_unicode(&self) -> bool {
1213 self.config.send_string_parameters_as_unicode
1214 }
1215
1216 /// Server's default collation, captured from ENVCHANGE during login.
1217 pub(crate) fn server_collation(&self) -> Option<&tds_protocol::token::Collation> {
1218 self.server_collation.as_ref()
1219 }
1220
1221 /// Shared implementation behind `query_stream` for both `Ready` and
1222 /// `InTransaction`. Sends the request, then pulls packets until the first
1223 /// result set's `ColMetaData` (resolving columns and any Always Encrypted
1224 /// decryptor up front) before handing back a [`RowStream`].
1225 pub(crate) async fn query_stream_inner<'a>(
1226 &'a mut self,
1227 sql: &str,
1228 params: &[&(dyn crate::ToSql + Sync)],
1229 ) -> Result<crate::row_stream::RowStream<'a, S>> {
1230 use crate::client::response::server_token_to_error;
1231 use crate::row_source::{Pull, RowSource};
1232 use tds_protocol::token::Token;
1233
1234 tracing::debug!(sql = sql, params_count = params.len(), "streaming query");
1235
1236 // Send the request (same wire format as the buffered path).
1237 if params.is_empty() {
1238 self.send_sql_batch(sql).await?;
1239 } else {
1240 let rpc = self.build_parameterized_rpc(sql, params).await?;
1241 self.send_rpc(&rpc).await?;
1242 }
1243 self.in_flight = true;
1244
1245 #[cfg(feature = "always-encrypted")]
1246 let encryption_enabled = self.encryption_context.is_some();
1247 #[cfg(not(feature = "always-encrypted"))]
1248 let encryption_enabled = false;
1249
1250 let mut source = RowSource::new(encryption_enabled);
1251
1252 // Prelude: pull packets until the first result set's ColMetaData (so the
1253 // columns and any Always Encrypted decryptor are resolved up front), or
1254 // until a terminal Done/Error if there is no result set.
1255 loop {
1256 match source.pull()? {
1257 Pull::Token(Token::ColMetaData(meta)) => {
1258 let columns = Self::build_columns(&meta);
1259 #[cfg(feature = "always-encrypted")]
1260 let decryptor = self
1261 .resolve_decryptor(&meta)
1262 .await?
1263 .map(std::sync::Arc::new);
1264 return Ok(crate::row_stream::RowStream::new(
1265 self,
1266 source,
1267 columns,
1268 meta,
1269 #[cfg(feature = "always-encrypted")]
1270 decryptor,
1271 ));
1272 }
1273 Pull::Token(Token::Error(err)) => {
1274 self.in_flight = false;
1275 return Err(server_token_to_error(&err));
1276 }
1277 Pull::Token(Token::Done(done)) => {
1278 if done.status.error {
1279 self.in_flight = false;
1280 return Err(Error::Query(
1281 "query failed (server set error flag in DONE token)".to_string(),
1282 ));
1283 }
1284 if !done.status.more {
1285 // No result set (e.g. an INSERT) — an empty stream.
1286 self.in_flight = false;
1287 return Ok(crate::row_stream::RowStream::empty(self));
1288 }
1289 // More results may follow; keep looking for ColMetaData.
1290 }
1291 Pull::Token(Token::EnvChange(env)) => {
1292 Self::process_transaction_env_change(&env, &mut self.transaction_descriptor);
1293 }
1294 Pull::Token(_) => {
1295 // Info / Order / DoneProc / DoneInProc, etc. — keep pulling.
1296 }
1297 Pull::NeedMore => match self.read_response_packet().await? {
1298 Some((payload, is_eom)) => source.push_packet(payload, is_eom),
1299 None => {
1300 self.in_flight = false;
1301 return Err(Error::ConnectionClosed);
1302 }
1303 },
1304 Pull::End => {
1305 self.in_flight = false;
1306 return Ok(crate::row_stream::RowStream::empty(self));
1307 }
1308 }
1309 }
1310 }
1311
1312 /// Shared implementation behind `query_stream_blob` for both `Ready` and
1313 /// `InTransaction`.
1314 pub(crate) async fn query_stream_blob_inner<'a>(
1315 &'a mut self,
1316 sql: &str,
1317 params: &[&(dyn crate::ToSql + Sync)],
1318 ) -> Result<crate::blob_stream::BlobStream<'a, S>> {
1319 let (meta, buf, eom, encryption_enabled) = self.open_blob_stream(sql, params).await?;
1320 let first_blob = Self::validate_blob_result_set(&meta)?;
1321 Ok(crate::blob_stream::BlobStream::new(
1322 self,
1323 buf,
1324 eom,
1325 encryption_enabled,
1326 meta,
1327 first_blob,
1328 // Single trailing MAX column; auto-position it so the existing
1329 // `next` → `copy_blob_to` flow works without an explicit `next_blob`.
1330 1,
1331 true,
1332 ))
1333 }
1334
1335 /// Shared implementation behind `query_stream_rows` for both `Ready` and
1336 /// `InTransaction`.
1337 pub(crate) async fn query_stream_rows_inner<'a>(
1338 &'a mut self,
1339 sql: &str,
1340 params: &[&(dyn crate::ToSql + Sync)],
1341 ) -> Result<crate::blob_stream::BlobStream<'a, S>> {
1342 let (meta, buf, eom, encryption_enabled) = self.open_blob_stream(sql, params).await?;
1343 let (first_blob, blob_count) = Self::validate_blob_rows_result_set(&meta)?;
1344 Ok(crate::blob_stream::BlobStream::new(
1345 self,
1346 buf,
1347 eom,
1348 encryption_enabled,
1349 meta,
1350 first_blob,
1351 blob_count,
1352 // Caller drives blobs explicitly via `next_blob`.
1353 false,
1354 ))
1355 }
1356
1357 /// Send the query and pull tokens until the first `ColMetaData`, returning
1358 /// the result-set metadata plus the unconsumed post-metadata wire bytes.
1359 /// Shared by the single-blob and multi-blob streaming paths.
1360 async fn open_blob_stream(
1361 &mut self,
1362 sql: &str,
1363 params: &[&(dyn crate::ToSql + Sync)],
1364 ) -> Result<(tds_protocol::token::ColMetaData, bytes::Bytes, bool, bool)> {
1365 use crate::client::response::server_token_to_error;
1366 use crate::row_source::{Pull, RowSource};
1367 use tds_protocol::token::Token;
1368
1369 if params.is_empty() {
1370 self.send_sql_batch(sql).await?;
1371 } else {
1372 let rpc = self.build_parameterized_rpc(sql, params).await?;
1373 self.send_rpc(&rpc).await?;
1374 }
1375 self.in_flight = true;
1376
1377 #[cfg(feature = "always-encrypted")]
1378 let encryption_enabled = self.encryption_context.is_some();
1379 #[cfg(not(feature = "always-encrypted"))]
1380 let encryption_enabled = false;
1381
1382 let mut source = RowSource::new(encryption_enabled);
1383
1384 loop {
1385 match source.pull()? {
1386 Pull::Token(Token::ColMetaData(meta)) => {
1387 let (buf, eom) = source.into_parts();
1388 return Ok((meta, buf, eom, encryption_enabled));
1389 }
1390 Pull::Token(Token::Error(err)) => {
1391 self.in_flight = false;
1392 return Err(server_token_to_error(&err));
1393 }
1394 Pull::Token(Token::Done(_)) => {
1395 self.in_flight = false;
1396 return Err(Error::Protocol(
1397 "blob streaming: query produced no result set".to_string(),
1398 ));
1399 }
1400 Pull::Token(_) => {}
1401 Pull::NeedMore => match self.read_response_packet().await? {
1402 Some((payload, is_eom)) => source.push_packet(payload, is_eom),
1403 None => {
1404 self.in_flight = false;
1405 return Err(Error::ConnectionClosed);
1406 }
1407 },
1408 Pull::End => {
1409 self.in_flight = false;
1410 return Err(Error::Protocol(
1411 "blob streaming: query produced no result set".to_string(),
1412 ));
1413 }
1414 }
1415 }
1416 }
1417
1418 /// Validate that a result set is shaped for [`query_stream_blob`] and return
1419 /// the index of its single trailing MAX column.
1420 fn validate_blob_result_set(meta: &tds_protocol::token::ColMetaData) -> Result<usize> {
1421 if meta.cek_table.is_some() {
1422 return Err(Error::Protocol(
1423 "query_stream_blob does not support Always Encrypted result sets".to_string(),
1424 ));
1425 }
1426 let max_cols: Vec<usize> = meta
1427 .columns
1428 .iter()
1429 .enumerate()
1430 .filter(|(_, c)| crate::blob_stream::is_plp_max(c))
1431 .map(|(i, _)| i)
1432 .collect();
1433 match max_cols.as_slice() {
1434 [] => Err(Error::Protocol(
1435 "query_stream_blob: result set has no MAX column — use query_stream".to_string(),
1436 )),
1437 [idx] if *idx == meta.columns.len() - 1 => Ok(*idx),
1438 [_] => Err(Error::Protocol(
1439 "query_stream_blob: the MAX column must be the last column".to_string(),
1440 )),
1441 _ => Err(Error::Protocol(
1442 "query_stream_blob: result set has more than one MAX column".to_string(),
1443 )),
1444 }
1445 }
1446
1447 /// Validate that a result set is shaped for [`query_stream_rows`] and return
1448 /// `(first_blob_index, blob_count)` — the start and length of the trailing
1449 /// run of MAX columns.
1450 ///
1451 /// Requires at least one MAX column and that every MAX column be trailing
1452 /// (no scalar column may follow a MAX column). The interleaved case (a
1453 /// scalar column after a MAX column) is rejected — supporting it needs a
1454 /// resumable per-column decoder (tracked in #258).
1455 fn validate_blob_rows_result_set(
1456 meta: &tds_protocol::token::ColMetaData,
1457 ) -> Result<(usize, usize)> {
1458 if meta.cek_table.is_some() {
1459 return Err(Error::Protocol(
1460 "query_stream_rows does not support Always Encrypted result sets".to_string(),
1461 ));
1462 }
1463 let first_blob = meta
1464 .columns
1465 .iter()
1466 .position(crate::blob_stream::is_plp_max)
1467 .ok_or_else(|| {
1468 Error::Protocol(
1469 "query_stream_rows: result set has no MAX column — use query_stream"
1470 .to_string(),
1471 )
1472 })?;
1473 // Every column from the first MAX column onward must itself be a MAX
1474 // column; a scalar column after a blob cannot be decoded until the blob
1475 // is consumed.
1476 if !meta.columns[first_blob..]
1477 .iter()
1478 .all(crate::blob_stream::is_plp_max)
1479 {
1480 return Err(Error::Protocol(
1481 "query_stream_rows: a non-MAX column follows a MAX column; interleaved MAX \
1482 columns are not supported (the MAX columns must be trailing)"
1483 .to_string(),
1484 ));
1485 }
1486 Ok((first_blob, meta.columns.len() - first_blob))
1487 }
1488}
1489
1490impl Client<Ready> {
1491 /// Mark this connection as needing a reset on next use.
1492 ///
1493 /// Called by the connection pool when a connection is returned.
1494 /// The next SQL batch or RPC will include the RESETCONNECTION flag
1495 /// in the TDS packet header, causing SQL Server to reset connection
1496 /// state (temp tables, SET options, transaction isolation level, etc.)
1497 /// before executing the command.
1498 ///
1499 /// This is more efficient than calling `sp_reset_connection` as a
1500 /// separate command because it's handled at the TDS protocol level.
1501 pub fn mark_needs_reset(&mut self) {
1502 self.needs_reset = true;
1503 }
1504
1505 /// Check if this connection needs a reset.
1506 ///
1507 /// Returns true if `mark_needs_reset()` was called and the reset
1508 /// hasn't been performed yet.
1509 #[must_use]
1510 pub fn needs_reset(&self) -> bool {
1511 self.needs_reset
1512 }
1513
1514 /// Execute a query and return a result set with lazy per-row decoding.
1515 ///
1516 /// Per ADR-007 the full response is buffered in memory and each row is
1517 /// *decoded* on demand as you iterate — this is not incremental network
1518 /// streaming, so peak memory tracks the response size. Use
1519 /// `.collect_all()` if you want all rows materialized into a `Vec` up
1520 /// front.
1521 ///
1522 /// # Example
1523 ///
1524 /// ```rust,no_run
1525 /// # use mssql_client::Row;
1526 /// # fn process(_: &Row) {}
1527 /// # async fn ex(client: &mut mssql_client::Client<mssql_client::Ready>) -> Result<(), mssql_client::Error> {
1528 /// // Streaming (synchronous iteration over the result set)
1529 /// let stream = client.query("SELECT * FROM users WHERE id = @p1", &[&1]).await?;
1530 /// for row in stream {
1531 /// let row = row?;
1532 /// process(&row);
1533 /// }
1534 ///
1535 /// // Buffered (loads all into memory)
1536 /// let rows: Vec<Row> = client
1537 /// .query("SELECT * FROM small_table", &[])
1538 /// .await?
1539 /// .collect_all()
1540 /// .await?;
1541 /// # let _ = rows;
1542 /// # Ok(())
1543 /// # }
1544 /// ```
1545 pub async fn query<'a>(
1546 &'a mut self,
1547 sql: &str,
1548 params: &[&(dyn crate::ToSql + Sync)],
1549 ) -> Result<QueryStream<'a>> {
1550 let deadline = self.command_deadline();
1551 self.query_inner(sql, params, deadline).await
1552 }
1553
1554 /// Shared query implementation with an explicit command deadline.
1555 async fn query_inner<'a>(
1556 &'a mut self,
1557 sql: &str,
1558 params: &[&(dyn crate::ToSql + Sync)],
1559 deadline: Option<std::time::Duration>,
1560 ) -> Result<QueryStream<'a>> {
1561 tracing::debug!(sql = sql, params_count = params.len(), "executing query");
1562
1563 #[cfg(feature = "otel")]
1564 let instrumentation = self.instrumentation.clone();
1565 #[cfg(feature = "otel")]
1566 let mut span = instrumentation.query_span(sql);
1567 #[cfg(feature = "otel")]
1568 let timer = crate::instrumentation::OperationTimer::start(
1569 crate::instrumentation::extract_operation(sql),
1570 );
1571
1572 let canceller = self.cancel_handle();
1573 let result = run_with_deadline(
1574 async {
1575 // Sends via the prepared-statement cache when enabled, else the
1576 // SQL batch / sp_executesql default.
1577 self.send_query_request(sql, params).await?;
1578
1579 // Read complete response including columns and rows
1580 self.read_query_response().await
1581 },
1582 deadline,
1583 canceller,
1584 )
1585 .await;
1586
1587 #[cfg(feature = "otel")]
1588 match &result {
1589 Ok(_) => InstrumentationContext::record_success(&mut span, None),
1590 Err(e) => InstrumentationContext::record_error(&mut span, e),
1591 }
1592 #[cfg(feature = "otel")]
1593 timer.finish(instrumentation.metrics(), result.is_ok());
1594
1595 // Drop the span before returning
1596 #[cfg(feature = "otel")]
1597 drop(span);
1598
1599 let resp = result?;
1600 #[cfg(feature = "always-encrypted")]
1601 {
1602 Ok(QueryStream::from_raw(
1603 resp.columns,
1604 resp.pending_rows,
1605 resp.meta,
1606 resp.decryptor,
1607 ))
1608 }
1609 #[cfg(not(feature = "always-encrypted"))]
1610 {
1611 Ok(QueryStream::from_raw(
1612 resp.columns,
1613 resp.pending_rows,
1614 resp.meta,
1615 ))
1616 }
1617 }
1618
1619 /// Execute a query and stream rows incrementally from the network.
1620 ///
1621 /// Unlike [`query`](Self::query) — which buffers the whole response in
1622 /// memory before returning — this reads TDS packets on demand as rows are
1623 /// pulled, so peak memory is roughly one packet plus one row regardless of
1624 /// result-set size. Use it for large result sets; use [`query`](Self::query)
1625 /// for the common small-result case where the buffered, synchronously
1626 /// iterable [`QueryStream`] is more convenient.
1627 ///
1628 /// The returned [`RowStream`](crate::RowStream) borrows the client for its
1629 /// lifetime, so no other request can run on this connection until the stream
1630 /// is consumed or dropped. Also available on `Client<InTransaction>` to
1631 /// stream within a transaction.
1632 ///
1633 /// # Example
1634 ///
1635 /// ```rust,no_run
1636 /// # async fn ex(client: &mut mssql_client::Client<mssql_client::Ready>) -> Result<(), mssql_client::Error> {
1637 /// let mut stream = client.query_stream("SELECT id FROM big_table", &[]).await?;
1638 /// while let Some(row) = stream.try_next().await? {
1639 /// let id: i32 = row.get_by_name("id")?;
1640 /// let _ = id;
1641 /// }
1642 /// # Ok(())
1643 /// # }
1644 /// ```
1645 pub async fn query_stream<'a>(
1646 &'a mut self,
1647 sql: &str,
1648 params: &[&(dyn crate::ToSql + Sync)],
1649 ) -> Result<crate::row_stream::RowStream<'a, Ready>> {
1650 self.query_stream_inner(sql, params).await
1651 }
1652
1653 /// Execute a query and stream a row's trailing MAX column from the network.
1654 ///
1655 /// For result sets whose last column is a single MAX type
1656 /// (`VARBINARY(MAX)`, `NVARCHAR(MAX)`, `VARCHAR(MAX)`, `XML`), this reads
1657 /// that column's bytes incrementally from the socket instead of
1658 /// materializing the cell — so a multi-GB BLOB can be streamed to a sink in
1659 /// bounded memory. The leading (scalar) columns are decoded eagerly into the
1660 /// per-row [`Row`](crate::Row).
1661 ///
1662 /// The MAX column must be the **last** column. The returned
1663 /// [`BlobStream`](crate::BlobStream) yields scalar [`Row`](crate::Row)s via
1664 /// [`next`](crate::BlobStream::next); read each row's blob with
1665 /// [`read_chunk`](crate::BlobStream::read_chunk) /
1666 /// [`copy_blob_to`](crate::BlobStream::copy_blob_to) before advancing. Also
1667 /// available on `Client<InTransaction>`.
1668 ///
1669 /// # Errors
1670 ///
1671 /// Returns an error if the result set has no trailing MAX column, has more
1672 /// than one MAX column, the MAX column is not last, or the result set uses
1673 /// Always Encrypted (not yet supported on this path).
1674 pub async fn query_stream_blob<'a>(
1675 &'a mut self,
1676 sql: &str,
1677 params: &[&(dyn crate::ToSql + Sync)],
1678 ) -> Result<crate::blob_stream::BlobStream<'a, Ready>> {
1679 self.query_stream_blob_inner(sql, params).await
1680 }
1681
1682 /// Execute a query and stream a row's **trailing MAX columns** from the
1683 /// network — the multi-column generalization of
1684 /// [`query_stream_blob`](Self::query_stream_blob).
1685 ///
1686 /// For result sets whose trailing columns are one or more MAX types
1687 /// (`VARBINARY(MAX)`, `NVARCHAR(MAX)`, `VARCHAR(MAX)`, `XML`), this decodes
1688 /// the leading scalar columns eagerly into the per-row [`Row`](crate::Row)
1689 /// and streams each trailing MAX column's bytes incrementally from the
1690 /// socket, in bounded memory. The returned
1691 /// [`BlobStream`](crate::BlobStream) yields scalar rows via
1692 /// [`next`](crate::BlobStream::next); within each row, iterate the trailing
1693 /// MAX columns with [`next_blob`](crate::BlobStream::next_blob), reading each
1694 /// with [`copy_blob_to`](crate::BlobStream::copy_blob_to) /
1695 /// [`read_chunk`](crate::BlobStream::read_chunk). Also available on
1696 /// `Client<InTransaction>`.
1697 ///
1698 /// # Errors
1699 ///
1700 /// Returns an error if the result set has no trailing MAX column, has a
1701 /// non-MAX column after a MAX column (interleaved MAX columns are not
1702 /// supported — the MAX columns must be trailing), or uses Always Encrypted
1703 /// (not yet supported on this path).
1704 pub async fn query_stream_rows<'a>(
1705 &'a mut self,
1706 sql: &str,
1707 params: &[&(dyn crate::ToSql + Sync)],
1708 ) -> Result<crate::blob_stream::BlobStream<'a, Ready>> {
1709 self.query_stream_rows_inner(sql, params).await
1710 }
1711
1712 /// Execute a query with a specific timeout.
1713 ///
1714 /// This overrides the default `command_timeout` from the connection configuration
1715 /// for this specific query. If the query does not complete within the specified
1716 /// duration, the driver sends an Attention packet to cancel it server-side,
1717 /// drains the acknowledgement, and returns [`Error::CommandTimeout`] with the
1718 /// connection left usable for the next request.
1719 ///
1720 /// # Arguments
1721 ///
1722 /// * `sql` - The SQL query to execute
1723 /// * `params` - Query parameters
1724 /// * `timeout_duration` - Maximum time to wait for the query to complete
1725 ///
1726 /// # Example
1727 ///
1728 /// ```rust,no_run
1729 /// # async fn ex(client: &mut mssql_client::Client<mssql_client::Ready>) -> Result<(), mssql_client::Error> {
1730 /// use std::time::Duration;
1731 ///
1732 /// // Execute with a 5-second timeout
1733 /// let rows = client
1734 /// .query_with_timeout(
1735 /// "SELECT * FROM large_table",
1736 /// &[],
1737 /// Duration::from_secs(5),
1738 /// )
1739 /// .await?;
1740 /// # let _ = rows;
1741 /// # Ok(())
1742 /// # }
1743 /// ```
1744 pub async fn query_with_timeout<'a>(
1745 &'a mut self,
1746 sql: &str,
1747 params: &[&(dyn crate::ToSql + Sync)],
1748 timeout_duration: std::time::Duration,
1749 ) -> Result<QueryStream<'a>> {
1750 self.query_inner(sql, params, Some(timeout_duration)).await
1751 }
1752
1753 /// Execute a batch that may return multiple result sets.
1754 ///
1755 /// This is useful for stored procedures or SQL batches that contain
1756 /// multiple SELECT statements.
1757 ///
1758 /// # Example
1759 ///
1760 /// ```rust,no_run
1761 /// # async fn ex(client: &mut mssql_client::Client<mssql_client::Ready>) -> Result<(), mssql_client::Error> {
1762 /// // Execute a batch with multiple SELECT statements
1763 /// let mut results = client.query_multiple(
1764 /// "SELECT 1 AS a; SELECT 2 AS b, 3 AS c;",
1765 /// &[]
1766 /// ).await?;
1767 ///
1768 /// // Process first result set
1769 /// while let Some(row) = results.next_row().await? {
1770 /// println!("Result 1: {:?}", row);
1771 /// }
1772 ///
1773 /// // Move to second result set
1774 /// if results.next_result().await? {
1775 /// while let Some(row) = results.next_row().await? {
1776 /// println!("Result 2: {:?}", row);
1777 /// }
1778 /// }
1779 /// # Ok(())
1780 /// # }
1781 /// ```
1782 pub async fn query_multiple<'a>(
1783 &'a mut self,
1784 sql: &str,
1785 params: &[&(dyn crate::ToSql + Sync)],
1786 ) -> Result<MultiResultStream<'a>> {
1787 tracing::debug!(
1788 sql = sql,
1789 params_count = params.len(),
1790 "executing multi-result query"
1791 );
1792
1793 #[cfg(feature = "otel")]
1794 let instrumentation = self.instrumentation.clone();
1795 #[cfg(feature = "otel")]
1796 let mut span = instrumentation.query_span(sql);
1797 #[cfg(feature = "otel")]
1798 let timer = crate::instrumentation::OperationTimer::start(
1799 crate::instrumentation::extract_operation(sql),
1800 );
1801
1802 let deadline = self.command_deadline();
1803 let canceller = self.connection_cancel_handle();
1804 let result = run_with_deadline(
1805 async {
1806 if params.is_empty() {
1807 // Simple batch without parameters - use SQL batch
1808 self.send_sql_batch(sql).await?;
1809 } else {
1810 // Parameterized query - sp_executesql (encrypts Always Encrypted params).
1811 let rpc = self.build_parameterized_rpc(sql, params).await?;
1812 self.send_rpc(&rpc).await?;
1813 }
1814
1815 // Read all result sets
1816 self.read_multi_result_response().await
1817 },
1818 deadline,
1819 canceller,
1820 )
1821 .await;
1822
1823 #[cfg(feature = "otel")]
1824 match &result {
1825 Ok(_) => InstrumentationContext::record_success(&mut span, None),
1826 Err(e) => InstrumentationContext::record_error(&mut span, e),
1827 }
1828 #[cfg(feature = "otel")]
1829 timer.finish(instrumentation.metrics(), result.is_ok());
1830 #[cfg(feature = "otel")]
1831 drop(span);
1832
1833 let result_sets = result?;
1834 Ok(MultiResultStream::new(result_sets))
1835 }
1836
1837 /// Execute a query that doesn't return rows.
1838 ///
1839 /// Returns the number of affected rows.
1840 pub async fn execute(
1841 &mut self,
1842 sql: &str,
1843 params: &[&(dyn crate::ToSql + Sync)],
1844 ) -> Result<u64> {
1845 let deadline = self.command_deadline();
1846 self.execute_inner(sql, params, deadline).await
1847 }
1848
1849 /// Shared execute implementation with an explicit command deadline.
1850 async fn execute_inner(
1851 &mut self,
1852 sql: &str,
1853 params: &[&(dyn crate::ToSql + Sync)],
1854 deadline: Option<std::time::Duration>,
1855 ) -> Result<u64> {
1856 tracing::debug!(
1857 sql = sql,
1858 params_count = params.len(),
1859 "executing statement"
1860 );
1861
1862 #[cfg(feature = "otel")]
1863 let instrumentation = self.instrumentation.clone();
1864 #[cfg(feature = "otel")]
1865 let mut span = instrumentation.query_span(sql);
1866 #[cfg(feature = "otel")]
1867 let timer = crate::instrumentation::OperationTimer::start(
1868 crate::instrumentation::extract_operation(sql),
1869 );
1870
1871 let canceller = self.cancel_handle();
1872 let result = run_with_deadline(
1873 async {
1874 if params.is_empty() {
1875 // Simple statement without parameters - use SQL batch
1876 self.send_sql_batch(sql).await?;
1877 } else {
1878 // Parameterized statement - sp_executesql (encrypts Always Encrypted params).
1879 let rpc = self.build_parameterized_rpc(sql, params).await?;
1880 self.send_rpc(&rpc).await?;
1881 }
1882
1883 // Read response and get row count
1884 self.read_execute_result().await
1885 },
1886 deadline,
1887 canceller,
1888 )
1889 .await;
1890
1891 #[cfg(feature = "otel")]
1892 match &result {
1893 Ok(rows) => InstrumentationContext::record_success(&mut span, Some(*rows)),
1894 Err(e) => InstrumentationContext::record_error(&mut span, e),
1895 }
1896 #[cfg(feature = "otel")]
1897 timer.finish(instrumentation.metrics(), result.is_ok());
1898
1899 // Drop the span before returning
1900 #[cfg(feature = "otel")]
1901 drop(span);
1902
1903 result
1904 }
1905
1906 /// Execute a statement with a specific timeout.
1907 ///
1908 /// This overrides the default `command_timeout` from the connection configuration
1909 /// for this specific statement. If the statement does not complete within the
1910 /// specified duration, the driver sends an Attention packet to cancel it
1911 /// server-side, drains the acknowledgement, and returns
1912 /// [`Error::CommandTimeout`] with the connection left usable.
1913 ///
1914 /// # Arguments
1915 ///
1916 /// * `sql` - The SQL statement to execute
1917 /// * `params` - Statement parameters
1918 /// * `timeout_duration` - Maximum time to wait for the statement to complete
1919 ///
1920 /// # Example
1921 ///
1922 /// ```rust,no_run
1923 /// # async fn ex(client: &mut mssql_client::Client<mssql_client::Ready>) -> Result<(), mssql_client::Error> {
1924 /// use std::time::Duration;
1925 ///
1926 /// // Execute with a 10-second timeout
1927 /// let rows_affected = client
1928 /// .execute_with_timeout(
1929 /// "UPDATE large_table SET status = @p1",
1930 /// &[&"processed"],
1931 /// Duration::from_secs(10),
1932 /// )
1933 /// .await?;
1934 /// # let _ = rows_affected;
1935 /// # Ok(())
1936 /// # }
1937 /// ```
1938 pub async fn execute_with_timeout(
1939 &mut self,
1940 sql: &str,
1941 params: &[&(dyn crate::ToSql + Sync)],
1942 timeout_duration: std::time::Duration,
1943 ) -> Result<u64> {
1944 self.execute_inner(sql, params, Some(timeout_duration))
1945 .await
1946 }
1947
1948 /// Begin a transaction.
1949 ///
1950 /// This transitions the client from `Ready` to `InTransaction` state.
1951 /// Per MS-TDS spec, the server returns a transaction descriptor in the
1952 /// BeginTransaction EnvChange token that must be included in subsequent
1953 /// ALL_HEADERS sections.
1954 pub async fn begin_transaction(mut self) -> Result<Client<InTransaction>> {
1955 tracing::debug!("beginning transaction");
1956
1957 #[cfg(feature = "otel")]
1958 let instrumentation = self.instrumentation.clone();
1959 #[cfg(feature = "otel")]
1960 let mut span = instrumentation.transaction_span("BEGIN");
1961
1962 // Execute BEGIN TRANSACTION and extract the transaction descriptor
1963 let result = async {
1964 self.send_sql_batch("BEGIN TRANSACTION").await?;
1965 self.read_transaction_begin_result().await
1966 }
1967 .await;
1968
1969 #[cfg(feature = "otel")]
1970 match &result {
1971 Ok(_) => InstrumentationContext::record_success(&mut span, None),
1972 Err(e) => InstrumentationContext::record_error(&mut span, e),
1973 }
1974
1975 // Drop the span before moving instrumentation
1976 #[cfg(feature = "otel")]
1977 drop(span);
1978
1979 let transaction_descriptor = result?;
1980
1981 Ok(Client {
1982 config: self.config,
1983 _state: PhantomData,
1984 connection: self.connection,
1985 server_version: self.server_version,
1986 current_database: self.current_database,
1987 server_collation: self.server_collation,
1988 statement_cache: self.statement_cache,
1989 transaction_descriptor, // Store the descriptor from server
1990 needs_reset: self.needs_reset,
1991 in_flight: self.in_flight,
1992 #[cfg(feature = "otel")]
1993 instrumentation: self.instrumentation,
1994 #[cfg(feature = "always-encrypted")]
1995 encryption_context: self.encryption_context,
1996 })
1997 }
1998
1999 /// Begin a transaction with a specific isolation level.
2000 ///
2001 /// This transitions the client from `Ready` to `InTransaction` state
2002 /// with the specified isolation level.
2003 ///
2004 /// # Example
2005 ///
2006 /// ```rust,no_run
2007 /// # async fn ex(client: mssql_client::Client<mssql_client::Ready>) -> Result<(), mssql_client::Error> {
2008 /// use mssql_client::IsolationLevel;
2009 ///
2010 /// let tx = client.begin_transaction_with_isolation(IsolationLevel::Serializable).await?;
2011 /// // All operations in this transaction use SERIALIZABLE isolation
2012 /// tx.commit().await?;
2013 /// # Ok(())
2014 /// # }
2015 /// ```
2016 pub async fn begin_transaction_with_isolation(
2017 mut self,
2018 isolation_level: crate::transaction::IsolationLevel,
2019 ) -> Result<Client<InTransaction>> {
2020 tracing::debug!(
2021 isolation_level = %isolation_level.name(),
2022 "beginning transaction with isolation level"
2023 );
2024
2025 #[cfg(feature = "otel")]
2026 let instrumentation = self.instrumentation.clone();
2027 #[cfg(feature = "otel")]
2028 let mut span = instrumentation.transaction_span("BEGIN");
2029
2030 // First set the isolation level
2031 let result = async {
2032 self.send_sql_batch(isolation_level.as_sql()).await?;
2033 self.read_execute_result().await?;
2034
2035 // Then begin the transaction
2036 self.send_sql_batch("BEGIN TRANSACTION").await?;
2037 self.read_transaction_begin_result().await
2038 }
2039 .await;
2040
2041 #[cfg(feature = "otel")]
2042 match &result {
2043 Ok(_) => InstrumentationContext::record_success(&mut span, None),
2044 Err(e) => InstrumentationContext::record_error(&mut span, e),
2045 }
2046
2047 #[cfg(feature = "otel")]
2048 drop(span);
2049
2050 let transaction_descriptor = result?;
2051
2052 Ok(Client {
2053 config: self.config,
2054 _state: PhantomData,
2055 connection: self.connection,
2056 server_version: self.server_version,
2057 current_database: self.current_database,
2058 server_collation: self.server_collation,
2059 statement_cache: self.statement_cache,
2060 transaction_descriptor,
2061 needs_reset: self.needs_reset,
2062 in_flight: self.in_flight,
2063 #[cfg(feature = "otel")]
2064 instrumentation: self.instrumentation,
2065 #[cfg(feature = "always-encrypted")]
2066 encryption_context: self.encryption_context,
2067 })
2068 }
2069
2070 /// Execute a simple query without parameters.
2071 ///
2072 /// This is useful for DDL statements and simple queries where you
2073 /// don't need to retrieve the affected row count.
2074 pub async fn simple_query(&mut self, sql: &str) -> Result<()> {
2075 tracing::debug!(sql = sql, "executing simple query");
2076
2077 // Send SQL batch
2078 self.send_sql_batch(sql).await?;
2079
2080 // Read and discard response
2081 let _ = self.read_execute_result().await?;
2082
2083 Ok(())
2084 }
2085
2086 /// Close the connection gracefully.
2087 pub async fn close(self) -> Result<()> {
2088 tracing::debug!("closing connection");
2089 Ok(())
2090 }
2091
2092 /// Get the current database name.
2093 #[must_use]
2094 pub fn database(&self) -> Option<&str> {
2095 self.config.database.as_deref()
2096 }
2097
2098 /// Get the server host.
2099 #[must_use]
2100 pub fn host(&self) -> &str {
2101 &self.config.host
2102 }
2103
2104 /// Get the server port.
2105 #[must_use]
2106 pub fn port(&self) -> u16 {
2107 self.config.port
2108 }
2109
2110 /// Check if the connection is currently in a transaction.
2111 ///
2112 /// This returns `true` if a transaction was started via raw SQL
2113 /// (`BEGIN TRANSACTION`) and has not yet been committed or rolled back.
2114 ///
2115 /// Note: This only tracks transactions started via raw SQL. Transactions
2116 /// started via the type-state API (`begin_transaction()`) result in a
2117 /// `Client<InTransaction>` which is a different type.
2118 ///
2119 /// # Example
2120 ///
2121 /// ```rust,no_run
2122 /// # async fn ex(client: &mut mssql_client::Client<mssql_client::Ready>) -> Result<(), mssql_client::Error> {
2123 /// client.execute("BEGIN TRANSACTION", &[]).await?;
2124 /// assert!(client.is_in_transaction());
2125 ///
2126 /// client.execute("COMMIT", &[]).await?;
2127 /// assert!(!client.is_in_transaction());
2128 /// # Ok(())
2129 /// # }
2130 /// ```
2131 #[must_use]
2132 pub fn is_in_transaction(&self) -> bool {
2133 self.transaction_descriptor != 0
2134 }
2135
2136 /// Check if a request is in-flight (sent but response not fully read).
2137 ///
2138 /// Used by the connection pool to detect dirty connections that were
2139 /// interrupted mid-query (e.g., by `tokio::select!` or a timeout).
2140 /// A connection with an in-flight request has unread data in the TCP
2141 /// buffer and must be discarded rather than returned to the pool.
2142 #[must_use]
2143 pub fn is_in_flight(&self) -> bool {
2144 self.in_flight
2145 }
2146
2147 /// Report whether an Always Encrypted key-store provider with the given
2148 /// name is currently reachable through this client's encryption context.
2149 ///
2150 /// Returns `false` when the `always-encrypted` feature isn't enabled, when
2151 /// the connection was opened without `column_encryption` configured, or
2152 /// when no matching provider was registered.
2153 #[cfg(feature = "always-encrypted")]
2154 #[must_use]
2155 pub fn has_encryption_provider(&self, name: &str) -> bool {
2156 self.encryption_context
2157 .as_ref()
2158 .is_some_and(|ctx| ctx.has_provider(name))
2159 }
2160
2161 /// Get a handle for cancelling the current query.
2162 ///
2163 /// The cancel handle can be cloned and sent to other tasks, enabling
2164 /// cancellation of long-running queries from a separate async context.
2165 ///
2166 /// # Example
2167 ///
2168 /// ```rust,no_run
2169 /// # async fn ex(client: &mut mssql_client::Client<mssql_client::Ready>) -> Result<(), mssql_client::Error> {
2170 /// use std::time::Duration;
2171 ///
2172 /// let cancel_handle = client.cancel_handle();
2173 ///
2174 /// // Spawn a task to cancel after 10 seconds
2175 /// let handle = tokio::spawn(async move {
2176 /// tokio::time::sleep(Duration::from_secs(10)).await;
2177 /// let _ = cancel_handle.cancel().await;
2178 /// });
2179 ///
2180 /// // This query will be cancelled if it runs longer than 10 seconds
2181 /// let result = client.query("SELECT * FROM very_large_table", &[]).await;
2182 /// # let _ = (handle, result);
2183 /// # Ok(())
2184 /// # }
2185 /// ```
2186 #[must_use]
2187 pub fn cancel_handle(&self) -> crate::cancel::CancelHandle {
2188 self.connection_cancel_handle()
2189 }
2190}
2191
2192/// # Drop Behavior
2193///
2194/// **`Client<InTransaction>` has no automatic rollback on drop.** If the client is
2195/// dropped without calling [`commit()`](Client::commit) or [`rollback()`](Client::rollback),
2196/// the transaction remains open on the server until the TCP connection closes
2197/// (at which point SQL Server automatically rolls back).
2198///
2199/// This is because `Drop` is synchronous and cannot perform the async I/O needed
2200/// to send a `ROLLBACK TRANSACTION` command.
2201///
2202/// ## Consequences of dropping without commit/rollback
2203///
2204/// - **Direct connections:** The transaction leaks until the OS TCP timeout
2205/// (potentially 30+ minutes), holding locks on any modified rows.
2206/// - **Pooled connections:** The pool detects the active transaction descriptor
2207/// and discards the connection rather than returning it to the idle pool
2208/// (see `PooledConnection::drop` in `mssql-driver-pool`).
2209///
2210/// ## Best practice
2211///
2212/// Always ensure `commit()` or `rollback()` is called. Use helper patterns
2213/// for error paths:
2214///
2215/// ```rust,no_run
2216/// # async fn do_work(_: &mssql_client::Client<mssql_client::InTransaction>) -> Result<(), mssql_client::Error> { Ok(()) }
2217/// # async fn ex(client: mssql_client::Client<mssql_client::Ready>) -> Result<(), mssql_client::Error> {
2218/// let tx = client.begin_transaction().await?;
2219/// match do_work(&tx).await {
2220/// Ok(_) => { tx.commit().await?; }
2221/// Err(e) => { tx.rollback().await?; return Err(e); }
2222/// }
2223/// # Ok(())
2224/// # }
2225/// ```
2226impl Client<InTransaction> {
2227 /// Execute a query within the transaction and return a streaming result set.
2228 ///
2229 /// See [`Client<Ready>::query`] for usage examples.
2230 pub async fn query<'a>(
2231 &'a mut self,
2232 sql: &str,
2233 params: &[&(dyn crate::ToSql + Sync)],
2234 ) -> Result<QueryStream<'a>> {
2235 let deadline = self.command_deadline();
2236 self.query_inner(sql, params, deadline).await
2237 }
2238
2239 /// Shared query implementation with an explicit command deadline.
2240 async fn query_inner<'a>(
2241 &'a mut self,
2242 sql: &str,
2243 params: &[&(dyn crate::ToSql + Sync)],
2244 deadline: Option<std::time::Duration>,
2245 ) -> Result<QueryStream<'a>> {
2246 tracing::debug!(
2247 sql = sql,
2248 params_count = params.len(),
2249 "executing query in transaction"
2250 );
2251
2252 #[cfg(feature = "otel")]
2253 let instrumentation = self.instrumentation.clone();
2254 #[cfg(feature = "otel")]
2255 let mut span = instrumentation.query_span(sql);
2256 #[cfg(feature = "otel")]
2257 let timer = crate::instrumentation::OperationTimer::start(
2258 crate::instrumentation::extract_operation(sql),
2259 );
2260
2261 let canceller = self.cancel_handle();
2262 let result = run_with_deadline(
2263 async {
2264 // Sends via the prepared-statement cache when enabled, else the
2265 // SQL batch / sp_executesql default.
2266 self.send_query_request(sql, params).await?;
2267
2268 // Read complete response including columns and rows
2269 self.read_query_response().await
2270 },
2271 deadline,
2272 canceller,
2273 )
2274 .await;
2275
2276 #[cfg(feature = "otel")]
2277 match &result {
2278 Ok(_) => InstrumentationContext::record_success(&mut span, None),
2279 Err(e) => InstrumentationContext::record_error(&mut span, e),
2280 }
2281 #[cfg(feature = "otel")]
2282 timer.finish(instrumentation.metrics(), result.is_ok());
2283
2284 // Drop the span before returning
2285 #[cfg(feature = "otel")]
2286 drop(span);
2287
2288 let resp = result?;
2289 #[cfg(feature = "always-encrypted")]
2290 {
2291 Ok(QueryStream::from_raw(
2292 resp.columns,
2293 resp.pending_rows,
2294 resp.meta,
2295 resp.decryptor,
2296 ))
2297 }
2298 #[cfg(not(feature = "always-encrypted"))]
2299 {
2300 Ok(QueryStream::from_raw(
2301 resp.columns,
2302 resp.pending_rows,
2303 resp.meta,
2304 ))
2305 }
2306 }
2307
2308 /// Stream rows incrementally from the network within the transaction.
2309 ///
2310 /// Identical to [`Client<Ready>::query_stream`] except the query runs inside
2311 /// the open transaction. The returned [`RowStream`](crate::RowStream)
2312 /// borrows the transaction client for its lifetime, so the stream must be
2313 /// consumed or dropped before the transaction can be committed or rolled
2314 /// back.
2315 pub async fn query_stream<'a>(
2316 &'a mut self,
2317 sql: &str,
2318 params: &[&(dyn crate::ToSql + Sync)],
2319 ) -> Result<crate::row_stream::RowStream<'a, InTransaction>> {
2320 self.query_stream_inner(sql, params).await
2321 }
2322
2323 /// Stream a row's trailing MAX column from the network within the
2324 /// transaction.
2325 ///
2326 /// See [`Client<Ready>::query_stream_blob`] for semantics and constraints;
2327 /// the only difference is that the query runs inside the open transaction.
2328 pub async fn query_stream_blob<'a>(
2329 &'a mut self,
2330 sql: &str,
2331 params: &[&(dyn crate::ToSql + Sync)],
2332 ) -> Result<crate::blob_stream::BlobStream<'a, InTransaction>> {
2333 self.query_stream_blob_inner(sql, params).await
2334 }
2335
2336 /// Stream a row's trailing MAX columns from the network within the
2337 /// transaction.
2338 ///
2339 /// See [`Client<Ready>::query_stream_rows`] for semantics and constraints;
2340 /// the only difference is that the query runs inside the open transaction.
2341 pub async fn query_stream_rows<'a>(
2342 &'a mut self,
2343 sql: &str,
2344 params: &[&(dyn crate::ToSql + Sync)],
2345 ) -> Result<crate::blob_stream::BlobStream<'a, InTransaction>> {
2346 self.query_stream_rows_inner(sql, params).await
2347 }
2348
2349 /// Execute a statement within the transaction.
2350 ///
2351 /// Returns the number of affected rows.
2352 pub async fn execute(
2353 &mut self,
2354 sql: &str,
2355 params: &[&(dyn crate::ToSql + Sync)],
2356 ) -> Result<u64> {
2357 let deadline = self.command_deadline();
2358 self.execute_inner(sql, params, deadline).await
2359 }
2360
2361 /// Shared execute implementation with an explicit command deadline.
2362 async fn execute_inner(
2363 &mut self,
2364 sql: &str,
2365 params: &[&(dyn crate::ToSql + Sync)],
2366 deadline: Option<std::time::Duration>,
2367 ) -> Result<u64> {
2368 tracing::debug!(
2369 sql = sql,
2370 params_count = params.len(),
2371 "executing statement in transaction"
2372 );
2373
2374 #[cfg(feature = "otel")]
2375 let instrumentation = self.instrumentation.clone();
2376 #[cfg(feature = "otel")]
2377 let mut span = instrumentation.query_span(sql);
2378 #[cfg(feature = "otel")]
2379 let timer = crate::instrumentation::OperationTimer::start(
2380 crate::instrumentation::extract_operation(sql),
2381 );
2382
2383 let canceller = self.cancel_handle();
2384 let result = run_with_deadline(
2385 async {
2386 if params.is_empty() {
2387 // Simple statement without parameters - use SQL batch
2388 self.send_sql_batch(sql).await?;
2389 } else {
2390 // Parameterized statement - sp_executesql (encrypts Always Encrypted params).
2391 let rpc = self.build_parameterized_rpc(sql, params).await?;
2392 self.send_rpc(&rpc).await?;
2393 }
2394
2395 // Read response and get row count
2396 self.read_execute_result().await
2397 },
2398 deadline,
2399 canceller,
2400 )
2401 .await;
2402
2403 #[cfg(feature = "otel")]
2404 match &result {
2405 Ok(rows) => InstrumentationContext::record_success(&mut span, Some(*rows)),
2406 Err(e) => InstrumentationContext::record_error(&mut span, e),
2407 }
2408 #[cfg(feature = "otel")]
2409 timer.finish(instrumentation.metrics(), result.is_ok());
2410
2411 // Drop the span before returning
2412 #[cfg(feature = "otel")]
2413 drop(span);
2414
2415 result
2416 }
2417
2418 /// Execute a query within the transaction with a specific timeout.
2419 ///
2420 /// See [`Client<Ready>::query_with_timeout`] for details.
2421 pub async fn query_with_timeout<'a>(
2422 &'a mut self,
2423 sql: &str,
2424 params: &[&(dyn crate::ToSql + Sync)],
2425 timeout_duration: std::time::Duration,
2426 ) -> Result<QueryStream<'a>> {
2427 self.query_inner(sql, params, Some(timeout_duration)).await
2428 }
2429
2430 /// Execute a statement within the transaction with a specific timeout.
2431 ///
2432 /// See [`Client<Ready>::execute_with_timeout`] for details.
2433 pub async fn execute_with_timeout(
2434 &mut self,
2435 sql: &str,
2436 params: &[&(dyn crate::ToSql + Sync)],
2437 timeout_duration: std::time::Duration,
2438 ) -> Result<u64> {
2439 self.execute_inner(sql, params, Some(timeout_duration))
2440 .await
2441 }
2442
2443 /// Open a FILESTREAM BLOB for async reading and/or writing.
2444 ///
2445 /// This method queries the server for the transaction context, then opens
2446 /// the FILESTREAM handle using the native Win32 `OpenSqlFilestream` API.
2447 ///
2448 /// # Arguments
2449 ///
2450 /// * `path` — The UNC path obtained from the T-SQL `column.PathName()` function.
2451 /// Query this yourself before calling `open_filestream`:
2452 /// ```sql
2453 /// SELECT Content.PathName() FROM dbo.Documents WHERE Id = @p1
2454 /// ```
2455 /// * `access` — Read, write, or read/write access mode.
2456 ///
2457 /// # Requirements
2458 ///
2459 /// - SQL Server must have FILESTREAM enabled (`sp_configure 'filestream access level', 2`)
2460 /// - The Microsoft OLE DB Driver for SQL Server must be installed on the client
2461 /// - The `FileStream` must be dropped before calling [`commit`] or [`rollback`]
2462 ///
2463 /// # Example
2464 ///
2465 /// ```text
2466 /// use mssql_client::FileStreamAccess;
2467 /// use tokio::io::AsyncReadExt;
2468 ///
2469 /// let mut tx = client.begin_transaction().await?;
2470 ///
2471 /// // Get the FILESTREAM path
2472 /// let rows = tx.query(
2473 /// "SELECT Content.PathName() FROM dbo.Documents WHERE Id = @p1",
2474 /// &[&doc_id],
2475 /// ).await?;
2476 /// let path: String = rows.into_iter().next().unwrap()?.get(0)?;
2477 ///
2478 /// // Open and read the BLOB
2479 /// let mut stream = tx.open_filestream(&path, FileStreamAccess::Read).await?;
2480 /// let mut data = Vec::new();
2481 /// stream.read_to_end(&mut data).await?;
2482 /// drop(stream);
2483 ///
2484 /// tx.commit().await?;
2485 /// ```
2486 #[cfg(all(windows, feature = "filestream"))]
2487 pub async fn open_filestream(
2488 &mut self,
2489 path: &str,
2490 access: crate::filestream::FileStreamAccess,
2491 ) -> Result<crate::filestream::FileStream> {
2492 tracing::debug!(path = path, ?access, "opening FILESTREAM BLOB");
2493
2494 // Get the transaction context from SQL Server.
2495 // This binds the file access to the current SQL transaction.
2496 let txn_context: Vec<u8> = {
2497 let rows = self
2498 .query("SELECT GET_FILESTREAM_TRANSACTION_CONTEXT()", &[])
2499 .await?;
2500 let mut ctx = None;
2501 for result in rows {
2502 let row = result?;
2503 ctx = Some(row.get::<Vec<u8>>(0)?);
2504 }
2505 ctx.ok_or_else(|| {
2506 Error::FileStream("GET_FILESTREAM_TRANSACTION_CONTEXT() returned no rows".into())
2507 })?
2508 };
2509
2510 crate::filestream::FileStream::open(path, access, &txn_context)
2511 }
2512
2513 /// Commit the transaction.
2514 ///
2515 /// This transitions the client back to `Ready` state.
2516 pub async fn commit(mut self) -> Result<Client<Ready>> {
2517 tracing::debug!("committing transaction");
2518
2519 #[cfg(feature = "otel")]
2520 let instrumentation = self.instrumentation.clone();
2521 #[cfg(feature = "otel")]
2522 let mut span = instrumentation.transaction_span("COMMIT");
2523
2524 // Execute COMMIT TRANSACTION
2525 let result = async {
2526 self.send_sql_batch("COMMIT TRANSACTION").await?;
2527 self.read_execute_result().await
2528 }
2529 .await;
2530
2531 #[cfg(feature = "otel")]
2532 match &result {
2533 Ok(_) => InstrumentationContext::record_success(&mut span, None),
2534 Err(e) => InstrumentationContext::record_error(&mut span, e),
2535 }
2536
2537 // Drop the span before moving instrumentation
2538 #[cfg(feature = "otel")]
2539 drop(span);
2540
2541 result?;
2542
2543 Ok(Client {
2544 config: self.config,
2545 _state: PhantomData,
2546 connection: self.connection,
2547 server_version: self.server_version,
2548 current_database: self.current_database,
2549 server_collation: self.server_collation,
2550 statement_cache: self.statement_cache,
2551 transaction_descriptor: 0, // Reset to auto-commit mode
2552 needs_reset: self.needs_reset,
2553 in_flight: self.in_flight,
2554 #[cfg(feature = "otel")]
2555 instrumentation: self.instrumentation,
2556 #[cfg(feature = "always-encrypted")]
2557 encryption_context: self.encryption_context,
2558 })
2559 }
2560
2561 /// Rollback the transaction.
2562 ///
2563 /// This transitions the client back to `Ready` state.
2564 pub async fn rollback(mut self) -> Result<Client<Ready>> {
2565 tracing::debug!("rolling back transaction");
2566
2567 #[cfg(feature = "otel")]
2568 let instrumentation = self.instrumentation.clone();
2569 #[cfg(feature = "otel")]
2570 let mut span = instrumentation.transaction_span("ROLLBACK");
2571
2572 // Execute ROLLBACK TRANSACTION
2573 let result = async {
2574 self.send_sql_batch("ROLLBACK TRANSACTION").await?;
2575 self.read_execute_result().await
2576 }
2577 .await;
2578
2579 #[cfg(feature = "otel")]
2580 match &result {
2581 Ok(_) => InstrumentationContext::record_success(&mut span, None),
2582 Err(e) => InstrumentationContext::record_error(&mut span, e),
2583 }
2584
2585 // Drop the span before moving instrumentation
2586 #[cfg(feature = "otel")]
2587 drop(span);
2588
2589 result?;
2590
2591 Ok(Client {
2592 config: self.config,
2593 _state: PhantomData,
2594 connection: self.connection,
2595 server_version: self.server_version,
2596 current_database: self.current_database,
2597 server_collation: self.server_collation,
2598 statement_cache: self.statement_cache,
2599 transaction_descriptor: 0, // Reset to auto-commit mode
2600 needs_reset: self.needs_reset,
2601 in_flight: self.in_flight,
2602 #[cfg(feature = "otel")]
2603 instrumentation: self.instrumentation,
2604 #[cfg(feature = "always-encrypted")]
2605 encryption_context: self.encryption_context,
2606 })
2607 }
2608
2609 /// Create a savepoint and return a handle for later rollback.
2610 ///
2611 /// The returned `SavePoint` handle contains the validated savepoint name.
2612 /// Use it with `rollback_to()` to partially undo transaction work.
2613 ///
2614 /// # Example
2615 ///
2616 /// ```rust,no_run
2617 /// # async fn ex(client: mssql_client::Client<mssql_client::Ready>) -> Result<(), mssql_client::Error> {
2618 /// let mut tx = client.begin_transaction().await?;
2619 /// tx.execute("INSERT INTO orders ...", &[]).await?;
2620 /// let sp = tx.save_point("before_items").await?;
2621 /// tx.execute("INSERT INTO items ...", &[]).await?;
2622 /// // Oops, rollback just the items
2623 /// tx.rollback_to(&sp).await?;
2624 /// tx.commit().await?;
2625 /// # Ok(())
2626 /// # }
2627 /// ```
2628 pub async fn save_point(&mut self, name: &str) -> Result<SavePoint> {
2629 crate::validation::validate_identifier(name)?;
2630 tracing::debug!(name = name, "creating savepoint");
2631
2632 // Execute SAVE TRANSACTION <name>
2633 // Note: name is validated by validate_identifier() to prevent SQL injection
2634 let sql = format!("SAVE TRANSACTION {name}");
2635 self.send_sql_batch(&sql).await?;
2636 self.read_execute_result().await?;
2637
2638 Ok(SavePoint::new(name.to_string()))
2639 }
2640
2641 /// Rollback to a savepoint.
2642 ///
2643 /// This rolls back all changes made after the savepoint was created,
2644 /// but keeps the transaction active. The savepoint remains valid and
2645 /// can be rolled back to again.
2646 ///
2647 /// # Example
2648 ///
2649 /// ```rust,no_run
2650 /// # async fn ex(mut tx: mssql_client::Client<mssql_client::InTransaction>) -> Result<(), mssql_client::Error> {
2651 /// let sp = tx.save_point("checkpoint").await?;
2652 /// // ... do some work ...
2653 /// tx.rollback_to(&sp).await?; // Undo changes since checkpoint
2654 /// // Transaction is still active, savepoint is still valid
2655 /// # Ok(())
2656 /// # }
2657 /// ```
2658 pub async fn rollback_to(&mut self, savepoint: &SavePoint) -> Result<()> {
2659 tracing::debug!(name = savepoint.name(), "rolling back to savepoint");
2660
2661 // Execute ROLLBACK TRANSACTION <name>
2662 // Note: savepoint name was validated during creation
2663 let sql = format!("ROLLBACK TRANSACTION {}", savepoint.name());
2664 self.send_sql_batch(&sql).await?;
2665 self.read_execute_result().await?;
2666
2667 Ok(())
2668 }
2669
2670 /// Release a savepoint (optional cleanup).
2671 ///
2672 /// Note: SQL Server doesn't have explicit savepoint release, but this
2673 /// method is provided for API completeness. The savepoint is automatically
2674 /// released when the transaction commits or rolls back.
2675 pub async fn release_savepoint(&mut self, savepoint: SavePoint) -> Result<()> {
2676 tracing::debug!(name = savepoint.name(), "releasing savepoint");
2677
2678 // SQL Server doesn't require explicit savepoint release
2679 // The savepoint is implicitly released on commit/rollback
2680 // This method exists for API completeness
2681 drop(savepoint);
2682 Ok(())
2683 }
2684
2685 /// Get a handle for cancelling the current query within the transaction.
2686 ///
2687 /// See [`Client<Ready>::cancel_handle`] for usage examples.
2688 #[must_use]
2689 pub fn cancel_handle(&self) -> crate::cancel::CancelHandle {
2690 self.connection_cancel_handle()
2691 }
2692}
2693
2694impl<S: ConnectionState> std::fmt::Debug for Client<S> {
2695 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2696 f.debug_struct("Client")
2697 .field("host", &self.config.host)
2698 .field("port", &self.config.port)
2699 .field("database", &self.config.database)
2700 .finish()
2701 }
2702}
2703
2704#[cfg(test)]
2705mod blob_result_set_validation_tests {
2706 use tds_protocol::token::{ColMetaData, ColumnData, TypeInfo};
2707 use tds_protocol::types::TypeId;
2708
2709 use super::{Client, Ready};
2710 use crate::error::Error;
2711
2712 /// A scalar (non-MAX) column.
2713 fn scalar(name: &str) -> ColumnData {
2714 col(name, TypeId::Int4, None)
2715 }
2716
2717 /// A MAX (PLP) column: `max_length == 0xFFFF` marks the MAX variant.
2718 fn blob(name: &str) -> ColumnData {
2719 col(name, TypeId::BigVarBinary, Some(0xFFFF))
2720 }
2721
2722 fn col(name: &str, type_id: TypeId, max_length: Option<u32>) -> ColumnData {
2723 ColumnData {
2724 name: name.to_string(),
2725 type_id,
2726 col_type: 0,
2727 flags: 0,
2728 user_type: 0,
2729 type_info: TypeInfo {
2730 max_length,
2731 ..Default::default()
2732 },
2733 crypto_metadata: None,
2734 }
2735 }
2736
2737 fn meta(columns: Vec<ColumnData>) -> ColMetaData {
2738 ColMetaData {
2739 columns,
2740 cek_table: None,
2741 }
2742 }
2743
2744 fn validate(columns: Vec<ColumnData>) -> Result<(usize, usize), Error> {
2745 Client::<Ready>::validate_blob_rows_result_set(&meta(columns))
2746 }
2747
2748 #[test]
2749 fn single_trailing_blob() {
2750 assert_eq!(validate(vec![scalar("id"), blob("doc")]).unwrap(), (1, 1));
2751 }
2752
2753 #[test]
2754 fn multiple_trailing_blobs() {
2755 assert_eq!(
2756 validate(vec![scalar("id"), blob("doc1"), blob("doc2")]).unwrap(),
2757 (1, 2)
2758 );
2759 }
2760
2761 #[test]
2762 fn all_columns_blobs() {
2763 assert_eq!(validate(vec![blob("a"), blob("b")]).unwrap(), (0, 2));
2764 }
2765
2766 #[test]
2767 fn no_max_column_is_rejected() {
2768 assert!(matches!(
2769 validate(vec![scalar("id"), scalar("j")]),
2770 Err(Error::Protocol(_))
2771 ));
2772 }
2773
2774 #[test]
2775 fn scalar_after_blob_is_rejected() {
2776 // Interleaved MAX columns are out of scope: a scalar after a blob.
2777 assert!(matches!(
2778 validate(vec![scalar("id"), blob("doc"), scalar("trailing")]),
2779 Err(Error::Protocol(_))
2780 ));
2781 // ...even when more blobs follow the interloping scalar.
2782 assert!(matches!(
2783 validate(vec![blob("doc1"), scalar("mid"), blob("doc2")]),
2784 Err(Error::Protocol(_))
2785 ));
2786 }
2787}