diff --git a/mssql-py-core/src/connection.rs b/mssql-py-core/src/connection.rs index eefaaff2..67c6a00b 100644 --- a/mssql-py-core/src/connection.rs +++ b/mssql-py-core/src/connection.rs @@ -6,14 +6,13 @@ use pyo3::prelude::*; use pyo3::types::PyDict; use std::{path::PathBuf, sync::Arc}; use tokio::runtime::Runtime; -use tokio::sync::Mutex; use crate::odbc_auth::odbc_authentication_transformer::transform_auth; use crate::odbc_auth::odbc_authentication_validator::validate_auth; +use crate::pyclient::{self, SharedClient}; use crate::python_logger_adapter::scoped_tracing_bridge; use mssql_tds::{ connection::client_context::{ClientContext, IPAddressPreference}, - connection::tds_client::TdsClient, connection_provider::tds_connection_provider::TdsConnectionProvider, core::{EncryptionOptions, EncryptionSetting}, message::login_options::ApplicationIntent, @@ -26,7 +25,7 @@ const DEFAULT_LIBRARY_NAME: &str = "MS-PYTHON"; pub struct PyCoreConnection { #[allow(dead_code)] // Used for async operations in cursor execute runtime: Runtime, - tds_client: Option>>, + tds_client: Option, is_closed: bool, } @@ -82,7 +81,7 @@ impl PyCoreConnection { tracing::info!("Successfully connected to SQL Server"); Ok(PyCoreConnection { runtime, - tds_client: Some(Arc::new(Mutex::new(client))), + tds_client: Some(pyclient::new_shared(client)), is_closed: false, }) } @@ -97,14 +96,13 @@ impl PyCoreConnection { fn close(&mut self) -> PyResult<()> { if !self.is_closed { - // Send TDS close to the server and shut down the TCP connection - if let Some(client) = self.tds_client.take() { - self.runtime.block_on(async { - let mut guard = client.lock().await; - if let Err(e) = guard.close_connection().await { - tracing::warn!("Error closing connection: {}", e); - } - }); + // Send TDS close to the server and shut down the TCP connection. + // Revert any live sync edge to async first (control-plane op). + if let Some(cell) = self.tds_client.take() { + let handle = self.runtime.handle().clone(); + if let Err(e) = pyclient::close_connection(&cell, &handle) { + tracing::warn!("Error closing connection: {}", e); + } } self.is_closed = true; } @@ -125,6 +123,22 @@ impl PyCoreConnection { } } + fn sync_cursor(&self) -> PyResult { + if self.is_closed { + return Err(PyRuntimeError::new_err("Connection is closed")); + } + + if let Some(client) = &self.tds_client { + let handle = self.runtime.handle().clone(); + Ok(crate::sync_cursor::PyCoreSyncCursor::new( + client.clone(), + handle, + )) + } else { + Err(PyRuntimeError::new_err("No active connection")) + } + } + fn commit(&mut self) -> PyResult<()> { if self.is_closed { return Err(PyRuntimeError::new_err("Connection is closed")); diff --git a/mssql-py-core/src/cursor.rs b/mssql-py-core/src/cursor.rs index e969fd7d..1ad11bba 100644 --- a/mssql-py-core/src/cursor.rs +++ b/mssql-py-core/src/cursor.rs @@ -4,7 +4,7 @@ use mssql_tds::connection::bulk_copy::{ BulkCopy, ColumnMapping as TdsColumnMapping, ColumnMappingSource, }; -use mssql_tds::connection::tds_client::{ExecuteOptions, ResultSet, StatementResult, TdsClient}; +use mssql_tds::connection::tds_client::{ExecuteOptions, ResultSet, StatementResult}; use mssql_tds::datatypes::column_values::ColumnValues; use mssql_tds::datatypes::sqldatatypes::VectorBaseType; use pyo3::prelude::*; @@ -12,11 +12,11 @@ use pyo3::types::{PyDict, PyIterator, PyList, PyTuple}; use std::sync::Arc; use std::time::Duration; use tokio::runtime::Handle; -use tokio::sync::Mutex; use tracing::{error, info}; use crate::arrow_bulkcopy::{ArrowBatchRowAdapter, ColumnPlan, build_column_plans}; use crate::bulkcopy::PythonRowAdapter; +use crate::pyclient::{self, SharedClient}; use crate::python_logger_adapter::scoped_tracing_bridge; use crate::utils::convert_tds_error; use arrow::array::{RecordBatch, StructArray}; @@ -26,15 +26,20 @@ use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema, from_ffi}; /// Python Cursor class for Core TDS backend #[pyclass] pub struct PyCoreCursor { - tds_client: Arc>, + tds_client: SharedClient, runtime_handle: Handle, has_resultset: bool, + rowcount: i64, } #[pymethods] impl PyCoreCursor { #[pyo3(signature = (query, params=None))] #[allow(unused_variables)] + // The std MutexGuard is intentionally held across the `block_on` future's + // awaits: `block_on` drives that future to completion on this one thread, so + // the `!Send` guard never crosses threads and cannot deadlock the cell. + #[allow(clippy::await_holding_lock)] fn execute( &mut self, py: Python, @@ -47,9 +52,13 @@ impl PyCoreCursor { let runtime_handle = self.runtime_handle.clone(); // Execute query asynchronously - py.detach(|| { + let rowcount = py.detach(|| { runtime_handle.block_on(async { - let mut client = tds_client.lock().await; + let mut guard = tds_client.lock().map_err(|_| { + pyo3::exceptions::PyRuntimeError::new_err("TDS client mutex was poisoned") + })?; + // Revert any live sync edge before the control-plane execute. + let client = pyclient::ensure_async(&mut guard)?; info!("execute: Locked TDS client"); // Close any open batch before executing a new query @@ -98,14 +107,19 @@ impl PyCoreCursor { } info!("execute: Query executed successfully"); - Ok::<_, PyErr>(()) + Ok::<_, PyErr>(client.last_rows_affected()) }) })?; self.has_resultset = true; + self.rowcount = rowcount; Ok(()) } + // The std MutexGuard is held across the `block_on` future's awaits by design; + // `block_on` completes the future on this single thread, so the `!Send` guard + // never crosses threads. + #[allow(clippy::await_holding_lock)] fn fetchone(&mut self, py: Python) -> PyResult>> { if !self.has_resultset { return Ok(None); @@ -119,7 +133,10 @@ impl PyCoreCursor { // Fetch one row via next_row_into → PyRowWriter (bypasses RowToken) let result = py.detach(|| { runtime_handle.block_on(async { - let mut client = tds_client.lock().await; + let mut guard = tds_client.lock().map_err(|_| { + pyo3::exceptions::PyRuntimeError::new_err("TDS client mutex was poisoned") + })?; + let client = pyclient::ensure_async(&mut guard)?; info!("fetchone: Locked TDS client"); if client.on_rows() { @@ -197,6 +214,11 @@ impl PyCoreCursor { Ok(results) } + #[getter] + fn rowcount(&self) -> i64 { + self.rowcount + } + fn close(&mut self) -> PyResult<()> { // TODO: Might need to drain the results. self.has_resultset = false; @@ -275,6 +297,9 @@ impl PyCoreCursor { /// ``` #[pyo3(signature = (table_name, data_source, batch_size=0, timeout=30, column_mappings=None, keep_identity=false, check_constraints=false, table_lock=false, keep_nulls=false, fire_triggers=false, use_internal_transaction=false, python_logger=None))] #[allow(clippy::too_many_arguments)] + // Std MutexGuard held across the `block_on` future's awaits by design (single + // thread, `!Send` guard never crosses threads). + #[allow(clippy::await_holding_lock)] fn bulkcopy( &mut self, py: Python, @@ -331,13 +356,16 @@ impl PyCoreCursor { let runtime_handle = self.runtime_handle.clone(); let result = runtime_handle.block_on(async { info!("bulkcopy: Inside async block, attempting to lock TDS client"); - // Lock the TDS client - let mut client = tds_client.lock().await; + // Lock the TDS client (revert any sync edge before control-plane op) + let mut guard = tds_client.lock().map_err(|_| { + pyo3::exceptions::PyRuntimeError::new_err("TDS client mutex was poisoned") + })?; + let client = pyclient::ensure_async(&mut guard)?; info!("bulkcopy: Successfully locked TDS client"); // Create BulkCopy instance info!("bulkcopy: Creating BulkCopy instance"); - let mut bulk_copy = BulkCopy::new(&mut client, table_name) + let mut bulk_copy = BulkCopy::new(&mut *client, table_name) .batch_size(options.batch_size) .timeout(options.timeout) .check_constraints(options.check_constraints) @@ -529,6 +557,9 @@ impl PyCoreCursor { /// Returns the same `dict` shape as `bulkcopy`. Errors mirror `bulkcopy`'s. #[pyo3(signature = (table_name, source, batch_size=0, timeout=30, column_mappings=None, keep_identity=false, check_constraints=false, table_lock=false, keep_nulls=false, fire_triggers=false, use_internal_transaction=false, python_logger=None))] #[allow(clippy::too_many_arguments)] + // Std MutexGuard held across the `block_on` future's awaits by design (single + // thread, `!Send` guard never crosses threads). + #[allow(clippy::await_holding_lock)] fn bulkcopy_arrow( &mut self, py: Python, @@ -578,9 +609,12 @@ impl PyCoreCursor { // lets other Python threads run during the (potentially long) transfer. let result = py.detach(|| { runtime_handle.block_on(async { - let mut client = tds_client.lock().await; + let mut guard = tds_client.lock().map_err(|_| { + pyo3::exceptions::PyRuntimeError::new_err("TDS client mutex was poisoned") + })?; + let client = pyclient::ensure_async(&mut guard)?; - let mut bulk_copy = BulkCopy::new(&mut client, table_name) + let mut bulk_copy = BulkCopy::new(&mut *client, table_name) .batch_size(options.batch_size) .timeout(options.timeout) .check_constraints(options.check_constraints) @@ -1222,11 +1256,12 @@ impl PyCoreCursor { } impl PyCoreCursor { - pub fn new(tds_client: Arc>, runtime_handle: Handle) -> Self { + pub(crate) fn new(tds_client: SharedClient, runtime_handle: Handle) -> Self { Self { tds_client, runtime_handle, has_resultset: false, + rowcount: -1, } } diff --git a/mssql-py-core/src/lib.rs b/mssql-py-core/src/lib.rs index 03163659..e6867c47 100644 --- a/mssql-py-core/src/lib.rs +++ b/mssql-py-core/src/lib.rs @@ -11,9 +11,11 @@ mod bulkcopy; mod connection; mod cursor; mod odbc_auth; +mod pyclient; mod python_entra_token_factory; mod python_logger_adapter; mod row_writer; +mod sync_cursor; mod tracing_init; mod types; mod utils; @@ -83,6 +85,7 @@ fn mssql_py_core(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; + m.add_class::()?; // Test-only hook to drive PythonEntraIdTokenFactory::create_token from // Python tests. Underscore-prefixed to mark as internal/test-only. diff --git a/mssql-py-core/src/pyclient.rs b/mssql-py-core/src/pyclient.rs new file mode 100644 index 00000000..d12c4064 --- /dev/null +++ b/mssql-py-core/src/pyclient.rs @@ -0,0 +1,243 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Shared TDS client cell backing the Python cursors. +//! +//! The connection owns one [`SharedClient`] — an `Arc>` holding +//! the connection in one of two interchangeable edges. The async [`TdsClient`] +//! drives every control-plane op (connect, execute, COLMETADATA, advance, +//! close, bulk copy). A sync cursor flips the cell to the reactor-free +//! [`TdsSyncClient`] for its row-pull hot loop and reverts before the next +//! control-plane op; the async cursor never flips and always uses the async arm. +//! +//! The flip runs IN PLACE under the lock via [`std::mem::replace`], so the `Arc` +//! clones each cursor holds keep pointing at the same cell. `into_sync`/ +//! `into_async` (which consume the client by value) operate on the value moved +//! OUT of the enum, never on the `Arc`, so `Arc::try_unwrap` is never needed — +//! the L6 structural blocker (live cursor clones) is sidestepped entirely. +//! +//! A [`std::sync::Mutex`] (not tokio) backs the cell so the sync arm acquires +//! the lock with no `.await`, keeping its fetch path reactor-free. Async-arm ops +//! still wrap their `.await` work in `block_on` while holding the guard on the +//! current thread; `block_on` runs to completion on one thread, so the `!Send` +//! guard held across it is sound. + +use std::sync::{Arc, Mutex}; + +use mssql_tds::connection::tds_client::{ExecuteOptions, ResultSet, StatementResult, TdsClient}; +use mssql_tds::connection::tds_sync_client::{SyncConversion, TdsSyncClient}; +use mssql_tds::datatypes::row_writer::RowWriter; +use pyo3::PyErr; +use pyo3::exceptions::PyRuntimeError; +use tokio::runtime::Handle; + +use crate::utils::convert_tds_error; + +/// The connection's TDS client in one of two interchangeable edges, plus two +/// transient sentinels that only ever exist between a take and a store. +pub(crate) enum PyClient { + /// Async edge: control-plane + async-core fetch (`block_on`). + Async(TdsClient), + /// Reactor-free edge: the sync cursor's row-pull hot loop. + Sync(TdsSyncClient), + /// Held only between a `mem::replace`-out and the store-back of a flip. + /// Observing it means a prior flip panicked mid-swap; treated as poisoned. + Transitioning, + /// Unrecoverable: a flip or revert failed and the connection is dead. + Dead(String), +} + +/// The shared, clonable client cell every cursor and the connection hold. +pub(crate) type SharedClient = Arc>; + +/// Wraps a freshly connected async client in a shared cell. +pub(crate) fn new_shared(client: TdsClient) -> SharedClient { + Arc::new(Mutex::new(PyClient::Async(client))) +} + +fn poisoned() -> PyErr { + PyRuntimeError::new_err("TDS client mutex was poisoned by a panic") +} + +fn dead(msg: &str) -> PyErr { + PyRuntimeError::new_err(format!("Connection is unusable: {msg}")) +} + +/// Reverts `guard` to the async edge in place and returns `&mut TdsClient`. +/// +/// No-op when already async. On revert failure the cell is marked [`PyClient::Dead`] +/// and an error is surfaced. A `Transitioning`/`Dead` cell yields an error without +/// re-entering the flip. This is the "revert before control-plane" primitive +/// mirrored from the L5 ODBC edge. +pub(crate) fn ensure_async(guard: &mut PyClient) -> Result<&mut TdsClient, PyErr> { + match std::mem::replace(guard, PyClient::Transitioning) { + PyClient::Async(c) => *guard = PyClient::Async(c), + PyClient::Sync(s) => match s.into_async() { + Ok(c) => *guard = PyClient::Async(c), + Err(e) => { + let msg = e.to_string(); + *guard = PyClient::Dead(msg.clone()); + return Err(dead(&msg)); + } + }, + PyClient::Transitioning => { + let msg = "client left in a transitioning state"; + *guard = PyClient::Dead(msg.to_string()); + return Err(dead(msg)); + } + PyClient::Dead(msg) => { + let err = dead(&msg); + *guard = PyClient::Dead(msg); + return Err(err); + } + } + match guard { + PyClient::Async(c) => Ok(c), + _ => unreachable!("ensure_async just stored the Async variant"), + } +} + +/// Locks the cell, reverts to the async edge, and runs a control-plane closure. +fn with_async( + cell: &SharedClient, + handle: &Handle, + f: impl FnOnce(&Handle, &mut TdsClient) -> Result, +) -> Result { + let mut guard = cell.lock().map_err(|_| poisoned())?; + let client = ensure_async(&mut guard)?; + f(handle, client) +} + +/// Reverts the cell to the async edge without running any op (error-recovery / +/// pre-control-plane revert). Used by the sync cursor after a fetch error or +/// before a control-plane transition. +pub(crate) fn revert_to_async(cell: &SharedClient) -> Result<(), PyErr> { + let mut guard = cell.lock().map_err(|_| poisoned())?; + ensure_async(&mut guard)?; + Ok(()) +} + +/// Flips the cell to the reactor-free sync edge in place. +/// +/// Runs `into_sync` under `handle.enter()` so `Handle::try_current()` captures +/// the connection's runtime (ruling 3). A TLS/non-raw transport reports +/// [`SyncConversion::NotEligible`] and the cell stays async — the caller then +/// transparently uses the `block_on` fallback. A failed flip marks the cell dead. +pub(crate) fn flip_to_sync(cell: &SharedClient, handle: &Handle) -> Result<(), PyErr> { + let mut guard = cell.lock().map_err(|_| poisoned())?; + match std::mem::replace(&mut *guard, PyClient::Transitioning) { + PyClient::Async(c) => { + let converted = { + let _entered = handle.enter(); + c.into_sync() + }; + match converted { + SyncConversion::Converted(s) => *guard = PyClient::Sync(s), + SyncConversion::NotEligible(c) => *guard = PyClient::Async(c), + SyncConversion::Failed(e) => { + let msg = e.to_string(); + *guard = PyClient::Dead(msg.clone()); + return Err(dead(&msg)); + } + } + } + PyClient::Sync(s) => *guard = PyClient::Sync(s), + PyClient::Transitioning => { + let msg = "client left in a transitioning state"; + *guard = PyClient::Dead(msg.to_string()); + return Err(dead(msg)); + } + PyClient::Dead(msg) => { + let err = dead(&msg); + *guard = PyClient::Dead(msg); + return Err(err); + } + } + Ok(()) +} + +/// Runs a query on the async edge and collapses forward to the first +/// row-returning result set. Returns `(rows_affected, on_rows)` captured on the +/// async client before any flip, so DML rowcount is oracle-faithful (rule C). +pub(crate) fn run_execute( + cell: &SharedClient, + handle: &Handle, + query: String, + timeout_secs: u32, +) -> Result<(i64, bool), PyErr> { + with_async(cell, handle, |handle, client| { + handle.block_on(async { + if client.has_open_batch() { + client.close_query().await.map_err(convert_tds_error)?; + } + let first = client + .execute( + query, + ExecuteOptions { + timeout: Some(timeout_secs), + ..Default::default() + }, + ) + .await + .map_err(convert_tds_error)?; + if !matches!(first, StatementResult::Rows) { + client.advance_to_rows().await.map_err(convert_tds_error)?; + } + Ok::<(i64, bool), PyErr>((client.last_rows_affected(), client.on_rows())) + }) + }) +} + +/// Whether the current edge is positioned on a row set. +pub(crate) fn is_on_rows(cell: &SharedClient) -> Result { + let guard = cell.lock().map_err(|_| poisoned())?; + Ok(match &*guard { + PyClient::Async(c) => c.on_rows(), + PyClient::Sync(s) => !s.get_metadata().is_empty(), + PyClient::Transitioning | PyClient::Dead(_) => false, + }) +} + +/// Column count of the current result-set metadata on whichever edge is active. +pub(crate) fn metadata_col_count(cell: &SharedClient) -> Result { + let guard = cell.lock().map_err(|_| poisoned())?; + Ok(match &*guard { + PyClient::Async(c) => c.get_metadata().len(), + PyClient::Sync(s) => s.get_metadata().len(), + PyClient::Transitioning | PyClient::Dead(_) => 0, + }) +} + +/// Pulls one row into `writer`. The reactor-free sync arm needs no `block_on`; +/// the async arm (TLS fallback) drives `next_row_into().await` via `block_on`. +/// Both arms route through the same shared parse body, so results are +/// byte-identical. +pub(crate) fn fetch_row_into( + cell: &SharedClient, + handle: &Handle, + writer: &mut (dyn RowWriter + Send), +) -> Result { + let mut guard = cell.lock().map_err(|_| poisoned())?; + match &mut *guard { + PyClient::Sync(s) => s.next_row_into(writer).map_err(convert_tds_error), + PyClient::Async(c) => handle + .block_on(async { c.next_row_into(writer).await }) + .map_err(convert_tds_error), + PyClient::Transitioning => Err(dead("client left in a transitioning state")), + PyClient::Dead(msg) => Err(dead(msg)), + } +} + +/// Reverts to the async edge (if flipped) and closes the current result set. +pub(crate) fn close_resultset(cell: &SharedClient, handle: &Handle) -> Result<(), PyErr> { + with_async(cell, handle, |handle, client| { + handle.block_on(async { client.close_query().await.map_err(convert_tds_error) }) + }) +} + +/// Reverts to the async edge (if flipped) and closes the connection. +pub(crate) fn close_connection(cell: &SharedClient, handle: &Handle) -> Result<(), PyErr> { + with_async(cell, handle, |handle, client| { + handle.block_on(async { client.close_connection().await.map_err(convert_tds_error) }) + }) +} diff --git a/mssql-py-core/src/sync_cursor.rs b/mssql-py-core/src/sync_cursor.rs new file mode 100644 index 00000000..6e1e0c58 --- /dev/null +++ b/mssql-py-core/src/sync_cursor.rs @@ -0,0 +1,171 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! First-class synchronous cursor over the reactor-free sync TDS core. +//! +//! [`PyCoreSyncCursor`] shares the same [`SharedClient`] cell as the async +//! [`PyCoreCursor`](crate::cursor::PyCoreCursor); it exposes an identical DB-API +//! surface but pulls rows through the reactor-free [`TdsSyncClient`] edge instead +//! of `block_on`-over-async. Control-plane work (execute/COLMETADATA/advance/ +//! close) stays async on the shared core — only the SELECT row-pull hot loop +//! flips to the sync edge. There is no protocol-parsing duplication: both cursors +//! drive the one shared token/parse body. +//! +//! Flip discipline: +//! - `execute` runs fully async (rule C), captures the rowcount oracle, then +//! flips to sync only when the statement is positioned on a row set. +//! - fetches use the sync edge (or, on a TLS/non-raw transport that reports +//! `NotEligible`, transparently fall back to the async edge via `block_on`). +//! - any control-plane transition (next `execute`, end-of-rows `close`, cursor +//! `close`, connection close) reverts to async first. + +use pyo3::prelude::*; +use tokio::runtime::Handle; +use tracing::{error, info}; + +use crate::pyclient::{self, SharedClient}; +use crate::row_writer::PyRowWriter; + +/// Python synchronous Cursor class driving the reactor-free sync core. +#[pyclass] +pub struct PyCoreSyncCursor { + tds_client: SharedClient, + runtime_handle: Handle, + has_resultset: bool, + rowcount: i64, +} + +#[pymethods] +impl PyCoreSyncCursor { + #[pyo3(signature = (query, params=None))] + #[allow(unused_variables)] + fn execute( + &mut self, + py: Python, + query: String, + params: Option>>, + ) -> PyResult<()> { + info!("sync execute: Executing query: {}", query); + + let cell = self.tds_client.clone(); + let handle = self.runtime_handle.clone(); + + // Control-plane execute stays async so the rowcount oracle (DML count) is + // captured on the async client before any flip to the sync edge. + let (rowcount, on_rows) = py.detach(|| pyclient::run_execute(&cell, &handle, query, 30))?; + + if on_rows { + // Flip to the reactor-free sync edge for the row-pull hot loop. + // A TLS/non-raw transport reports NotEligible and stays async; the + // fetch path then falls back to block_on transparently. + py.detach(|| pyclient::flip_to_sync(&cell, &handle))?; + } + + self.has_resultset = true; + self.rowcount = rowcount; + Ok(()) + } + + fn fetchone(&mut self, py: Python) -> PyResult>> { + if !self.has_resultset { + return Ok(None); + } + + let cell = self.tds_client.clone(); + let handle = self.runtime_handle.clone(); + + let result: Option = py.detach(|| { + if !pyclient::is_on_rows(&cell)? { + return Ok::<_, PyErr>(None); + } + + let col_count = pyclient::metadata_col_count(&cell)?; + let mut writer = PyRowWriter::new(col_count); + + match pyclient::fetch_row_into(&cell, &handle, &mut writer) { + Ok(true) => Ok(Some(writer)), + Ok(false) => { + // End of rows: revert to async and close the result set. + pyclient::close_resultset(&cell, &handle)?; + Ok(None) + } + Err(e) => { + error!("sync fetchone: fetch failed, reverting to async: {}", e); + // Recover the shared cell onto a live async edge; surface the + // original fetch error (do not poison silently, ruling 4). + let _ = pyclient::revert_to_async(&cell); + Err(e) + } + } + })?; + + if let Some(writer) = result { + Python::attach(|py| { + let py_tuple = writer.to_py_tuple(py)?; + Ok(Some(py_tuple.into())) + }) + } else { + self.has_resultset = false; + Ok(None) + } + } + + fn fetchall(&mut self, py: Python) -> PyResult>> { + if !self.has_resultset { + return Ok(vec![]); + } + + let mut results = Vec::new(); + while let Some(row) = self.fetchone(py)? { + results.push(row); + } + Ok(results) + } + + fn fetchmany(&mut self, py: Python, size: Option) -> PyResult>> { + let fetch_size = size.unwrap_or(1); + let mut results = Vec::new(); + + for _ in 0..fetch_size { + if let Some(row) = self.fetchone(py)? { + results.push(row); + } else { + break; + } + } + Ok(results) + } + + fn close(&mut self, py: Python) -> PyResult<()> { + self.has_resultset = false; + let cell = self.tds_client.clone(); + // Revert any live sync edge so the connection can resume control-plane + // work; best-effort, matching the async cursor's non-draining close. + py.detach(|| { + if let Err(e) = pyclient::revert_to_async(&cell) { + error!("sync close: failed to revert to async edge: {}", e); + } + }); + Ok(()) + } + + #[getter] + fn rowcount(&self) -> i64 { + self.rowcount + } + + fn __repr__(&self) -> String { + "PyCoreSyncCursor()".to_string() + } +} + +impl PyCoreSyncCursor { + pub(crate) fn new(tds_client: SharedClient, runtime_handle: Handle) -> Self { + Self { + tds_client, + runtime_handle, + has_resultset: false, + rowcount: -1, + } + } +} diff --git a/mssql-py-core/tests/rs-only-tests/test_sync_async_cursor_mock.py b/mssql-py-core/tests/rs-only-tests/test_sync_async_cursor_mock.py new file mode 100644 index 00000000..85b09727 --- /dev/null +++ b/mssql-py-core/tests/rs-only-tests/test_sync_async_cursor_mock.py @@ -0,0 +1,274 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""L6 sync + async cursor tests against the mock TDS server. + +These exercise the shared reactor-free sync core wiring without a live SQL +Server: + +- The **sync** cursor (`conn.sync_cursor()`) flips the shared client cell to the + reactor-free `TdsSyncClient` edge for its row-pull hot loop on a **plaintext** + connection, then reverts to the async edge for control-plane work. +- The **async** cursor (`conn.cursor()`) drives the async core via `block_on` + and never flips (regression guard — its public API is unchanged). +- `rowcount` is an additive read-only property sourced identically on both + cursors, so sync == async by construction. +- On a **TLS** connection the sync edge is `NotEligible`, so the sync cursor + transparently falls back to the async `block_on` path (byte-identical rows). + +The mock's built-in query registry answers `SELECT 1` and +`SELECT CAST(1 AS BIGINT), 2, 3`; unknown SELECTs return an empty result set. +""" + +import pytest + +try: + import mssql_mock_tds + + MOCK_TDS_PY_AVAILABLE = True +except ImportError: + MOCK_TDS_PY_AVAILABLE = False + +pytestmark = pytest.mark.skipif( + not MOCK_TDS_PY_AVAILABLE, + reason="mssql_mock_tds not available. Build it with: cd mssql-mock-tds-py && maturin develop", +) + + +def _ctx(server, encryption="Optional"): + """Build a sql_auth client context pointed at the mock server.""" + return { + "server": server.sql_address, + "database": "master", + "user_name": "sa", + "password": "unused-by-mock", + "encryption": encryption, + "trust_server_certificate": True, + } + + +@pytest.fixture +def plaintext_server(): + """A plaintext (non-TLS) mock TDS server; the sync edge is eligible here.""" + server = mssql_mock_tds.PyMockTdsServer(port=0, tls=False) + with server: + yield server + + +@pytest.fixture +def tls_server(): + """A TLS mock TDS server; the sync edge reports NotEligible here. + + TLS against the mock is environment-sensitive (native-tls handshake); when + the local host cannot complete it, the dependent test skips rather than + hanging the suite — mirroring the existing FedAuth TLS tests. + """ + try: + server = mssql_mock_tds.PyMockTdsServer(port=0, tls=True) + except Exception as exc: # noqa: BLE001 - surfaced as a skip below + pytest.skip(f"TLS mock server unavailable in this environment: {exc}") + with server: + yield server + + +def _connect(ctx): + import mssql_py_core + + return mssql_py_core.PyCoreConnection(ctx) + + +# --------------------------------------------------------------------------- +# Sync cursor — reactor-free edge on a plaintext connection +# --------------------------------------------------------------------------- + + +def test_sync_cursor_fetchone_plaintext(plaintext_server): + conn = _connect(_ctx(plaintext_server)) + try: + cur = conn.sync_cursor() + cur.execute("SELECT 1") + assert cur.fetchone() == (1,) + assert cur.fetchone() is None + cur.close() + finally: + conn.close() + + +def test_sync_cursor_fetchall_plaintext(plaintext_server): + conn = _connect(_ctx(plaintext_server)) + try: + cur = conn.sync_cursor() + cur.execute("SELECT CAST(1 AS BIGINT), 2, 3") + assert cur.fetchall() == [(1, 2, 3)] + cur.close() + finally: + conn.close() + + +def test_sync_cursor_fetchmany_plaintext(plaintext_server): + conn = _connect(_ctx(plaintext_server)) + try: + cur = conn.sync_cursor() + cur.execute("SELECT 1") + assert cur.fetchmany(1) == [(1,)] + assert cur.fetchmany(1) == [] + cur.close() + finally: + conn.close() + + +# --------------------------------------------------------------------------- +# Async cursor — unchanged public API (regression guard) +# --------------------------------------------------------------------------- + + +def test_async_cursor_fetchone_plaintext(plaintext_server): + conn = _connect(_ctx(plaintext_server)) + try: + cur = conn.cursor() + cur.execute("SELECT 1") + assert cur.fetchone() == (1,) + assert cur.fetchone() is None + cur.close() + finally: + conn.close() + + +def test_sync_matches_async_rows_plaintext(plaintext_server): + """Sync fetch is byte-identical to the async path for the same statement.""" + conn = _connect(_ctx(plaintext_server)) + try: + query = "SELECT CAST(1 AS BIGINT), 2, 3" + + ac = conn.cursor() + ac.execute(query) + async_rows = ac.fetchall() + ac.close() + + sc = conn.sync_cursor() + sc.execute(query) + sync_rows = sc.fetchall() + sc.close() + + assert sync_rows == async_rows + finally: + conn.close() + + +# --------------------------------------------------------------------------- +# rowcount parity (additive property; same oracle on both cursors) +# --------------------------------------------------------------------------- + + +def test_rowcount_sync_equals_async_select(plaintext_server): + conn = _connect(_ctx(plaintext_server)) + try: + ac = conn.cursor() + ac.execute("SELECT 1") + async_rowcount = ac.rowcount + ac.fetchall() + ac.close() + + sc = conn.sync_cursor() + sc.execute("SELECT 1") + sync_rowcount = sc.rowcount + sc.fetchall() + sc.close() + + assert sync_rowcount == async_rowcount + finally: + conn.close() + + +# --------------------------------------------------------------------------- +# Flip / revert discipline — the shared cell recovers to the async edge +# --------------------------------------------------------------------------- + + +def test_sync_then_async_reuse(plaintext_server): + """After a sync fetch, a control-plane op on the async cursor reverts the + shared cell and succeeds (revert-before-control-plane).""" + conn = _connect(_ctx(plaintext_server)) + try: + sc = conn.sync_cursor() + sc.execute("SELECT 1") + assert sc.fetchone() == (1,) + sc.close() + + ac = conn.cursor() + ac.execute("SELECT 1") + assert ac.fetchone() == (1,) + ac.close() + finally: + conn.close() + + +def test_async_then_sync_reuse(plaintext_server): + conn = _connect(_ctx(plaintext_server)) + try: + ac = conn.cursor() + ac.execute("SELECT 1") + assert ac.fetchone() == (1,) + ac.close() + + sc = conn.sync_cursor() + sc.execute("SELECT 1") + assert sc.fetchone() == (1,) + sc.close() + finally: + conn.close() + + +def test_sync_cursor_close_before_exhaust_reverts(plaintext_server): + """Closing a sync cursor mid-result-set reverts the cell so the connection + stays usable for the next statement.""" + conn = _connect(_ctx(plaintext_server)) + try: + sc = conn.sync_cursor() + sc.execute("SELECT CAST(1 AS BIGINT), 2, 3") + sc.close() # close without draining rows -> revert_to_async + + again = conn.sync_cursor() + again.execute("SELECT 1") + assert again.fetchone() == (1,) + again.close() + finally: + conn.close() + + +def test_sync_cursor_reexecute_reverts(plaintext_server): + """Re-executing on the same sync cursor reverts the previous sync edge + before driving the new control-plane execute.""" + conn = _connect(_ctx(plaintext_server)) + try: + sc = conn.sync_cursor() + sc.execute("SELECT 1") + assert sc.fetchone() == (1,) + + sc.execute("SELECT CAST(1 AS BIGINT), 2, 3") + assert sc.fetchall() == [(1, 2, 3)] + sc.close() + finally: + conn.close() + + +# --------------------------------------------------------------------------- +# TLS connection — sync edge NotEligible, transparent async fallback +# --------------------------------------------------------------------------- + + +def test_sync_cursor_tls_fallback(tls_server): + """On a TLS connection the sync edge is NotEligible; the sync cursor falls + back to the async block_on path and returns byte-identical rows.""" + ctx = _ctx(tls_server, encryption="Optional") + try: + conn = _connect(ctx) + except Exception as exc: # noqa: BLE001 - env-sensitive TLS handshake + pytest.skip(f"TLS connect to mock unavailable in this environment: {exc}") + try: + sc = conn.sync_cursor() + sc.execute("SELECT 1") + assert sc.fetchone() == (1,) + sc.close() + finally: + conn.close() diff --git a/mssql-py-core/tests/test_sync_cursor.py b/mssql-py-core/tests/test_sync_cursor.py new file mode 100644 index 00000000..738d44c0 --- /dev/null +++ b/mssql-py-core/tests/test_sync_cursor.py @@ -0,0 +1,150 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Integration tests for the sync cursor and rowcount against a live server. + +These require a reachable SQL Server (via the standard `client_context` +fixture / `.env`) and are marked `integration`, so they skip when no server is +configured. The mock-server-backed behavioral tests for the sync/async cursors +live in `tests/rs-only-tests/test_sync_async_cursor_mock.py`. +""" + +import pytest + +import mssql_py_core + + +@pytest.mark.integration +def test_sync_cursor_execute(client_context): + conn = mssql_py_core.PyCoreConnection(client_context) + try: + cursor = conn.sync_cursor() + cursor.execute("SELECT 1 AS value") + result = cursor.fetchone() + assert result is not None + assert result[0] == 1 + cursor.close() + finally: + conn.close() + + +@pytest.mark.integration +def test_sync_cursor_fetchall(client_context): + conn = mssql_py_core.PyCoreConnection(client_context) + try: + cursor = conn.sync_cursor() + cursor.execute("SELECT 1 AS value UNION ALL SELECT 2 UNION ALL SELECT 3") + results = cursor.fetchall() + assert [row[0] for row in results] == [1, 2, 3] + cursor.close() + finally: + conn.close() + + +@pytest.mark.integration +def test_sync_cursor_matches_async(client_context): + """The sync fetch path returns byte-identical rows to the async path.""" + query = "SELECT 1 AS a, CAST('x' AS NVARCHAR(10)) AS b UNION ALL SELECT 2, 'y'" + + conn = mssql_py_core.PyCoreConnection(client_context) + try: + ac = conn.cursor() + ac.execute(query) + async_rows = ac.fetchall() + ac.close() + + sc = conn.sync_cursor() + sc.execute(query) + sync_rows = sc.fetchall() + sc.close() + + assert sync_rows == async_rows + finally: + conn.close() + + +@pytest.mark.integration +def test_rowcount_sync_equals_async_select(client_context): + query = "SELECT 1 AS value UNION ALL SELECT 2 UNION ALL SELECT 3" + + conn = mssql_py_core.PyCoreConnection(client_context) + try: + ac = conn.cursor() + ac.execute(query) + async_rowcount = ac.rowcount + ac.fetchall() + ac.close() + + sc = conn.sync_cursor() + sc.execute(query) + sync_rowcount = sc.rowcount + sc.fetchall() + sc.close() + + assert sync_rowcount == async_rowcount + finally: + conn.close() + + +@pytest.mark.integration +def test_dml_rowcount_async(client_context): + """A count-bearing DML statement captures rowcount on the async path.""" + conn = mssql_py_core.PyCoreConnection(client_context) + try: + cursor = conn.cursor() + cursor.execute("CREATE TABLE #l6_rowcount (id INT)") + cursor.execute("INSERT INTO #l6_rowcount (id) VALUES (1), (2), (3)") + assert cursor.rowcount == 3 + cursor.close() + finally: + conn.close() + + +@pytest.mark.integration +def test_sync_cursor_error_mid_fetch_recovers(client_context): + """A fetch error on the sync edge surfaces, then the cell reverts to async + so the connection stays usable (ruling 4: recover via into_async drain).""" + # The conversion trips on the third row, after rows have started streaming. + bad_query = ( + "SELECT CAST(value AS INT) AS n " + "FROM (VALUES ('1'), ('2'), ('notanumber')) AS t(value)" + ) + + conn = mssql_py_core.PyCoreConnection(client_context) + try: + sc = conn.sync_cursor() + sc.execute(bad_query) + with pytest.raises(RuntimeError): + sc.fetchall() + sc.close() + + # The connection reverted to the async edge and remains usable. + ac = conn.cursor() + ac.execute("SELECT 1") + assert ac.fetchone() == (1,) + ac.close() + finally: + conn.close() + + +@pytest.mark.integration +def test_dml_rowcount_sync_equals_async(client_context): + """DML never flips to the sync edge, so both cursors capture the same count.""" + dml = "INSERT INTO #l6_rowcount_parity (id) VALUES (1), (2)" + + conn = mssql_py_core.PyCoreConnection(client_context) + try: + ac = conn.cursor() + ac.execute("CREATE TABLE #l6_rowcount_parity (id INT)") + ac.execute(dml) + async_rowcount = ac.rowcount + ac.close() + + sc = conn.sync_cursor() + sc.execute(dml) + sync_rowcount = sc.rowcount + sc.close() + + assert sync_rowcount == async_rowcount + finally: + conn.close()