AQE: Async Runtime Infrastructure

[Subrata (subrata-ms)]

Description

This pull request introduces a new asynchronous API surface to the mssql-py-core package, providing Python users with async connection and cursor types for interacting with SQL Server using asyncio. The async implementation is marked as a preview and is isolated from the existing synchronous API. Key changes include the addition of async connection/cursor types, a shared Tokio runtime for async operations, and integration of these types into the Python module initialization.

New Asynchronous API:

• Added PyAsyncConnection and PyAsyncCursor types, providing Python awaitable methods for connecting, closing, committing, rolling back, and creating cursors in an async context. These types use the shared Tokio runtime and return awaitables compatible with Python's asyncio. The API is marked as unstable/preview and emits a FutureWarning on first use. (mssql-py-core/src/async_connection.rs, mssql-py-core/src/async_cursor.rs) [1] [2]

• Integrated the new async types into the Python extension module, registering both classes and initializing the shared Tokio runtime at module load. This ensures async connections and awaitables work consistently across the process. (mssql-py-core/src/lib.rs) [1] [2] [3]

Async Runtime Management:

• Added a process-wide, lazily-initialized Tokio runtime specifically for the async API. This avoids creating multiple thread pools for each connection and ensures all async operations share the same executor. (mssql-py-core/src/async_runtime.rs)

Dependency and Utility Updates:

• Added pyo3-async-runtimes as a dependency to bridge Rust futures to Python awaitables, and updated Cargo.toml accordingly. (mssql-py-core/Cargo.toml)
• Made dict_to_client_context public within the crate to allow async connection creation to reuse the same client context parsing as the sync code. (mssql-py-core/src/connection.rs)
• Added a TODO to the TDS error conversion utility for future compliance with Python DB-API error classes. (mssql-py-core/src/utils.rs)

Related Issues

ADO work item: https://sqlclientdrivers.visualstudio.com/mssql-python/_workitems/edit/46956 , https://sqlclientdrivers.visualstudio.com/mssql-python/_workitems/edit/46957

Checklist

• cargo bfmt passes
• cargo bclippy passes
• cargo btest passes
• New/changed functionality has tests
• Public API changes are documented

[Copilot]

Pull request overview

Adds preview-gated Python async infrastructure for mssql-py-core.

Changes:

• Adds shared Tokio/asyncio runtime integration.
• Introduces async connection lifecycle and transaction methods.
• Exposes an initial async cursor scaffold.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
mssql-py-core/Cargo.toml	Adds async dependency and feature flag.
mssql-py-core/src/lib.rs	Registers preview modules and classes.
mssql-py-core/src/connection.rs	Shares context parsing internally.
mssql-py-core/src/async_runtime.rs	Configures the shared Tokio runtime.
mssql-py-core/src/async_connection.rs	Implements async connection operations.
mssql-py-core/src/async_cursor.rs	Adds the async cursor scaffold.

Suppressed comments (3)

mssql-py-core/src/async_connection.rs:164

