diff --git a/src/client.rs b/src/client.rs index c6176d4d..a96dbcdd 100644 --- a/src/client.rs +++ b/src/client.rs @@ -20,8 +20,8 @@ pub use cancellation::CancellationToken; pub use config::*; pub(crate) use connection::*; pub use cursor::{ - Cursor, CursorConcurrencyOptions, CursorHandle, CursorOpenOptions, CursorScrollOptions, Fetch, - PreparedCursor, + Cursor, CursorConcurrencyOptions, CursorHandle, CursorOpenOptions, CursorPrepExecOutcome, + CursorScrollOptions, DirectResultSet, DirectResults, Fetch, PreparedCursor, }; pub use prepared::{PreparedHandle, PreparedStatement}; pub use rpc_response::OutputValue; @@ -541,21 +541,31 @@ impl Client { /// /// `param_defs` follows the same format as [`prepare`](Self::prepare); /// pass an empty string when the statement has no parameters. + /// + /// The return value is a [`CursorPrepExecOutcome`]: usually the server + /// opens a cursor ([`Cursor`](CursorPrepExecOutcome::Cursor)), but when the + /// statement is opened read-only with + /// [`CursorConcurrencyOptions::AllowDirect`] the server may take the + /// *AllowDirect* fast path — preparing the statement but skipping the + /// cursor and streaming the result sets inline + /// ([`Direct`](CursorPrepExecOutcome::Direct)). Either way the prepared + /// handle must be released (via [`PreparedCursor::unprepare`] or + /// [`DirectResults::unprepare`]). pub async fn cursor_prep_exec<'a>( &mut self, sql: impl Into>, options: cursor::CursorOpenOptions, param_defs: impl Into>, params: &[&'a dyn ToSql], - ) -> crate::Result { + ) -> crate::Result { self.connection.flush_stream().await?; let rpc_params = cursor::build_cursorprepexec_params(sql.into(), options, param_defs.into(), params); self.send_rpc(RpcProcId::CursorPrepExec, rpc_params).await?; - let (outputs, _status, metadata) = - rpc_response::collect_rpc_outputs_with_metadata(&mut self.connection).await?; - cursor::prepared_cursor_from_outputs(&outputs, metadata) + let (result_sets, outputs, _status) = + rpc_response::collect_rpc_result_sets(&mut self.connection).await?; + cursor::cursor_prep_exec_outcome(&outputs, result_sets) } /// Prepare and execute a SQL statement in a single round trip. Returns diff --git a/src/client/cursor.rs b/src/client/cursor.rs index 26a19844..65db6af8 100644 --- a/src/client/cursor.rs +++ b/src/client/cursor.rs @@ -9,15 +9,18 @@ //! until the connection closes; a warning is emitted via `tracing`. use std::borrow::Cow; +use std::sync::Arc; use enumflags2::{bitflags, BitFlags}; use futures_util::io::{AsyncRead, AsyncWrite}; use tracing::{event, Level}; -use crate::client::rpc_response::{collect_metadata_only_rpc, collect_rpc_outputs, OutputValue}; +use crate::client::rpc_response::{ + collect_metadata_only_rpc, collect_rpc_outputs, BufferedResultSet, OutputValue, +}; use crate::tds::codec::{ColumnData, RpcParam, RpcProcId, RpcStatus}; use crate::tds::stream::{QueryStream, TokenStream}; -use crate::{Client, Column, PreparedHandle, ToSql}; +use crate::{Client, Column, PreparedHandle, Row, ToSql}; /// Scroll options for `sp_cursoropen` / `sp_cursorprepexec` (TDS §2.2.6.7). /// @@ -349,6 +352,172 @@ impl Drop for PreparedCursor { } } +/// The outcome of [`Client::cursor_prep_exec`](crate::Client::cursor_prep_exec). +/// +/// `sp_cursorprepexec` normally opens a server-side cursor, but when the +/// statement is opened read-only with the +/// [`AllowDirect`](CursorConcurrencyOptions::AllowDirect) concurrency option +/// the server may elect the *AllowDirect* fast path: it prepares the statement +/// but skips opening a cursor (`@cursor == 0`) and instead streams the result +/// sets inline during the RPC response. This enum reports which path the +/// server took. +#[derive(Debug)] +pub enum CursorPrepExecOutcome { + /// The server opened a cursor (`@cursor` nonzero). Page through it with + /// [`PreparedCursor::fetch`]. + Cursor(PreparedCursor), + /// The server executed the statement directly (`@cursor == 0`), returning + /// the result sets inline. The prepared handle still needs releasing — see + /// [`DirectResults::unprepare`]. + Direct(DirectResults), +} + +impl CursorPrepExecOutcome { + /// The server-assigned prepared statement handle, which is present in both + /// outcomes. + pub fn prepared_handle(&self) -> PreparedHandle { + match self { + CursorPrepExecOutcome::Cursor(c) => c.prepared_handle(), + CursorPrepExecOutcome::Direct(d) => d.prepared_handle(), + } + } + + /// `true` if the server took the AllowDirect fast path. + pub fn is_direct(&self) -> bool { + matches!(self, CursorPrepExecOutcome::Direct(_)) + } + + /// Consume the outcome, returning the [`PreparedCursor`] if the server + /// opened a cursor. + /// + /// An AllowDirect response is returned intact as `Err(DirectResults)` so + /// its prepared handle and buffered rows remain available for cleanup and + /// consumption. + pub fn into_cursor(self) -> std::result::Result { + match self { + CursorPrepExecOutcome::Cursor(c) => Ok(c), + CursorPrepExecOutcome::Direct(d) => Err(d), + } + } + + /// Consume the outcome, returning the [`DirectResults`] if the server took + /// the AllowDirect fast path. + /// + /// A normal cursor response is returned intact as `Err(PreparedCursor)` so + /// both server-side handles remain available for fetching and cleanup. + pub fn into_direct(self) -> std::result::Result { + match self { + CursorPrepExecOutcome::Direct(d) => Ok(d), + CursorPrepExecOutcome::Cursor(c) => Err(c), + } + } +} + +/// The result sets a server returned inline from an AllowDirect +/// `sp_cursorprepexec`, together with the prepared handle that must still be +/// released. +/// +/// No cursor was opened, so there is nothing to fetch or close — the buffered +/// [`results`](Self::results) are the complete response. Call +/// [`unprepare`](Self::unprepare) to release the prepared handle (and take +/// ownership of the result sets); dropping without unpreparing leaks the +/// handle until the connection closes and logs a warning. +#[derive(Debug)] +pub struct DirectResults { + prepared_handle: PreparedHandle, + results: Vec, + scrollopt: BitFlags, + ccopt: BitFlags, + row_count: i32, + released: bool, +} + +impl DirectResults { + /// The server-assigned prepared statement handle. + pub fn prepared_handle(&self) -> PreparedHandle { + self.prepared_handle + } + + /// Negotiated scroll flags returned by the server. + pub fn scroll_options(&self) -> BitFlags { + self.scrollopt + } + + /// Negotiated concurrency flags returned by the server (includes the + /// [`AllowDirect`](CursorConcurrencyOptions::AllowDirect) bit). + pub fn concurrency_options(&self) -> BitFlags { + self.ccopt + } + + /// Server-reported row count, or `-1` if unknown. + pub fn row_count(&self) -> i32 { + self.row_count + } + + /// The buffered result sets, in the order the server streamed them. + pub fn results(&self) -> &[DirectResultSet] { + &self.results + } + + /// Release the prepared handle via `sp_cursorunprepare`, returning the + /// owned result sets. + /// + /// There is no cursor to close, so this only unprepares the statement. + pub async fn unprepare( + mut self, + client: &mut Client, + ) -> crate::Result> + where + S: AsyncRead + AsyncWrite + Unpin + Send, + { + client.connection.flush_stream().await?; + let handle_param = build_cursorunprepare_param(self.prepared_handle); + client + .send_rpc(RpcProcId::CursorUnprepare, vec![handle_param]) + .await?; + self.released = true; + collect_rpc_outputs(&mut client.connection).await?; + Ok(std::mem::take(&mut self.results)) + } +} + +impl Drop for DirectResults { + fn drop(&mut self) { + if !self.released { + event!( + Level::WARN, + prepared_handle = self.prepared_handle.as_i32(), + "DirectResults dropped without unprepare; server-side prepared handle will leak until the connection closes" + ); + } + } +} + +/// A single result set from an AllowDirect execution: its column metadata and +/// the rows the server streamed for it. +#[derive(Debug)] +pub struct DirectResultSet { + columns: Arc>, + rows: Vec, +} + +impl DirectResultSet { + /// The column metadata for this result set. + pub fn columns(&self) -> &[Column] { + &self.columns + } + + /// The buffered rows, in order. + pub fn rows(&self) -> &[Row] { + &self.rows + } + + /// Consume this result set, returning its owned rows. + pub fn into_rows(self) -> Vec { + self.rows + } +} + /// A server-side cursor handle. /// /// Obtain via [`Client::open_cursor`](crate::Client::open_cursor). Page @@ -712,10 +881,87 @@ pub(crate) fn prepared_cursor_from_outputs( }) } +/// Interpret an `sp_cursorprepexec` response, distinguishing a normal +/// prepared cursor from an AllowDirect direct-result response. +/// +/// The server signals AllowDirect by returning a nonzero `@prepared_handle` +/// but a zero `@cursor`, streaming the result sets inline (buffered here into +/// `result_sets`). A nonzero `@cursor` is the normal cursor path and is +/// handled exactly as before by [`prepared_cursor_from_outputs`]; the buffered +/// rows (a normal open streams only metadata) are used solely to seed the +/// cursor's cached metadata. +pub(crate) fn cursor_prep_exec_outcome( + outputs: &[OutputValue], + result_sets: Vec, +) -> crate::Result { + let lookup_named = |name: &str| -> Option { + outputs + .iter() + .find(|o| !o.name().is_empty() && o.matches_name(name)) + .and_then(|o| o.get::().ok().flatten()) + }; + let by_pos = |idx: usize| -> Option { + outputs.get(idx).and_then(|o| o.get::().ok().flatten()) + }; + + let prepared_handle = lookup_named("prepared_handle") + .or_else(|| lookup_named("handle")) + .or_else(|| by_pos(0)) + .ok_or_else(|| { + crate::Error::Protocol( + "sp_cursorprepexec: missing @prepared_handle output parameter".into(), + ) + })?; + if prepared_handle == 0 { + return Err(crate::Error::Protocol( + "sp_cursorprepexec: server returned a zero @prepared_handle".into(), + )); + } + + // A zero @cursor (with a valid prepared handle) is the AllowDirect signal: + // the server executed directly and streamed the result sets inline. A + // nonzero or absent @cursor falls through to the normal cursor path. + match lookup_named("cursor").or_else(|| by_pos(1)) { + Some(0) => { + let scrollopt = lookup_named("scrollopt").or_else(|| by_pos(2)).unwrap_or(0); + let ccopt = lookup_named("ccopt").or_else(|| by_pos(3)).unwrap_or(0); + let row_count = lookup_named("rowcount").or_else(|| by_pos(4)).unwrap_or(-1); + + let results = result_sets + .into_iter() + .map(|rs| DirectResultSet { + columns: rs.columns, + rows: rs.rows, + }) + .collect(); + + Ok(CursorPrepExecOutcome::Direct(DirectResults { + prepared_handle: PreparedHandle::from_i32(prepared_handle), + results, + scrollopt: i32_to_scroll_flags(scrollopt), + ccopt: i32_to_cc_flags(ccopt), + row_count, + released: false, + })) + } + _ => { + // Preserve the pre-AllowDirect behaviour: seed the cursor's cached + // metadata from the first result set (only non-empty sets are + // buffered, so `first` is the first non-empty COLMETADATA). + let metadata = result_sets.first().map(|rs| (*rs.columns).clone()); + Ok(CursorPrepExecOutcome::Cursor(prepared_cursor_from_outputs( + outputs, metadata, + )?)) + } + } +} + #[cfg(test)] mod tests { use super::*; - use crate::tds::codec::{BaseMetaDataColumn, FixedLenType, TokenReturnValue, TypeInfo}; + use crate::tds::codec::{ + BaseMetaDataColumn, FixedLenType, TokenReturnValue, TokenRow, TypeInfo, + }; #[test] fn fetch_encodes_next() { @@ -921,4 +1167,170 @@ mod tests { assert!(param.flags.is_empty()); assert!(matches!(param.value, ColumnData::I32(Some(77)))); } + + fn direct_set(result_index: usize, values: &[i32]) -> BufferedResultSet { + let columns = Arc::new(vec![Column::new("v".to_string(), crate::ColumnType::Int4)]); + let rows = values + .iter() + .map(|&v| { + let mut data = TokenRow::new(); + data.push(ColumnData::I32(Some(v))); + Row { + columns: columns.clone(), + data, + result_index, + } + }) + .collect(); + BufferedResultSet { columns, rows } + } + + #[test] + fn cursor_prep_exec_outcome_returns_cursor_for_nonzero_cursor() { + let outputs = vec![ + output("", 11), + output("", 22), + output("", CursorScrollOptions::ForwardOnly as i32), + output("", CursorConcurrencyOptions::ReadOnly as i32), + output("", 3), + ]; + + let outcome = cursor_prep_exec_outcome(&outputs, Vec::new()).unwrap(); + assert!(!outcome.is_direct()); + assert_eq!(outcome.prepared_handle().as_i32(), 11); + + let cursor = outcome.into_cursor().expect("expected cursor outcome"); + assert_eq!(cursor.prepared_handle().as_i32(), 11); + assert_eq!(cursor.cursor_handle().as_i32(), 22); + assert_eq!(cursor.row_count(), 3); + } + + #[test] + fn cursor_prep_exec_outcome_returns_direct_for_zero_cursor() { + let outputs = vec![ + output("", 11), + output("", 0), + output("", CursorScrollOptions::ForwardOnly as i32), + output("", CursorConcurrencyOptions::AllowDirect as i32), + output("", 3), + ]; + + let outcome = cursor_prep_exec_outcome(&outputs, vec![direct_set(0, &[1, 2, 3])]).unwrap(); + assert!(outcome.is_direct()); + assert_eq!(outcome.prepared_handle().as_i32(), 11); + + let direct = outcome.into_direct().expect("expected direct outcome"); + assert_eq!(direct.prepared_handle().as_i32(), 11); + assert_eq!(direct.row_count(), 3); + assert!(direct + .concurrency_options() + .contains(CursorConcurrencyOptions::AllowDirect)); + assert_eq!(direct.results().len(), 1); + assert_eq!(direct.results()[0].columns().len(), 1); + assert_eq!(direct.results()[0].columns()[0].name(), "v"); + assert_eq!(direct.results()[0].rows().len(), 3); + assert_eq!(direct.results()[0].rows()[0].get::(0), Some(1)); + assert_eq!(direct.results()[0].rows()[2].get::(0), Some(3)); + } + + #[test] + fn into_cursor_preserves_direct_results_on_mismatch() { + let outputs = vec![output("", 11), output("", 0)]; + let outcome = cursor_prep_exec_outcome(&outputs, vec![direct_set(0, &[7])]).unwrap(); + + let mut direct = outcome + .into_cursor() + .expect_err("expected intact direct results"); + assert_eq!(direct.prepared_handle().as_i32(), 11); + assert_eq!(direct.results().len(), 1); + assert_eq!(direct.results()[0].rows()[0].get::(0), Some(7)); + + // This is a synthetic unit-test handle, so suppress the intentional + // leak warning after proving the alternate variant survived intact. + direct.released = true; + } + + #[test] + fn into_direct_preserves_prepared_cursor_on_mismatch() { + let outputs = vec![output("", 11), output("", 22)]; + let outcome = cursor_prep_exec_outcome(&outputs, Vec::new()).unwrap(); + + let mut cursor = outcome + .into_direct() + .expect_err("expected intact prepared cursor"); + assert_eq!(cursor.prepared_handle().as_i32(), 11); + assert_eq!(cursor.cursor_handle().as_i32(), 22); + + // These are synthetic unit-test handles, so suppress the intentional + // leak warnings after proving the alternate variant survived intact. + cursor.released = true; + if let Some(inner) = cursor.cursor.as_mut() { + inner.closed = true; + } + } + + #[test] + fn cursor_prep_exec_outcome_preserves_multiple_direct_sets_in_order() { + let outputs = vec![ + output("", 11), + output("", 0), + output("", 0), + output("", 0), + output("", 5), + ]; + + let outcome = cursor_prep_exec_outcome( + &outputs, + vec![direct_set(0, &[1, 2, 3]), direct_set(1, &[4, 5])], + ) + .unwrap(); + let direct = outcome.into_direct().expect("expected direct outcome"); + + assert_eq!(direct.results().len(), 2); + assert_eq!(direct.results()[0].rows().len(), 3); + assert_eq!(direct.results()[0].rows()[0].get::(0), Some(1)); + assert_eq!(direct.results()[1].rows().len(), 2); + assert_eq!(direct.results()[1].rows()[0].get::(0), Some(4)); + assert_eq!(direct.results()[1].rows()[1].get::(0), Some(5)); + } + + #[test] + fn direct_result_set_into_rows_returns_owned_rows() { + let columns = Arc::new(vec![Column::new("v".to_string(), crate::ColumnType::Int4)]); + let mut data = TokenRow::new(); + data.push(ColumnData::I32(Some(9))); + let rows = vec![Row { + columns: columns.clone(), + data, + result_index: 0, + }]; + let rs = DirectResultSet { columns, rows }; + + assert_eq!(rs.columns().len(), 1); + assert_eq!(rs.rows().len(), 1); + + let owned = rs.into_rows(); + assert_eq!(owned.len(), 1); + assert_eq!(owned[0].get::(0), Some(9)); + } + + #[test] + fn cursor_prep_exec_outcome_errors_on_zero_prepared_handle() { + let outputs = vec![output("", 0), output("", 0)]; + let err = cursor_prep_exec_outcome(&outputs, Vec::new()).unwrap_err(); + assert!(matches!(err, crate::Error::Protocol(_))); + } + + #[test] + fn cursor_prep_exec_outcome_errors_on_missing_cursor() { + // A prepared handle but no @cursor anywhere falls through to the normal + // cursor path, which reports the missing-@cursor protocol error. + let outputs = vec![output("@prepared_handle", 11)]; + match cursor_prep_exec_outcome(&outputs, Vec::new()).unwrap_err() { + crate::Error::Protocol(msg) => { + assert!(msg.contains("@cursor"), "unexpected message: {msg}") + } + other => panic!("expected protocol error, got {:?}", other), + } + } } diff --git a/src/client/rpc_response.rs b/src/client/rpc_response.rs index 50f4f2a5..4db88863 100644 --- a/src/client/rpc_response.rs +++ b/src/client/rpc_response.rs @@ -13,7 +13,7 @@ use crate::tds::codec::{ TokenTabName, }; use crate::tds::stream::{ReceivedToken, TokenStream}; -use crate::{Column, FromSql, SqlReadBytes, TokenType}; +use crate::{Column, FromSql, Row, SqlReadBytes, TokenType}; use futures_util::io::{AsyncRead, AsyncWrite}; use futures_util::stream::{Stream, TryStreamExt}; use std::convert::TryFrom; @@ -182,6 +182,120 @@ where Ok((outputs, status, metadata)) } +/// A single result set buffered from an RPC response: its column metadata and +/// every row that belongs to it. +/// +/// Unlike [`collect_rpc_outputs_with_metadata`], which keeps only the first +/// COLMETADATA and discards row data, this retains the rows so callers such as +/// the AllowDirect `sp_cursorprepexec` path can own the streamed result sets. +#[derive(Debug)] +pub(crate) struct BufferedResultSet { + pub columns: Arc>, + pub rows: Vec, +} + +/// Drain an RPC response, buffering every result set (metadata + rows, in +/// order) alongside the `RETURNVALUE` outputs and the return status. +/// +/// This is the row-preserving counterpart of +/// [`collect_rpc_outputs_with_metadata`]: `sp_cursorprepexec` normally streams +/// only COLMETADATA on open (rows arrive later via `sp_cursorfetch`), but when +/// the server selects the AllowDirect fast path it streams the result sets +/// inline, and those rows must be kept. +pub(crate) async fn collect_rpc_result_sets( + conn: &mut Connection, +) -> crate::Result<(Vec, Vec, Option)> +where + S: AsyncRead + AsyncWrite + Unpin + Send, +{ + let ts = TokenStream::new(conn); + let stream = ts.try_unfold(); + collect_rpc_result_sets_from_stream(stream).await +} + +/// Lower-level variant that drains an arbitrary token stream, so the buffering +/// logic can be unit-tested with synthetic inputs (mirrors +/// [`collect_rpc_outputs_from_stream`]). +async fn collect_rpc_result_sets_from_stream( + mut stream: S, +) -> crate::Result<(Vec, Vec, Option)> +where + S: Stream> + Unpin, +{ + let mut results: Vec = Vec::new(); + let mut columns: Option>> = None; + let mut current: Vec = Vec::new(); + let mut outputs: Vec = Vec::new(); + let mut status = None; + let mut last_error: Option = None; + + // Flush the in-progress result set into `results`, but only if it carries + // column metadata — a result set is defined by its COLMETADATA, so an + // empty-COLMETADATA placeholder (0 columns) is dropped rather than recorded. + fn flush( + results: &mut Vec, + columns: &mut Option>>, + current: &mut Vec, + ) { + if let Some(cols) = columns.take() { + if !cols.is_empty() { + results.push(BufferedResultSet { + columns: cols, + rows: std::mem::take(current), + }); + } else { + current.clear(); + } + } + } + + while let Some(token) = stream.try_next().await? { + match token { + ReceivedToken::NewResultset(meta) => { + flush(&mut results, &mut columns, &mut current); + columns = Some(Arc::new(meta.columns().collect::>())); + } + ReceivedToken::Row(data) => { + if let Some(cols) = &columns { + current.push(Row { + columns: cols.clone(), + data, + // The index this set will occupy once flushed. + result_index: results.len(), + }); + } + } + ReceivedToken::ReturnValue(rv) => outputs.push(rv.into()), + ReceivedToken::ReturnStatus(s) => status = Some(s), + ReceivedToken::Error(e) => { + if last_error.is_none() { + last_error = Some(crate::Error::Server(e)); + } + } + ReceivedToken::DoneInProc(done) => { + // A DoneInProc without `More` closes one statement's result set + // inside the proc; more tokens (ReturnValue, DoneProc) follow. + if !done.status().contains(DoneStatus::More) { + flush(&mut results, &mut columns, &mut current); + } + } + ReceivedToken::DoneProc(done) | ReceivedToken::Done(done) => { + if !done.status().contains(DoneStatus::More) { + flush(&mut results, &mut columns, &mut current); + break; + } + } + _ => {} + } + } + + if let Some(err) = last_error { + return Err(err); + } + + Ok((results, outputs, status)) +} + /// Drain a metadata-only cursor-fetch RPC response without constructing a /// [`QueryStream`](crate::QueryStream). The metadata probe requests zero rows, /// so this walker returns as soon as it sees the first non-empty COLMETADATA @@ -294,7 +408,7 @@ where mod tests { use super::*; use crate::tds::codec::{ - BaseMetaDataColumn, FixedLenType, MetaDataColumn, TokenDone, TokenError, TypeInfo, + BaseMetaDataColumn, FixedLenType, MetaDataColumn, TokenDone, TokenError, TokenRow, TypeInfo, }; use enumflags2::BitFlags; use futures_util::stream::iter; @@ -514,4 +628,115 @@ mod tests { assert!(outputs.is_empty()); assert!(status.is_none()); } + + fn mk_row(value: i32) -> ReceivedToken { + let mut row = TokenRow::new(); + row.push(ColumnData::I32(Some(value))); + ReceivedToken::Row(row) + } + + #[tokio::test] + async fn collect_result_sets_buffers_single_set_with_rows_and_outputs() { + let s = synthetic(vec![ + ReceivedToken::NewResultset(mk_metadata("v")), + mk_row(1), + mk_row(2), + mk_row(3), + ReceivedToken::ReturnValue(mk_return_value( + "@prepared_handle", + 1, + ColumnData::I32(Some(11)), + )), + ReceivedToken::ReturnValue(mk_return_value("@cursor", 2, ColumnData::I32(Some(0)))), + ReceivedToken::ReturnStatus(0), + mk_done_proc_final(), + ]); + + let (results, outputs, status) = collect_rpc_result_sets_from_stream(s).await.unwrap(); + + assert_eq!(results.len(), 1); + assert_eq!(results[0].columns.len(), 1); + assert_eq!(results[0].columns[0].name(), "v"); + assert_eq!(results[0].rows.len(), 3); + assert_eq!(results[0].rows[0].get::(0), Some(1)); + assert_eq!(results[0].rows[2].get::(0), Some(3)); + assert_eq!(results[0].rows[0].result_index(), 0); + assert_eq!(outputs.len(), 2); + assert_eq!(status, Some(0)); + } + + #[tokio::test] + async fn collect_result_sets_preserves_multiple_sets_in_order() { + let s = synthetic(vec![ + ReceivedToken::NewResultset(mk_metadata("v")), + mk_row(1), + mk_row(2), + ReceivedToken::DoneInProc(TokenDone::with_rows(2)), + ReceivedToken::NewResultset(mk_metadata("v")), + mk_row(3), + ReceivedToken::ReturnStatus(0), + mk_done_proc_final(), + ]); + + let (results, _outputs, _status) = collect_rpc_result_sets_from_stream(s).await.unwrap(); + + assert_eq!(results.len(), 2); + assert_eq!(results[0].rows.len(), 2); + assert_eq!(results[0].rows[0].get::(0), Some(1)); + assert_eq!(results[0].rows[0].result_index(), 0); + assert_eq!(results[1].rows.len(), 1); + assert_eq!(results[1].rows[0].get::(0), Some(3)); + assert_eq!(results[1].rows[0].result_index(), 1); + } + + #[tokio::test] + async fn collect_result_sets_drops_empty_colmetadata_placeholder() { + let s = synthetic(vec![ + ReceivedToken::NewResultset(Arc::new(TokenColMetaData { + columns: Vec::new(), + })), + ReceivedToken::NewResultset(mk_metadata("v")), + mk_row(1), + mk_done_proc_final(), + ]); + + let (results, _outputs, _status) = collect_rpc_result_sets_from_stream(s).await.unwrap(); + + assert_eq!(results.len(), 1); + assert_eq!(results[0].rows[0].get::(0), Some(1)); + assert_eq!(results[0].rows[0].result_index(), 0); + } + + #[tokio::test] + async fn collect_result_sets_surfaces_first_error() { + let err = TokenError::new(50000, 1, 16, "boom", "srv", "proc", 1); + let s = synthetic(vec![ + ReceivedToken::NewResultset(mk_metadata("v")), + mk_row(1), + ReceivedToken::Error(err), + ReceivedToken::ReturnStatus(0), + mk_done_proc_final(), + ]); + + match collect_rpc_result_sets_from_stream(s).await { + Err(crate::Error::Server(te)) => assert_eq!(te.code, 50000), + other => panic!("expected Server error, got {:?}", other), + } + } + + #[tokio::test] + async fn collect_result_sets_no_result_sets_outputs_only() { + // The normal (non-AllowDirect) shape: no inline rows, just outputs. + let s = synthetic(vec![ + ReceivedToken::ReturnValue(mk_return_value("@handle", 1, ColumnData::I32(Some(7)))), + ReceivedToken::ReturnStatus(0), + mk_done_proc_final(), + ]); + + let (results, outputs, status) = collect_rpc_result_sets_from_stream(s).await.unwrap(); + + assert!(results.is_empty()); + assert_eq!(outputs.len(), 1); + assert_eq!(status, Some(0)); + } } diff --git a/src/lib.rs b/src/lib.rs index 644208cf..28ea47f9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -271,8 +271,8 @@ pub mod server; pub use client::{ AuthMethod, CancellationToken, Client, Config, Cursor, CursorConcurrencyOptions, CursorHandle, - CursorOpenOptions, CursorScrollOptions, Fetch, OutputValue, PreparedCursor, PreparedHandle, - PreparedStatement, + CursorOpenOptions, CursorPrepExecOutcome, CursorScrollOptions, DirectResultSet, DirectResults, + Fetch, OutputValue, PreparedCursor, PreparedHandle, PreparedStatement, }; pub(crate) use error::Error; pub use from_sql::{FromSql, FromSqlOwned}; diff --git a/tests/query.rs b/tests/query.rs index 59ecac68..3907c2e6 100644 --- a/tests/query.rs +++ b/tests/query.rs @@ -8,8 +8,9 @@ use std::sync::Once; use tiberius::FromSql; use tiberius::{ - numeric::Numeric, xml::XmlData, ColumnData, ColumnFlag, ColumnType, Query, QueryItem, Result, - TvpColumn, TvpData, TypeInfo, UdtData, VarLenContext, VarLenType, VariantData, + numeric::Numeric, xml::XmlData, ColumnData, ColumnFlag, ColumnType, CursorConcurrencyOptions, + CursorOpenOptions, CursorScrollOptions, Query, QueryItem, Result, TvpColumn, TvpData, TypeInfo, + UdtData, VarLenContext, VarLenType, VariantData, }; use uuid::Uuid; @@ -565,6 +566,61 @@ where Ok(()) } +#[test_on_runtimes] +async fn cursor_prep_exec_allow_direct_returns_owned_results( + mut conn: tiberius::Client, +) -> Result<()> +where + S: AsyncRead + AsyncWrite + Unpin + Send, +{ + let outcome = conn + .cursor_prep_exec( + "SELECT CAST(101 AS int) AS first_value; \ + SELECT CAST(202 AS int) AS second_value;", + CursorOpenOptions::new( + CursorScrollOptions::ForwardOnly, + CursorConcurrencyOptions::ReadOnly | CursorConcurrencyOptions::AllowDirect, + ), + "", + &[], + ) + .await?; + + let direct = match outcome.into_direct() { + Ok(direct) => direct, + Err(cursor) => { + cursor.close_and_unprepare(&mut conn).await?; + panic!("SQL Server opened a cursor instead of taking AllowDirect"); + } + }; + + assert_ne!(direct.prepared_handle().as_i32(), 0); + assert_eq!(direct.results().len(), 2); + assert_eq!(direct.results()[0].columns()[0].name(), "first_value"); + assert_eq!(direct.results()[0].rows().len(), 1); + assert_eq!(direct.results()[0].rows()[0].result_index(), 0); + assert_eq!(direct.results()[0].rows()[0].get::(0), Some(101)); + assert_eq!(direct.results()[1].columns()[0].name(), "second_value"); + assert_eq!(direct.results()[1].rows().len(), 1); + assert_eq!(direct.results()[1].rows()[0].result_index(), 1); + assert_eq!(direct.results()[1].rows()[0].get::(0), Some(202)); + + let owned = direct.unprepare(&mut conn).await?; + assert_eq!(owned.len(), 2); + + // A successful follow-up request proves the direct RPC response and + // CursorUnprepare response were both fully drained. + let row = conn + .query("SELECT CAST(303 AS int)", &[]) + .await? + .into_row() + .await? + .unwrap(); + assert_eq!(row.get::(0), Some(303)); + + Ok(()) +} + #[test_on_runtimes] async fn bool_type(mut conn: tiberius::Client) -> Result<()> where diff --git a/tests/rpc_helpers.rs b/tests/rpc_helpers.rs index a51a0bf2..abf10fa3 100644 --- a/tests/rpc_helpers.rs +++ b/tests/rpc_helpers.rs @@ -166,6 +166,10 @@ struct SharedState { cursor_fetch_log: Mutex>, cursorprepexec_param_defs_log: Mutex>>, cursorprepexec_send_metadata: Mutex, + /// When set, `sp_cursorprepexec` takes the AllowDirect fast path: it + /// prepares the statement but returns `@cursor == 0` and streams the + /// result sets inline instead of opening a cursor. + cursorprepexec_allow_direct: Mutex, /// When set, a metadata-only (`n_rows == 0`) `sp_cursorfetch` emits no /// tokens and instead polls for an attention, simulating a stalled probe. /// Used to exercise cancellation of the metadata-fetch read path. @@ -182,6 +186,7 @@ impl SharedState { cursor_fetch_log: Mutex::new(Vec::new()), cursorprepexec_param_defs_log: Mutex::new(Vec::new()), cursorprepexec_send_metadata: Mutex::new(false), + cursorprepexec_allow_direct: Mutex::new(false), stall_metadata_fetch: Mutex::new(false), rpc_log: Mutex::new(Vec::new()), } @@ -705,6 +710,39 @@ impl RpcHandler for SpecialRpc { ColumnData::I32(Some(v)) => *v, _ => 0, }; + + if *self.0.cursorprepexec_allow_direct.lock().unwrap() { + // AllowDirect: prepare the statement but do NOT open a + // cursor. Stream the result sets inline and return + // @cursor == 0. Emit two sets to exercise ordering. + write_int_result_set(client, &[1, 2, 3], FinalDone::InProc).await?; + write_int_result_set(client, &[4, 5, 6], FinalDone::InProc).await?; + + let outputs = vec![ + OutputParameter::from_input( + &all[0], + ColumnData::I32(Some(prepared_handle.as_i32())), + ) + .with_ordinal(1), + OutputParameter::from_input(&all[1], ColumnData::I32(Some(0))) + .with_ordinal(2), + OutputParameter::from_input(&all[4], ColumnData::I32(Some(scrollopt))) + .with_ordinal(5), + OutputParameter::from_input(&all[5], ColumnData::I32(Some(ccopt))) + .with_ordinal(6), + OutputParameter::from_input(&all[6], ColumnData::I32(Some(6))) + .with_ordinal(7), + ]; + send_output_params(client, outputs).await?; + send_return_status(client, 0).await?; + client + .send(TdsBackendMessage::Token(BackendToken::DoneProc( + TokenDone::with_rows(0), + ))) + .await?; + return Ok(()); + } + let cursor_entry = CursorEntry::new(sql, scrollopt, ccopt, rows.len() as i32); let cursor_handle = self.0.cursors.lock().unwrap().open(cursor_entry); self.0 @@ -1102,7 +1140,9 @@ fn cursor_prep_exec_fetch_close_unprepare() { &[], ) .await - .unwrap(); + .unwrap() + .into_cursor() + .expect("expected prepared cursor"); assert_ne!(cursor.prepared_handle().as_i32(), 0); assert_ne!(cursor.cursor_handle().as_i32(), 0); assert_eq!(cursor.row_count(), 3); @@ -1158,6 +1198,73 @@ fn cursor_prep_exec_fetch_close_unprepare() { }); } +#[test] +fn cursor_prep_exec_allow_direct_returns_owned_results() { + smol::block_on(async { + with_server(|addr, state| async move { + *state.cursorprepexec_allow_direct.lock().unwrap() = true; + let mut client = connect_client(addr).await.unwrap(); + + let outcome = client + .cursor_prep_exec( + "SELECT 1 AS v UNION ALL SELECT 2 AS v UNION ALL SELECT 3 AS v", + CursorOpenOptions::new( + CursorScrollOptions::ForwardOnly, + tiberius::CursorConcurrencyOptions::ReadOnly + | tiberius::CursorConcurrencyOptions::AllowDirect, + ), + "", + &[], + ) + .await + .unwrap(); + + // The server took the AllowDirect fast path: no cursor, results + // streamed inline. + assert!(outcome.is_direct()); + assert_ne!(outcome.prepared_handle().as_i32(), 0); + + let direct = outcome.into_direct().expect("expected AllowDirect results"); + assert!(direct + .concurrency_options() + .contains(tiberius::CursorConcurrencyOptions::AllowDirect)); + + // Both result sets are preserved, in order, with metadata and rows. + assert_eq!(direct.results().len(), 2); + assert_eq!(direct.results()[0].columns().len(), 1); + assert_eq!(direct.results()[0].columns()[0].name(), "v"); + + let first: Vec = direct.results()[0] + .rows() + .iter() + .map(|r| r.get::(0).unwrap()) + .collect(); + assert_eq!(first, vec![1, 2, 3]); + + let second: Vec = direct.results()[1] + .rows() + .iter() + .map(|r| r.get::(0).unwrap()) + .collect(); + assert_eq!(second, vec![4, 5, 6]); + + // Cleanup releases the prepared handle and hands back the owned + // result sets — no cursor close is issued. + let owned = direct.unprepare(&mut client).await.unwrap(); + assert_eq!(owned.len(), 2); + assert_eq!(owned[0].rows().len(), 3); + assert_eq!(owned[1].rows().len(), 3); + + let seen_rpcs = state.rpc_log.lock().unwrap().clone(); + assert_eq!( + seen_rpcs, + vec![RpcProcId::CursorPrepExec, RpcProcId::CursorUnprepare] + ); + }) + .await; + }); +} + #[test] fn cursor_prep_exec_fetch_metadata_close_unprepare() { smol::block_on(async { @@ -1172,7 +1279,9 @@ fn cursor_prep_exec_fetch_metadata_close_unprepare() { &[], ) .await - .unwrap(); + .unwrap() + .into_cursor() + .expect("expected prepared cursor"); let columns = cursor.fetch_metadata(&mut client).await.unwrap(); assert_eq!(columns.len(), 1); @@ -1223,7 +1332,9 @@ fn cursor_prep_exec_fetch_metadata_reuses_initial_metadata() { &[], ) .await - .unwrap(); + .unwrap() + .into_cursor() + .expect("expected prepared cursor"); let columns = cursor.fetch_metadata(&mut client).await.unwrap(); assert_eq!(columns.len(), 1); @@ -1437,7 +1548,7 @@ fn cancel_interrupts_packetless_cursor_metadata_fetch() { *state.stall_metadata_fetch.lock().unwrap() = true; let mut client = connect_client(addr).await.unwrap(); - let mut cursor = client + let cursor = client .cursor_prep_exec( "SELECT 1 AS v UNION ALL SELECT 2 AS v UNION ALL SELECT 3 AS v", CursorOpenOptions::default(), @@ -1445,7 +1556,9 @@ fn cancel_interrupts_packetless_cursor_metadata_fetch() { &[], ) .await - .unwrap(); + .unwrap() + .into_cursor() + .expect("expected prepared cursor"); let token = client.cancellation_token(); let canceller = smol::spawn(async move {