Rewire mssql-odbc fetch hot path onto the reactor-free sync core (L5)

mssql-mock-tds/src/protocol.rs

@@ -674,13 +674,20 @@ pub fn build_login_ack() -> BytesMut {
 
 /// Build a DONE token
 pub fn build_done_token(row_count: u64) -> BytesMut {
+    build_done_token_with_status(0x0000, row_count)
+}
+
+/// DONE token with an explicit status word. `0x0000` is `DONE_FINAL` (the batch
+/// ends); `0x0001` is `DONE_MORE`, signalling that another result set follows in
+/// the same batch (what a multi-statement `SELECT …; SELECT …` produces).
+pub fn build_done_token_with_status(status: u16, row_count: u64) -> BytesMut {
     let mut token_data = BytesMut::new();
 
     // DONE token (0xFD)
     token_data.put_u8(TokenType::Done as u8);
 
-    // Status: DONE_FINAL (0x00) - little-endian
-    token_data.put_u16_le(0x0000);
+    // Status - little-endian
+    token_data.put_u16_le(status);
 
     // CurCmd: SELECT (0xC1) - little-endian
     token_data.put_u16_le(0x00C1);
@@ -822,6 +829,41 @@ pub fn build_info_token(info: &crate::query_response::InfoMessage) -> BytesMut {
 pub fn build_query_result(response: &crate::query_response::QueryResponse) -> BytesMut {
     let mut result = BytesMut::new();
 
+    // Flatten the primary set and any additional sets (recursively) into batch
+    // order. All but the last close with DONE_MORE so the client advances across
+    // the boundary; the last closes with DONE_FINAL. Single-set responses (the
+    // common case) have no additional sets and serialize byte-identically.
+    let mut sets = Vec::new();
+    collect_result_sets(response, &mut sets);
+    let last = sets.len() - 1;
+    for (idx, set) in sets.into_iter().enumerate() {
+        serialize_result_set(&mut result, set, idx == last);
+    }
+
+    wrap_in_packet(PacketType::TabularResult, result)
+}
+
+/// Flatten a response and its `additional_sets` (depth-first, in order) into a
+/// linear batch of result sets so their terminal DONE statuses can be assigned.
+fn collect_result_sets<'a>(
+    response: &'a crate::query_response::QueryResponse,
+    out: &mut Vec<&'a crate::query_response::QueryResponse>,
+) {
+    out.push(response);
+    for set in &response.additional_sets {
+        collect_result_sets(set, out);
+    }
+}
+
+/// Serialize one result set's tokens (ColMetadata, injected Info, rows, terminal
+/// DONE) into `result`. `is_last` selects the DONE status: the final set in a
+/// batch closes `DONE_FINAL`, earlier sets `DONE_MORE`. An `error_after` set is
+/// always terminal (its ERROR + DONE end the batch), so it ignores `is_last`.
+fn serialize_result_set(
+    result: &mut BytesMut,
+    response: &crate::query_response::QueryResponse,
+    is_last: bool,
+) {
     // ColMetadata token (0x81)
     result.put_u8(TokenType::ColMetadata as u8);
     result.put_u16_le(response.columns.len() as u16); // Column count
@@ -860,7 +902,7 @@ pub fn build_query_result(response: &crate::query_response::QueryResponse) -> By
             for row in response.rows.iter().take(err.after_rows) {
                 result.put_u8(TokenType::Row as u8);
                 for value in &row.values {
-                    value.write_to_buffer(&mut result);
+                    value.write_to_buffer(result);
                 }
             }
             result.extend_from_slice(&build_error_token(
@@ -880,16 +922,18 @@ pub fn build_query_result(response: &crate::query_response::QueryResponse) -> By
             for row in &response.rows {
                 result.put_u8(TokenType::Row as u8);
                 for value in &row.values {
-                    value.write_to_buffer(&mut result);
+                    value.write_to_buffer(result);
                 }
             }
 
-            // DONE token
-            result.extend_from_slice(&build_done_token(response.rows.len() as u64));
+            // DONE_MORE when another set follows, else terminal DONE_FINAL.
+            let status = if is_last { 0x0000 } else { 0x0001 };
+            result.extend_from_slice(&build_done_token_with_status(
+                status,
+                response.rows.len() as u64,
+            ));
         }
     }
-
-    wrap_in_packet(PacketType::TabularResult, result)
 }
 
 /// Build a bare ERROR token (0xAA) with no surrounding DONE or packet framing,

mssql-mock-tds/src/query_response.rs

@@ -187,6 +187,13 @@ pub struct QueryResponse {
     /// When set, only the first `after_rows` rows are streamed, then an ERROR
     /// token and a terminal DONE are emitted (no trailing rows).
     pub error_after: Option<MidStreamError>,
+    /// Further result sets streamed in the same batch after this one. When
+    /// non-empty, this set's terminal DONE carries `DONE_MORE` and each
+    /// subsequent set is emitted in turn (the last one closing with a terminal
+    /// `DONE_FINAL`) — the multi-result-set shape a `SELECT …; SELECT …` batch
+    /// produces, exercising `SQLMoreResults`/`advance()` across the boundary.
+    /// Empty by default, so single-set responses serialize byte-identically.
+    pub additional_sets: Vec<QueryResponse>,
 }
 
 impl QueryResponse {
@@ -197,6 +204,7 @@ impl QueryResponse {
             rows,
             info_tokens: Vec::new(),
             error_after: None,
+            additional_sets: Vec::new(),
         }
     }
 
@@ -211,13 +219,22 @@ impl QueryResponse {
         self
     }
 
+    /// Append another result set to this batch. The current terminal DONE
+    /// becomes `DONE_MORE`; `next` is streamed after it (recursively, if it too
+    /// carries additional sets), with the final set closing on `DONE_FINAL`.
+    pub fn with_additional_result_set(mut self, next: QueryResponse) -> Self {
+        self.additional_sets.push(next);
+        self
+    }
+
     /// Helper to create a response for SELECT 1
     pub fn select_one() -> Self {
         Self {
             columns: vec![ColumnDefinition::new("", SqlDataType::Int)],
             rows: vec![Row::new(vec![ColumnValue::Int(1)])],
             info_tokens: Vec::new(),
             error_after: None,
+            additional_sets: Vec::new(),
         }
     }
 
@@ -236,6 +253,7 @@ impl QueryResponse {
             ])],
             info_tokens: Vec::new(),
             error_after: None,