• Cancelling the Python Future returned by future_into_py drops this Rust future. Because the client has already been removed from self, cancellation before close_connection() runs makes close() non-retryable; if a cursor still owns the Arc, the transport can remain open until that cursor is dropped. Start cleanup in a cancellation-independent task or retain recoverable connection state until shutdown is guaranteed.

        let client_opt = self.tds_client.take();

        pyo3_async_runtimes::tokio::future_into_py(py, async move {

mssql-py-core/src/async_connection.rs:248

• The rollback path has the same cancellation hole as commit: dropping the Rust future after TM_ROLLBACK is sent can leave its response unread while releasing the client mutex. Subsequent operations may then read the rollback response as their own. Ensure cancellation drains the response or invalidates/closes the connection before it can be reused.

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            tracing::info!("PyAsyncConnection::rollback: sending TM_ROLLBACK");
            let mut guard = client.lock().await;
            guard.rollback_transaction(None, None).await.map_err(|e| {

mssql-py-core/src/async_connection.rs:213

• future_into_py propagates asyncio cancellation by dropping this Rust future. If cancellation occurs after TM_COMMIT is written but before its response is consumed, the mutex is released while unread TDS tokens remain, so the next operation can consume the stale response and desynchronize the session. Run the wire command in a cancellation-safe task and drain/close the connection before permitting reuse.

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            tracing::info!("PyAsyncConnection::commit: sending TM_COMMIT");
            let mut guard = client.lock().await;
            guard.commit_transaction(None, None).await.map_err(|e| {

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Saurabh Singh (saurabh500)]

Reviewed the async preview infrastructure. The feature gating is clean, symmetric, and correctly defaults off, and the Arc<Mutex<TdsClient>> ownership pattern (clone/take synchronously before entering future_into_py) is exactly right for producing 'static + Send futures. Verified it compiles clean under cargo clippy --features async-preview --all-targets -- -D warnings.

Submitting as Comment. Two items worth resolving before cursor I/O lands on the shared client (details inline):

1. Cancellation safety of the async lifecycle methods vs. future_into_py's cancel-on-drop behavior — the main one.
2. async_runtime.rs module doc describes a sync-runtime consolidation that hasn't happened yet.

Also a heads up: the PR body has no linked GitHub issue / ADO work item, and there are no tests exercising the async surface (it builds in CI via --all-features but nothing runs it). Not blocking for a gated preview, but flagging.

[David Engel (David-Engel)]

Reviewed the diff independently and then read the existing threads — most of what I found was already covered by Copilot and Saurabh Singh (@saurabh500), so I'm only posting the four items nobody has raised yet. I verified the branch compiles and lints clean with cargo clippy --all-features --all-targets in a scratch worktree.

Explicitly not re-raising (already well covered): cancellation safety of the wire commands, the missing tests, the swallowed PyErr::warn, the async_runtime.rs doc-vs-reality gap, &mut self on commit/rollback, and the non-optional pyo3-async-runtimes dependency. I agree with all of them, and with Saurabh Singh (@saurabh500) that the cancellation question is the one to settle before cursor I/O piles onto the same client.

[github-actions]

📊 Code Coverage Report

🔥 Diff Coverage 92%

🎯 Overall Coverage 91.5%

📦 Project: mssql-tds + mssql-odbc + mssql-py-core ℹ️ Note: diff coverage is reported, not enforced.

---

Diff Coverage

Diff: main...HEAD, staged and unstaged changes

• mssql-py-core/src/async_connection.rs (91.0%): Missing lines 88,131,133-134,183,223-224,259-260
• mssql-py-core/src/async_cursor.rs (100%)
• mssql-py-core/src/async_runtime.rs (100%)
• mssql-py-core/src/connection.rs (100%)
• mssql-py-core/src/lib.rs (100%)

Summary

• Total: 114 lines
• Missing: 9 lines
• Coverage: 92%

mssql-py-core/src/async_connection.rs

  84     /// across `.await` points without corrupting the TDS byte stream.
  85     tds_client: Option<Arc<Mutex<TdsClient>>>,
  86 }
  87 
! 88 #[pymethods]
  89 impl PyAsyncConnection {
  90     /// Establish a TDS connection asynchronously.
  91     ///
  92     /// ```python

---

  127             let client = provider
  128                 .create_client(context, &datasource, None)
  129                 .await
  130                 .map_err(|e| {
! 131                     tracing::error!("PyAsyncConnection::connect: failed: {}", e);
  132                     // TODO(User Story 47181): map TdsError to a DB-API-compliant exception, preserving SQLSTATE + server error number.
! 133                     PyRuntimeError::new_err(format!("Failed to connect to SQL Server: {e}"))
! 134                 })?;
  135 
  136             tracing::info!("PyAsyncConnection::connect: connection established");
  137             Python::attach(|py| {
  138                 Py::new(

---

  179             if let Err(e) = guard.close_connection().await {
  180                 // Match sync-path semantics: log and swallow. The connection
  181                 // is treated as closed regardless — the transport will be
  182                 // dropped when the Arc's last reference goes away.
! 183                 tracing::warn!(
  184                     "PyAsyncConnection::close: error during graceful shutdown: {}",
  185                     e
  186                 );
  187             }

---

  219                 tracing::error!("PyAsyncConnection::commit: failed: {}", e);
  220                 // TODO(User Story 47181): map TdsError to a DB-API-compliant exception, preserving SQLSTATE + server error number.
  221                 PyRuntimeError::new_err(format!("Commit failed: {e}"))
  222             })?;
! 223             tracing::info!("PyAsyncConnection::commit: transaction committed");
! 224             Python::attach(|py| Ok(py.None()))
  225         })
  226     }
  227 
  228     /// Roll back the current TDS transaction asynchronously.

---

  255                 tracing::error!("PyAsyncConnection::rollback: failed: {}", e);
  256                 // TODO(User Story 47181): map TdsError to a DB-API-compliant exception, preserving SQLSTATE + server error number.
  257                 PyRuntimeError::new_err(format!("Rollback failed: {e}"))
  258             })?;
! 259             tracing::info!("PyAsyncConnection::rollback: transaction rolled back");
! 260             Python::attach(|py| Ok(py.None()))
  261         })
  262     }
  263 
  264     /// Create an async cursor bound to this connection.

