From 0fbfb5ed9f37b15da2108d5633909d24d59e1715 Mon Sep 17 00:00:00 2001 From: Saurabh Singh <1623701+saurabh500@users.noreply.github.com> Date: Sun, 9 Aug 2026 23:14:55 -0700 Subject: [PATCH] Add first-class coroutine PyCoreAsyncCursor to mssql-py-core Add PyCoreAsyncCursor, a genuine asyncio coroutine cursor built on pyo3-async-runtimes future_into_py. execute/fetchone/fetchmany/fetchall/ close return awaitables that drive the TDS I/O on the connection's tokio runtime with no block_on on the coroutine path. The !Send client is checked out of the shared cell and the guard dropped before .await, so the Python event loop stays free during fetches. Strictly additive: the default PyCoreCursor and opt-in PyCoreSyncCursor are byte-identical to their prior form. Also scrubs stale 'async cursor' doc labels now that a real coroutine cursor exists. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: e2c378f8-3ba1-4b48-9ebe-5a4ea2bd2761 --- mssql-py-core/Cargo.toml | 1 + mssql-py-core/src/async_cursor.rs | 264 ++++++++++++++++++ mssql-py-core/src/connection.rs | 16 ++ mssql-py-core/src/lib.rs | 2 + mssql-py-core/src/pyclient.rs | 76 ++++- .../rs-only-tests/test_async_cursor_mock.py | 162 +++++++++++ .../test_sync_async_cursor_mock.py | 14 +- .../tests/test_async_cursor_integration.py | 76 +++++ 8 files changed, 600 insertions(+), 11 deletions(-) create mode 100644 mssql-py-core/src/async_cursor.rs create mode 100644 mssql-py-core/tests/rs-only-tests/test_async_cursor_mock.py create mode 100644 mssql-py-core/tests/test_async_cursor_integration.py diff --git a/mssql-py-core/Cargo.toml b/mssql-py-core/Cargo.toml index 3a39ce07..21fde92f 100644 --- a/mssql-py-core/Cargo.toml +++ b/mssql-py-core/Cargo.toml @@ -10,6 +10,7 @@ crate-type = ["cdylib"] [dependencies] pyo3 = { version = "0.29.0", features = ["extension-module"] } +pyo3-async-runtimes = { version = "0.29", features = ["tokio-runtime"] } mssql-tds = { path = "../mssql-tds" } tokio = { version = "1", features = ["full", "rt-multi-thread"] } async-trait = "0.1" diff --git a/mssql-py-core/src/async_cursor.rs b/mssql-py-core/src/async_cursor.rs new file mode 100644 index 00000000..547ee794 --- /dev/null +++ b/mssql-py-core/src/async_cursor.rs @@ -0,0 +1,264 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! First-class asynchronous (coroutine) cursor over the shared TDS core. +//! +//! [`PyCoreAsyncCursor`] shares the same [`SharedClient`] cell as the synchronous +//! [`PyCoreCursor`](crate::cursor::PyCoreCursor), but exposes genuine `asyncio` +//! coroutines: `execute`/`fetchone`/`fetchall`/`fetchmany`/`close` each return a +//! Python awaitable built by [`pyo3_async_runtimes::tokio::future_into_py`]. There +//! is no protocol-parsing duplication — both cursors drive the one shared token +//! and parse body; this cursor only ever uses the async [`TdsClient`] arm and +//! never flips to the reactor-free sync edge. +//! +//! Non-blocking discipline: the coroutine path does **no** `block_on`. The actual +//! TDS I/O (`next_row_into().await`, `execute().await`, …) is spawned onto the +//! owning connection's tokio runtime — the runtime its socket is registered with +//! — via [`Handle::spawn`], and the coroutine simply `.await`s the resulting +//! [`JoinHandle`](tokio::task::JoinHandle). While that awaits, the Python event +//! loop stays free to run other tasks, so a fetch never blocks the loop. +//! +//! Because the shared cell is backed by a [`std::sync::Mutex`] whose guard is +//! `!Send`, the spawned task uses [`pyclient::with_async_client`], which checks +//! the owned client out of the cell (dropping the guard before any `.await`) and +//! stores it back afterwards — even on error — so a mid-fetch failure leaves the +//! connection usable rather than poisoned. + +use mssql_tds::connection::tds_client::{ExecuteOptions, ResultSet, StatementResult, TdsClient}; +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::*; +use pyo3::types::PyList; +use tokio::runtime::Handle; +use tokio::task::JoinError; + +use crate::pyclient::{self, SharedClient}; +use crate::row_writer::PyRowWriter; +use crate::utils::convert_tds_error; + +/// Python asynchronous Cursor class driving the shared async TDS core. +#[pyclass] +pub struct PyCoreAsyncCursor { + tds_client: SharedClient, + runtime_handle: Handle, +} + +#[pymethods] +impl PyCoreAsyncCursor { + #[pyo3(signature = (query, params=None))] + #[allow(unused_variables)] + fn execute<'py>( + &self, + py: Python<'py>, + query: String, + params: Option>>, + ) -> PyResult> { + let cell = self.tds_client.clone(); + let handle = self.runtime_handle.clone(); + + pyo3_async_runtimes::tokio::future_into_py(py, async move { + handle + .spawn(pyclient::with_async_client(cell, |mut client| async move { + let out = run_execute_on(&mut client, query).await; + (client, out) + })) + .await + .map_err(join_err)??; + Python::attach(|py| Ok(py.None())) + }) + } + + fn fetchone<'py>(&self, py: Python<'py>) -> PyResult> { + let cell = self.tds_client.clone(); + let handle = self.runtime_handle.clone(); + + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let row = handle + .spawn(pyclient::with_async_client(cell, |mut client| async move { + let out = fetch_one_on(&mut client).await; + (client, out) + })) + .await + .map_err(join_err)??; + Python::attach(|py| match row { + Some(writer) => Ok(writer.to_py_tuple(py)?.into_any().unbind()), + None => Ok(py.None()), + }) + }) + } + + fn fetchall<'py>(&self, py: Python<'py>) -> PyResult> { + let cell = self.tds_client.clone(); + let handle = self.runtime_handle.clone(); + + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let rows = handle + .spawn(pyclient::with_async_client(cell, |mut client| async move { + let out = fetch_all_on(&mut client).await; + (client, out) + })) + .await + .map_err(join_err)??; + Python::attach(|py| rows_to_py_list(py, &rows)) + }) + } + + #[pyo3(signature = (size=None))] + fn fetchmany<'py>(&self, py: Python<'py>, size: Option) -> PyResult> { + let cell = self.tds_client.clone(); + let handle = self.runtime_handle.clone(); + let fetch_size = size.unwrap_or(1); + + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let rows = handle + .spawn(pyclient::with_async_client( + cell, + move |mut client| async move { + let out = fetch_many_on(&mut client, fetch_size).await; + (client, out) + }, + )) + .await + .map_err(join_err)??; + Python::attach(|py| rows_to_py_list(py, &rows)) + }) + } + + fn close<'py>(&self, py: Python<'py>) -> PyResult> { + let cell = self.tds_client.clone(); + let handle = self.runtime_handle.clone(); + + pyo3_async_runtimes::tokio::future_into_py(py, async move { + handle + .spawn(pyclient::with_async_client(cell, |mut client| async move { + let out = close_on(&mut client).await; + (client, out) + })) + .await + .map_err(join_err)??; + Python::attach(|py| Ok(py.None())) + }) + } + + fn __repr__(&self) -> String { + "PyCoreAsyncCursor()".to_string() + } +} + +impl PyCoreAsyncCursor { + pub(crate) fn new(tds_client: SharedClient, runtime_handle: Handle) -> Self { + Self { + tds_client, + runtime_handle, + } + } +} + +/// Maps a spawned-task join failure (panic/cancellation) onto a Python error. +fn join_err(e: JoinError) -> PyErr { + PyRuntimeError::new_err(format!("async cursor task failed: {e}")) +} + +/// Builds a Python list of row tuples from decoded writers under the GIL. +fn rows_to_py_list<'py>(py: Python<'py>, rows: &[PyRowWriter]) -> PyResult> { + let list = PyList::empty(py); + for writer in rows { + list.append(writer.to_py_tuple(py)?)?; + } + Ok(list.into_any().unbind()) +} + +/// Runs a query on the async edge and collapses forward to the first +/// row-returning result set. +async fn run_execute_on(client: &mut TdsClient, query: String) -> Result<(), PyErr> { + if client.has_open_batch() { + client.close_query().await.map_err(convert_tds_error)?; + } + let first = client + .execute( + query, + ExecuteOptions { + timeout: Some(30), + ..Default::default() + }, + ) + .await + .map_err(convert_tds_error)?; + if !matches!(first, StatementResult::Rows) { + client.advance_to_rows().await.map_err(convert_tds_error)?; + } + Ok(()) +} + +/// Pulls one row on the async edge, closing the result set at end-of-rows. +async fn fetch_one_on(client: &mut TdsClient) -> Result, PyErr> { + if !client.on_rows() { + return Ok(None); + } + let col_count = client.get_metadata().len(); + let mut writer = PyRowWriter::new(col_count); + if client + .next_row_into(&mut writer) + .await + .map_err(convert_tds_error)? + { + Ok(Some(writer)) + } else { + client.close_query().await.map_err(convert_tds_error)?; + Ok(None) + } +} + +/// Pulls up to `size` rows on the async edge, closing the result set if the row +/// stream is exhausted before `size` is reached. +async fn fetch_many_on(client: &mut TdsClient, size: usize) -> Result, PyErr> { + let mut rows = Vec::new(); + if !client.on_rows() { + return Ok(rows); + } + let col_count = client.get_metadata().len(); + for _ in 0..size { + let mut writer = PyRowWriter::new(col_count); + if client + .next_row_into(&mut writer) + .await + .map_err(convert_tds_error)? + { + rows.push(writer); + } else { + client.close_query().await.map_err(convert_tds_error)?; + break; + } + } + Ok(rows) +} + +/// Drains the whole result set on the async edge into decoded row writers. +async fn fetch_all_on(client: &mut TdsClient) -> Result, PyErr> { + let mut rows = Vec::new(); + if !client.on_rows() { + return Ok(rows); + } + let col_count = client.get_metadata().len(); + loop { + let mut writer = PyRowWriter::new(col_count); + if client + .next_row_into(&mut writer) + .await + .map_err(convert_tds_error)? + { + rows.push(writer); + } else { + client.close_query().await.map_err(convert_tds_error)?; + break; + } + } + Ok(rows) +} + +/// Closes the current result set on the async edge, if one is open. +async fn close_on(client: &mut TdsClient) -> Result<(), PyErr> { + if client.has_open_batch() { + client.close_query().await.map_err(convert_tds_error)?; + } + Ok(()) +} diff --git a/mssql-py-core/src/connection.rs b/mssql-py-core/src/connection.rs index 67c6a00b..03949eed 100644 --- a/mssql-py-core/src/connection.rs +++ b/mssql-py-core/src/connection.rs @@ -139,6 +139,22 @@ impl PyCoreConnection { } } + fn async_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::async_cursor::PyCoreAsyncCursor::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/lib.rs b/mssql-py-core/src/lib.rs index e6867c47..24dd19cf 100644 --- a/mssql-py-core/src/lib.rs +++ b/mssql-py-core/src/lib.rs @@ -7,6 +7,7 @@ use std::sync::OnceLock; use mssql_tds::connection::client_context::DriverVersion; mod arrow_bulkcopy; +mod async_cursor; mod bulkcopy; mod connection; mod cursor; @@ -86,6 +87,7 @@ fn mssql_py_core(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; 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 index d12c4064..aacd4fd0 100644 --- a/mssql-py-core/src/pyclient.rs +++ b/mssql-py-core/src/pyclient.rs @@ -8,7 +8,9 @@ //! 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. +//! control-plane op; the default `block_on` cursor and the coroutine +//! [`PyCoreAsyncCursor`](crate::async_cursor::PyCoreAsyncCursor) never flip and +//! always use 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`/ @@ -17,10 +19,12 @@ //! 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. +//! the lock with no `.await`, keeping its fetch path reactor-free. The default +//! cursor's async-arm ops 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. The coroutine cursor instead +//! checks the owned client out of the cell (dropping the guard before any +//! `.await`) via [`with_async_client`], so its path never calls `block_on`. use std::sync::{Arc, Mutex}; @@ -241,3 +245,65 @@ pub(crate) fn close_connection(cell: &SharedClient, handle: &Handle) -> Result<( handle.block_on(async { client.close_connection().await.map_err(convert_tds_error) }) }) } + +/// Checks the async [`TdsClient`] out of the cell, runs an async op on it, then +/// stores it back — used by the coroutine [`PyCoreAsyncCursor`](crate::async_cursor::PyCoreAsyncCursor). +/// +/// This is the async analog of [`with_async`]. Because the cell is backed by a +/// [`std::sync::Mutex`] (whose guard is `!Send`), an async task cannot hold the +/// guard across `.await`. Instead this "checks out" the owned client via +/// [`std::mem::replace`] (leaving [`PyClient::Transitioning`]), drops the guard, +/// awaits `f` on the owned value (which is `Send`), then re-locks and stores it +/// back as [`PyClient::Async`]. The store-back runs even when `f` errors, so a +/// mid-fetch error leaves the connection usable (ruling 4) rather than poisoned. +/// +/// `f` returns the client alongside its result so ownership round-trips cleanly. +/// Must be polled inside a runtime context — the caller spawns it on the +/// connection's runtime so `TdsClient`'s I/O is driven by its own reactor. +pub(crate) async fn with_async_client(cell: SharedClient, f: F) -> Result +where + F: FnOnce(TdsClient) -> Fut, + Fut: std::future::Future)>, +{ + let taken = { + let mut guard = cell.lock().map_err(|_| poisoned())?; + std::mem::replace(&mut *guard, PyClient::Transitioning) + }; + let client = match taken { + PyClient::Async(c) => c, + // The async cursor never flips to sync, but a sync cursor sharing the + // cell may have left it on the sync edge; revert it here. + PyClient::Sync(s) => match s.into_async() { + Ok(c) => c, + Err(e) => { + let msg = e.to_string(); + if let Ok(mut guard) = cell.lock() { + *guard = PyClient::Dead(msg.clone()); + } + return Err(dead(&msg)); + } + }, + PyClient::Transitioning => { + let msg = "client left in a transitioning state"; + if let Ok(mut guard) = cell.lock() { + *guard = PyClient::Dead(msg.to_string()); + } + return Err(dead(msg)); + } + PyClient::Dead(msg) => { + let err = dead(&msg); + if let Ok(mut guard) = cell.lock() { + *guard = PyClient::Dead(msg); + } + return Err(err); + } + }; + + let (client, result) = f(client).await; + + { + let mut guard = cell.lock().map_err(|_| poisoned())?; + *guard = PyClient::Async(client); + } + result +} diff --git a/mssql-py-core/tests/rs-only-tests/test_async_cursor_mock.py b/mssql-py-core/tests/rs-only-tests/test_async_cursor_mock.py new file mode 100644 index 00000000..164df474 --- /dev/null +++ b/mssql-py-core/tests/rs-only-tests/test_async_cursor_mock.py @@ -0,0 +1,162 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Coroutine cursor (`conn.async_cursor()` -> PyCoreAsyncCursor) tests against +the mock TDS server. + +These exercise the genuine `asyncio` coroutine path: `execute`/`fetchone`/ +`fetchall`/`fetchmany`/`close` each return a Python awaitable built by +`pyo3_async_runtimes::tokio::future_into_py`. There is no `block_on` on this +path — the TDS I/O is spawned onto the connection's own tokio runtime and the +coroutine `.await`s the resulting join handle, so the asyncio event loop stays +free to run other tasks while a fetch is in flight. + +The mock answers `SELECT 1` and `SELECT CAST(1 AS BIGINT), 2, 3`; unknown +SELECTs return an empty result set. The discriminating "loop stays free during +a long server-side WAITFOR" proof lives in the integration suite (it needs a +live server); here we prove correctness of the coroutine API plus that the +coroutine actually suspends onto the event loop rather than running inline. +""" + +import asyncio + +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.""" + server = mssql_mock_tds.PyMockTdsServer(port=0, tls=False) + with server: + yield server + + +def _connect(ctx): + import mssql_py_core + + return mssql_py_core.PyCoreConnection(ctx) + + +# --------------------------------------------------------------------------- +# Coroutine API — awaitable execute / fetch* / close +# --------------------------------------------------------------------------- + + +async def test_async_cursor_fetchone_plaintext(plaintext_server): + conn = _connect(_ctx(plaintext_server)) + try: + cur = conn.async_cursor() + await cur.execute("SELECT 1") + assert await cur.fetchone() == (1,) + assert await cur.fetchone() is None + await cur.close() + finally: + conn.close() + + +async def test_async_cursor_fetchall_plaintext(plaintext_server): + conn = _connect(_ctx(plaintext_server)) + try: + cur = conn.async_cursor() + await cur.execute("SELECT CAST(1 AS BIGINT), 2, 3") + assert await cur.fetchall() == [(1, 2, 3)] + await cur.close() + finally: + conn.close() + + +async def test_async_cursor_fetchmany_plaintext(plaintext_server): + conn = _connect(_ctx(plaintext_server)) + try: + cur = conn.async_cursor() + await cur.execute("SELECT 1") + assert await cur.fetchmany(1) == [(1,)] + assert await cur.fetchmany(1) == [] + await cur.close() + finally: + conn.close() + + +async def test_async_matches_sync_rows_plaintext(plaintext_server): + """The coroutine cursor decodes byte-identical rows to the sync cursor for + the same statement — both drive the one shared parse body.""" + conn = _connect(_ctx(plaintext_server)) + try: + query = "SELECT CAST(1 AS BIGINT), 2, 3" + + ac = conn.async_cursor() + await ac.execute(query) + async_rows = await ac.fetchall() + await ac.close() + + sc = conn.sync_cursor() + sc.execute(query) + sync_rows = sc.fetchall() + sc.close() + + assert async_rows == sync_rows + finally: + conn.close() + + +# --------------------------------------------------------------------------- +# The coroutine suspends onto the event loop (does not run inline) +# --------------------------------------------------------------------------- + + +async def test_async_cursor_yields_to_event_loop(plaintext_server): + """A concurrent ticker advances while a fetch coroutine is awaited, proving + the coroutine hands control back to the asyncio loop rather than running to + completion inline on the loop thread. + + The mock answers in microseconds, so this asserts the weaker "yields at + least once" property; the wall-clock "loop stays free during a long fetch" + proof is the WAITFOR integration test. + """ + ticks = 0 + stop = False + + async def ticker(): + nonlocal ticks + while not stop: + ticks += 1 + await asyncio.sleep(0) + + conn = _connect(_ctx(plaintext_server)) + ticker_task = asyncio.create_task(ticker()) + try: + cur = conn.async_cursor() + for _ in range(50): + await cur.execute("SELECT 1") + await cur.fetchall() + await cur.close() + finally: + stop = True + await ticker_task + conn.close() + + assert ticks > 0 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 index 85b09727..a4a24a2a 100644 --- 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 @@ -1,7 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -"""L6 sync + async cursor tests against the mock TDS server. +"""L6 sync + default cursor tests against the mock TDS server. These exercise the shared reactor-free sync core wiring without a live SQL Server: @@ -9,10 +9,12 @@ - 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). +- The **default** cursor (`conn.cursor()`) drives the async core via `block_on` + and never flips (regression guard — its public API is unchanged). It is a + blocking cursor; the genuine coroutine cursor is `conn.async_cursor()` + (`PyCoreAsyncCursor`), covered by the async-focused tests. - `rowcount` is an additive read-only property sourced identically on both - cursors, so sync == async by construction. + cursors, so sync == default 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). @@ -118,7 +120,7 @@ def test_sync_cursor_fetchmany_plaintext(plaintext_server): # --------------------------------------------------------------------------- -# Async cursor — unchanged public API (regression guard) +# Default cursor — unchanged public API (regression guard) # --------------------------------------------------------------------------- @@ -186,7 +188,7 @@ def test_rowcount_sync_equals_async_select(plaintext_server): def test_sync_then_async_reuse(plaintext_server): - """After a sync fetch, a control-plane op on the async cursor reverts the + """After a sync fetch, a control-plane op on the default cursor reverts the shared cell and succeeds (revert-before-control-plane).""" conn = _connect(_ctx(plaintext_server)) try: diff --git a/mssql-py-core/tests/test_async_cursor_integration.py b/mssql-py-core/tests/test_async_cursor_integration.py new file mode 100644 index 00000000..d6340e50 --- /dev/null +++ b/mssql-py-core/tests/test_async_cursor_integration.py @@ -0,0 +1,76 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Integration tests for the coroutine cursor (`PyCoreAsyncCursor`). + +These require a live SQL Server (`client_context` is built from `.env` / +environment by `conftest.py` and skips when credentials are absent). They are +marked `integration` and run in CI where a server is available. +""" + +import asyncio +import time + +import pytest + +import mssql_py_core + + +@pytest.mark.integration +async def test_async_cursor_execute_fetch(client_context): + conn = mssql_py_core.PyCoreConnection(client_context) + try: + cur = conn.async_cursor() + await cur.execute("SELECT 1 AS value") + row = await cur.fetchone() + assert row is not None + assert row[0] == 1 + await cur.close() + finally: + conn.close() + + +@pytest.mark.integration +async def test_async_cursor_nonblocking_during_waitfor(client_context): + """The event loop stays free while a coroutine fetch is in flight. + + A server-side `WAITFOR DELAY` holds the fetch open for ~2s. A concurrent + ticker increments every ~10ms. If the coroutine were `block_on`-backed it + would pin the event-loop thread for the whole delay and the ticker would + barely advance; because the fetch is spawned onto the connection's tokio + runtime and merely awaited, the loop keeps running and the ticker racks up + ticks proportional to the elapsed wall-clock time. + """ + delay_secs = 2 + tick_interval = 0.01 + + ticks = 0 + stop = False + + async def ticker(): + nonlocal ticks + while not stop: + ticks += 1 + await asyncio.sleep(tick_interval) + + conn = mssql_py_core.PyCoreConnection(client_context) + ticker_task = asyncio.create_task(ticker()) + try: + cur = conn.async_cursor() + started = time.monotonic() + await cur.execute(f"WAITFOR DELAY '00:00:0{delay_secs}'; SELECT 1 AS value") + row = await cur.fetchone() + elapsed = time.monotonic() - started + await cur.close() + finally: + stop = True + await ticker_task + conn.close() + + assert row is not None and row[0] == 1 + # The fetch really did take about the WAITFOR duration. + assert elapsed >= delay_secs * 0.5 + # A blocked loop would yield only a handful of ticks; a free loop yields + # roughly elapsed / tick_interval. Use a conservative floor to stay robust + # against scheduler jitter while still failing hard if the loop was pinned. + assert ticks >= (delay_secs / tick_interval) * 0.25