+            additional_sets: Vec::new(),
         }
     }
 }

mssql-odbc/Cargo.toml

@@ -46,5 +46,9 @@ windows = { version = "0.58", features = [
 # `TdsClient` driven by scripted TDS tokens (see `mssql_tds::test_client_support`)
 # to cover client-driven ODBC paths without a live server.
 mssql-tds = { path = "../mssql-tds", features = ["test-util"] }
+# A real raw-TCP mock TDS peer for the sync-fetch integration tests: scripted
+# token clients are never `into_sync`-eligible (their transport yields no
+# blocking parts), so the sync arm can only be exercised over an actual socket.
+mssql-mock-tds = { path = "../mssql-mock-tds" }
 
 [build-dependencies]

mssql-odbc/src/api/close_cursor.rs

@@ -131,7 +131,7 @@ fn sql_free_stmt_close_safe(statement_handle: SqlHandle, stmt: &StmtHandle) -> S
 /// Resets cursor state on the statement (cursor is no longer open, metadata cleared).
 pub(super) fn reset_cursor_state(stmt_state: &mut crate::handles::stmt::StmtState) {
     stmt_state.clear_state(STMT_STATE_CURSOR_OPEN | STMT_STATE_EXEC_CONTEXT);
-    stmt_state.current_row = None;
+    stmt_state.reset_fetch_state();
     stmt_state.column_metadata.clear();
     stmt_state.pending_row_counts.clear();
 }
@@ -167,7 +167,7 @@ pub(super) fn drain_and_release(stmt: &StmtHandle, statement_handle: SqlHandle)
         dbc_state.client.take()
     };
 
-    let Some(mut client) = client else {
+    let Some(client) = client else {
         error!("drain_and_release: no TDS client to drain — this is a bug");
         if let Ok(mut ds) = dbc.inner.lock()
             && ds.active_stmt == Some(statement_handle)
@@ -177,6 +177,26 @@ pub(super) fn drain_and_release(stmt: &StmtHandle, statement_handle: SqlHandle)
         return DrainOutcome::Failed;
     };
 
+    // Cursor close is a control-plane drain: revert a sync fetch cursor to the
+    // async edge first (rule C) so `close_query` runs on the reactor and the
+    // `sp_prepexec` `@handle` capture below works. A no-op when already async; a
+    // revert failure poisons the connection.
+    let mut client = match client.into_async() {
+        Ok(client) => client,
+        Err(e) => {
+            error!(%e, "drain_and_release: reverting sync cursor to async failed — connection lost");
+            if let Ok(mut stmt_state) = stmt.inner.lock() {
+                post_tds_error(&mut stmt_state, &e, SQLSTATE_HY000);
+            }
+            if let Ok(mut ds) = dbc.inner.lock()
+                && ds.active_stmt == Some(statement_handle)
+            {
+                ds.active_stmt = None;
+            }
+            return DrainOutcome::Failed;
+        }
+    };
+
     if let Err(e) = dbc.runtime.block_on(client.close_query()) {
         error!(%e, "drain_and_release: failed to drain TDS stream — connection may be broken");
         // Surface the failure as a diagnostic so the app is not told the close
@@ -185,7 +205,7 @@ pub(super) fn drain_and_release(stmt: &StmtHandle, statement_handle: SqlHandle)
             post_tds_error(&mut stmt_state, &e, SQLSTATE_HY000);
         }
         if let Ok(mut ds) = dbc.inner.lock() {
-            ds.client = Some(client);
+            ds.store_async(client);
             if ds.active_stmt == Some(statement_handle) {
                 ds.active_stmt = None;
             }
@@ -203,7 +223,7 @@ pub(super) fn drain_and_release(stmt: &StmtHandle, statement_handle: SqlHandle)
         Err(_) => {
             error!("drain_and_release: stmt mutex poisoned while posting info messages");
             if let Ok(mut dbc_state) = dbc.inner.lock() {
-                dbc_state.client = Some(client);
+                dbc_state.store_async(client);
                 if dbc_state.active_stmt == Some(statement_handle) {
                     dbc_state.active_stmt = None;
                 }
@@ -215,7 +235,7 @@ pub(super) fn drain_and_release(stmt: &StmtHandle, statement_handle: SqlHandle)
     // Drain complete: return client and release busy claim atomically.
     super::exec_common::capture_prepared_handle(stmt, &mut client);
     if let Ok(mut dbc_state) = dbc.inner.lock() {
-        dbc_state.client = Some(client);
+        dbc_state.store_async(client);
         if dbc_state.active_stmt == Some(statement_handle) {
             dbc_state.active_stmt = None;
         }

mssql-odbc/src/api/driver_connect.rs

@@ -338,7 +338,11 @@ fn do_connect(
 
     let has_server_info = post_tds_info_messages(state, &info_messages);
 
-    state.client = Some(client);
+    // Cache the negotiated server version so SQLGetInfo(SQL_DBMS_VER) never has
+    // to touch the live client — it stays reportable even while a sync fetch
+    // cursor owns the connection.
+    state.server_version = client.server_version();
+    state.store_async(client);
     state.connection_state = ConnectionState::Connected;
     debug!("SQLDriverConnectW: connected successfully");