---

---

🔗 Quick Links

View Azure DevOps Build · Coverage Report

[Sumit Sarabhai (sumitmsft)]

This PR maps every server/transport failure on connect, commit, and rollback to a bare PyRuntimeError (e.g. "Failed to connect to SQL Server", "Commit failed", "Rollback failed"). I want to flag this now so it doesn't harden into the shipped contract.

The design doc is explicit that we must be DB-API 2.0 compliant on error mapping. All async errors surface as subclasses of mssql_python.exceptions, the same hierarchy the sync path uses, so existing "except" handlers work.

Every abnormal outcome maps to exactly one Python exception class from a small, documented hierarchy. No RuntimeError catch-all in user-facing paths. The server errors must raise from the DB-API hierarchy (DatabaseError / ProgrammingError / OperationalError), preserving message and, where available, SQLSTATE / server error number.

A raw PyRuntimeError breaks that contract in two concrete ways: it isn't part of the DB-API exception tree, so a caller's except mssql_python.DatabaseError / except mssql_python.OperationalError won't catch it and flattening everything to one string discards the SQLSTATE / server error number.

I understand the real mapping lands with the error-taxonomy work, not this infrastructure PR, but since these strings are user-visible the moment the preview is exercised, could we add a TODO at each of these three sites and the single shared convert_tds_error style function, so it's unambiguous these are placeholders and not the intended DB-API-compliant contract? That keeps us honest that the async and sync paths must converge on one hierarchy.

[David Engel (David-Engel)]

Re-reviewed at 8ee6a76e. All four of my earlier items are addressed or consciously deferred with tracking, and the async surface now has tests that actually execute in CI — nice turnaround.

Verified locally: cargo clippy --all-targets -- -D warnings clean with default features; built the extension via dev/test-python.sh --skip-integration and all 5 non-integration async tests pass, including test_future_warning_propagates_when_promoted_to_error, which confirms the PyErr::warn propagation fix works end to end rather than just reading correctly. On the CI side, the ADO validation build 166415 task "Test Python bindings (mssql-py-core)" succeeded, and since dev/test-python.sh runs pytest tests/ -m "not longhaul and not smoke" against the live sql1 container, the integration-marked async tests genuinely ran against a server. Dropping the feature gate resolved my earlier "CI never runs this code" concern without needing any pipeline edit.

Everything below is non-blocking polish. All of it is the same theme Sumit Sarabhai (@sumitmsft) raised in his comment on the DB-API error mapping: placeholder semantics on a preview surface tend to harden into the shipped contract, so it's cheapest to mark or fix them now. His comment is the more consequential instance of this; mine are smaller siblings. In particular, one of the inline notes below walks back a test suggestion I made earlier that would have cut against his ask.

[David Engel (David-Engel)]

Sumit Sarabhai (@sumitmsft) +1 on getting TODOs at those three sites plus the shared conversion function before the strings become the de facto contract. One factual refinement that I think strengthens the ask rather than weakening it.

The comment says flattening to one string "discards the SQLSTATE / server error number." That is not quite what happens today. SqlErrorInfo's Display in mssql-tds/src/error/mod.rs (L37-L51) renders:

Sql Error: {number}: Class {class}: State {state}: {message} on {server} in {proc} at line {line}

so the number and state do survive into the PyRuntimeError message — just as unparseable text a caller would have to regex.

The real loss is structural, and it is one layer up. Error::SqlServerError { diagnostics } already carries a Vec<SqlErrorInfo> with number, class, state, server_name, proc_name, and line_number as typed fields, and multiple errors per batch are preserved. All of that structure is discarded by format!("Commit failed: {e}") at the pyo3 boundary.

So the mapping function you are asking for has a fully populated, already-structured source to work from — this is a boundary problem, not a data-availability problem, which should make the eventual DB-API mapping mostly mechanical. Worth capturing in the TODO text so whoever picks up the taxonomy work knows the diagnostics are already there and does not go re-plumbing the TDS layer.

Related: I left a note on tests/test_async_connection.py walking back a suggestion I had made to pin the transaction tests to pytest.raises(RuntimeError, match="Commit failed"). That would have hardened exactly the placeholder you are flagging, so I switched it to matching on the server error number (3902 / 3903), which stays valid after the taxonomy lands.