AQE: Async Runtime Infrastructure

[Subrata (subrata-ms)]

Description

This pull request introduces an asynchronous (async/await) API as a gated preview feature to the mssql-py-core Rust crate, enabling users to interact with SQL Server asynchronously from Python via PyO3. The async API is isolated behind a new Cargo feature flag (async-preview) and is not enabled by default, ensuring stability for production users while allowing early adopters to experiment. The implementation includes new async connection and cursor types, process-wide Tokio runtime management, and integration with Python's asyncio via pyo3-async-runtimes.

The most important changes are:

Async API Preview Feature

• Introduced the async-preview Cargo feature: This feature gates all new asynchronous APIs (PyAsyncConnection, PyAsyncCursor, and the shared Tokio runtime). It is off by default, so production builds remain stable unless explicitly opted in.

• Added new dependencies for async support: Integrated the pyo3-async-runtimes crate with the tokio-runtime feature, enabling seamless bridging between Rust's Tokio and Python's asyncio.

Async Connection and Cursor Types

• Implemented PyAsyncConnection and PyAsyncCursor: These new types provide asynchronous connection and cursor functionality, including async connect, close, commit, rollback, and cursor methods, all returning Python awaitables. The API is clearly marked as unstable and emits a FutureWarning on first use. [1] [2]

Shared Tokio Runtime

• Added a process-wide Tokio runtime: The new async_runtime module creates a single, lazily-initialized, multi-threaded Tokio runtime shared by all async connections and cursors, avoiding per-connection thread overhead and ensuring efficient resource usage.

Library Integration

• Conditional compilation and module registration: The async modules are included and registered with PyO3 only when the async-preview feature is enabled, ensuring no impact on the stable sync API. [1] [2] [3]

Internal Refactoring

• Exposed dict_to_client_context for async use: The helper for extracting connection parameters from Python was made pub(crate) to allow reuse in the async connection logic.

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.

[Saurabh Singh (saurabh500)]

Why gate the feature. These are anyhow parallel set of APIs which will leverage async-preview. This by design makes the new APIs opt-in and we don't need to create 2 different py-core binaries for testing in the test pipeline also.

Rust features only make sense when we are OK with 2 binaries being built separately. That also means engineering the pipelines to test with 2 different feature sets. I am hoping we can keep this simple and remove the feature gating for mssql-py-core

[Saurabh Singh (saurabh500)]

Doc-vs-reality gap: this says the runtime is "shared by every PyCoreConnection" and that it replaces "the previous model where each PyCoreConnection owned its own Runtime." Neither is true in this PR — the sync PyCoreConnection in connection.rs still calls Runtime::new() per connection and never touches this shared runtime. Today the process runs two regimes side by side: N per-connection sync runtimes + this one shared async runtime.

For the async surface the single-shared-runtime claim is accurate (future_into_py resolves the process-global runtime via get_runtime()), so suggest scoping the wording to the async surface (e.g. "shared by every async connection and the asyncio bridge") and dropping/qualifying the sync-consolidation claim until that migration actually happens.

[Saurabh Singh (saurabh500)]

Cancellation safety: in pyo3-async-runtimes 0.29, future_into_py cancels the Rust future when the returned asyncio future is cancelled (the Cancellable wrapper in generic.rs, "new behaviour in v0.15"). So await asyncio.wait_for(conn.commit(), timeout=...) — or any cancel — can drop this future while it is suspended inside commit_transaction().await, after the TM_COMMIT bytes are on the wire but before the server's DONE token is read.

The MutexGuard is released, but the shared TdsClient's TDS byte stream is now desynced: a subsequent rollback() (or, later, a cursor execute) on the same connection will read the stale response as its own -> protocol corruption on a live connection. The mutex gives mutual exclusion, not cancellation atomicity. Same applies to rollback (L245) and close (L164), and will apply to cursor I/O once it shares this client.

For a gated preview, documenting this is acceptable, but the GA design should either run these ops non-cancellably (spawn + await the join handle) or poison the connection when a future is dropped mid-flight. Cheapest to decide before cursor I/O piles onto the same client.

[Saurabh Singh (saurabh500)]

Subrata (@subrata-ms) you could add a todo here and revisit this part of the design. This is critical though

[Saurabh Singh (saurabh500)]

Nit: commit (and rollback at L235) take &mut self but only read self.tds_client via .as_ref() and clone the Arc. &self would suffice — as cursor() already does — and lets Python hold a shared rather than exclusive borrow. Harmless today, just slightly over-restrictive.

[Subrata (subrata-ms)]

Agreed. &mut self here gives real zero protection. It probably misleads the reader. Changing it to &self

[Saurabh Singh (saurabh500)]

Nit: the comment says init "Must run before any code that touches Tokio," but nothing else in this function touches Tokio, so the real constraint is just "before the first connection" — which holds regardless of position. Slightly overstates the ordering coupling. (For the record: pyo3_async_runtimes::tokio::init panics if called twice; the Once guard in async_runtime::init makes that unreachable, which is correct.)

[Subrata (subrata-ms)]

fair point fixed in next commit.