Skip to content

[py-core] Async Connection Foundation #288

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Subrata (subrata-ms) wants to merge 8 commits into
base: main
Choose a base branch
from
+340 −168
Open

[py-core] Async Connection Foundation #288

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
d740d6e
[TASK 46956] Add python_logger and tracing improvements to PyAsyncCon…
subrata-ms Aug 14, 2026
b41bf3f
[TASK 46956] Add entry and completion tracing to PyAsyncConnection.close
subrata-ms Aug 14, 2026
070a741
[TASK 46956] Trim docstrings in async_connection.rs to match repo style
subrata-ms Aug 14, 2026
b0b732f
[TASK 46956] Add default query timeout to PyAsyncConnection
subrata-ms Aug 14, 2026
6a75af1
[TASK 46956] Add closed property and is_connected() to PyAsyncConnection
subrata-ms Aug 14, 2026
2b152dd
[TASK 46956] Add __aenter__ / __aexit__ to PyAsyncConnection
subrata-ms Aug 14, 2026
7b40950
[TASK 46956] Add __repr__ and factor TDS-error mapping in PyAsyncConn…
subrata-ms Aug 14, 2026
280ad92
[TASK 46956] Consolidate PyAsyncConnection tests and cover cursor() path
subrata-ms Aug 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
264 changes: 124 additions & 140 deletions mssql-py-core/src/async_connection.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,23 +3,9 @@

//! Asynchronous connection API for the Core TDS backend.
//!
//! # ⚠️ Preview API — unstable
//! Preview API — unstable. First use emits a `FutureWarning`.
//!
//! The types and methods in this module are **not** part of the stable
//! `mssql-py-core` surface. Signatures, error behavior, and internal
//! semantics may change without notice in any release. First use in a
//! Python process emits a [`FutureWarning`] via `warnings.warn`.
//!
//! Sibling of `connection.rs` (the synchronous surface). Every type defined
//! here submits its I/O to the shared process-wide Tokio runtime via
//! [`crate::async_runtime`] and returns Python awaitables through
//! `pyo3_async_runtimes::tokio::future_into_py`, so callers can `await` the
//! results from `asyncio`.
//!
//! Invariant: one async connection maps to exactly one async cursor, one
//! `TdsClient`, and one TDS wire session.
//!
//! [`FutureWarning`]: https://docs.python.org/3/library/exceptions.html#FutureWarning
//! Invariant: one async connection ↔ one async cursor ↔ one `TdsClient`.

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
Expand All @@ -34,9 +20,9 @@ use mssql_tds::connection_provider::tds_connection_provider::TdsConnectionProvid

use crate::async_cursor::PyAsyncCursor;
use crate::connection::PyCoreConnection;
use crate::python_logger_adapter::scoped_tracing_bridge;

/// Emit a `FutureWarning` the first time any async API is exercised in this
/// process. Silenceable by callers via `warnings.filterwarnings(...)`.
/// One-shot `FutureWarning` per process; silenceable via `warnings.filterwarnings`.
static PREVIEW_WARNED: AtomicBool = AtomicBool::new(false);

fn emit_preview_warning(py: Python<'_>) -> PyResult<()> {
Expand All @@ -55,115 +41,141 @@ fn emit_preview_warning(py: Python<'_>) -> PyResult<()> {
Ok(())
}

/// Asynchronous Python connection backed by the Core TDS client.
///
/// # ⚠️ Preview API — unstable
/// Map a TDS error to a Python exception with per-operation context.
///
/// Preview surface: API, method signatures, error behavior, and internal
/// semantics may change without notice in minor releases. Do not depend on
/// it from production code.
/// TODO(User Story 47181): map TdsError to a DB-API-compliant exception,
/// preserving SQLSTATE + server error number.
/// <https://sqlclientdrivers.visualstudio.com/mssql-python/_workitems/edit/47181>
fn map_tds_error(op: &str, user_msg: &str, e: impl std::fmt::Display) -> PyErr {
tracing::error!("PyAsyncConnection::{op}: failed: {e}");
PyRuntimeError::new_err(format!("{user_msg}: {e}"))
}

/// Asynchronous Python connection backed by the Core TDS client.
///
/// Instances are created via [`PyAsyncConnection::connect`], which returns a
/// Python awaitable. The awaitable resolves on the caller's `asyncio` loop
/// once the TCP + TLS + login handshake has completed on the shared Tokio
/// runtime.
/// Preview API — unstable.
///
/// TODO(User Story 47180 [mssql-python] Cancel API and Cancellation Bridge):
/// cancellation of a suspended `commit`, `rollback`, or `close` future can
/// desync the TDS byte stream (bytes written, response not yet read), so a
/// subsequent operation on the same connection may read a stale response
/// and corrupt the wire. Callers must not cancel these awaitables against a
/// connection they intend to keep using. Cancellation-safe semantics are
/// tracked at
/// <https://sqlclientdrivers.visualstudio.com/mssql-python/_workitems/edit/47180>.
/// desync the TDS byte stream. Callers must not cancel these awaitables
/// against a connection they intend to keep using.
/// <https://sqlclientdrivers.visualstudio.com/mssql-python/_workitems/edit/47180>
#[pyclass]
pub struct PyAsyncConnection {
/// Wrapped in `Option` so `close()` can take ownership and drop the
/// client; wrapped in `Arc<tokio::sync::Mutex<...>>` so the (upcoming)
/// async cursor and connection-level lifecycle methods can share access
/// across `.await` points without corrupting the TDS byte stream.
/// `Option` so `close()` can `take()`; `Arc<Mutex<>>` for cursor sharing.
tds_client: Option<Arc<Mutex<TdsClient>>>,
/// Default query timeout (seconds) applied to cursors created from this
/// connection. `0` = no timeout, per pyodbc/ODBC `SQL_ATTR_QUERY_TIMEOUT`.
/// Pure Python-side state — the setter performs no I/O.
default_query_timeout: u32,
}

#[pymethods]
impl PyAsyncConnection {
/// Establish a TDS connection asynchronously.
///
/// ```python
/// conn = await PyAsyncConnection.connect(client_context_dict)
/// ```
///
/// Dictionary parsing runs synchronously on the calling thread (it needs
/// the GIL). The network handshake is submitted to the shared Tokio
/// runtime and driven concurrently with the caller's asyncio loop.
/// Establish a TDS connection. Dict parsing is synchronous; the network
/// handshake runs on the shared Tokio runtime.
#[classmethod]
#[pyo3(signature = (client_context_dict, python_logger=None))]
fn connect<'py>(
cls: &Bound<'py, PyType>,
client_context_dict: &Bound<'_, PyDict>,
python_logger: Option<&Bound<'_, PyAny>>,
) -> PyResult<Bound<'py, PyAny>> {
let py = cls.py();

// Preview API: emit a one-shot FutureWarning so callers see the
// instability signal at runtime.
// `DefaultGuard` is `!Send`, so its coverage ends when this method returns.
let _guard = python_logger

@saurabh500 Saurabh Singh (saurabh500) Aug 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This python_logger bridge only covers the synchronous part of connect. scoped_tracing_bridge returns a thread-local DefaultGuard, and _guard is dropped the moment connect returns — but the actual TCP+TLS+login handshake runs later, inside the future_into_py block, on a Tokio worker thread. So every event in that block bypasses the caller's logger: the connection established line and, more importantly, the connect-failure error from map_tds_error("connect", ...) (line 121) — which is exactly the log someone who bothered to pass a python_logger would most want.

The sync PyCoreConnection::new doesn't hit this because it drives the handshake with runtime.block_on(...) inline, under the guard, on the same thread.

Can we move the bridge into the async block (hand the Arc<Py<PyAny>> into the future and set_default there), or — if capturing only the prelude is intentional for this foundation — document that per-call python_logger doesn't cover the async handshake? As written, passing a logger silently drops the failure diagnostics.

.map(|logger| scoped_tracing_bridge(Arc::new(logger.clone().unbind()), file!()));
Comment on lines +86 to +88
Comment on lines +86 to +88

@David-Engel David Engel (David-Engel) Aug 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Blocking: this bridge does not cover any of the async work, including the connect-failure log.

scoped_tracing_bridge calls tracing::subscriber::set_default, which is thread-local and is undone when the guard drops at the end of connect. The future_into_py body runs later, on a Tokio worker thread. So of the events in connect, the Python logger receives only the four synchronous pre-flight info! lines, and never receives:

  • "connection established" (L123)
  • map_tds_error("connect", ...)'s tracing::error! (L50) - the handshake-failure diagnostic, which is the whole reason a caller passes a logger

This is the opposite of the sync path it is modeled on: PyCoreConnection::new uses runtime.block_on on the same thread, so its tracing::error!("Failed to connect to SQL Server: {}", e) is inside the guard. Separately, close/commit/rollback/cursor take no logger and are never bridged at all, so the connection is unlogged for its entire lifetime after connect returns.

Concrete fix - propagate a Dispatch across the await point instead of a thread-local guard. Add alongside the existing helpers in python_logger_adapter.rs:

pub fn tracing_bridge_dispatch(logger: Arc<Py<PyAny>>, module_name: &'static str) -> tracing::Dispatch {
    let layer = PythonLoggerLayer::new(logger, module_name)
        .with_filter(filter_fn(|m| m.target().starts_with("mssql_py_core")));
    tracing::Dispatch::new(Registry::default().with(layer))
}

then here:

use tracing::instrument::WithSubscriber;

let dispatch = python_logger
    .map(|l| tracing_bridge_dispatch(Arc::new(l.clone().unbind()), file!()));
// ... sync tracing calls stay inside a set_default guard as today ...
pyo3_async_runtimes::tokio::future_into_py(py, async move {
    let fut = async move { /* existing body */ };
    match dispatch { Some(d) => fut.with_subscriber(d).await, None => fut.await }
})

Dispatch is Clone + Send + Sync, so it can also be stored on the pyclass and reused by close/commit/rollback - which is what makes the parameter actually useful.

If that is too much for this PR, I would accept: drop the python_logger param for now, or keep it and (a) state the limitation in the Python-visible docstring rather than only a code comment, and (b) file a follow-up work item. What I would rather not ship is a python_logger= that silently swallows the failure path.

Supporting point: there is no test that passes a python_logger. A test asserting which records reach py_core_log for both a successful and a failed connect would have caught this, and is the right regression guard for whichever fix you pick.


emit_preview_warning(py)?;

tracing::info!("PyAsyncConnection::connect: initiating async connection");

@David-Engel David Engel (David-Engel) Aug 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: "initiating async connection" is immediately followed by "extracting client context" at the same level - the first adds nothing. Also, the sync path uses debug! for the dict conversion; info! for every step is chatty on a per-connection path.

@saurabh500 Saurabh Singh (saurabh500) Aug 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These two info! lines fire back to back with no work between them (initiating async connection, then extracting client context). Drop the first — it's noise.


tracing::info!("PyAsyncConnection::connect: extracting client context");
let context = PyCoreConnection::dict_to_client_context(client_context_dict)?;
let datasource = context.data_source.clone();

tracing::info!(
"PyAsyncConnection::connect: encryption mode={:?}, trust_server_certificate={}, host_name_in_cert={:?}",
"PyAsyncConnection::connect: encryption mode={:?}, trust_server_certificate={}, host_name_in_cert={:?}, server_certificate={:?}",
context.encryption_options.mode,
context.encryption_options.trust_server_certificate,
context.encryption_options.host_name_in_cert,
context.encryption_options.server_certificate,
);

tracing::info!(
"PyAsyncConnection::connect: authentication method={:?}",
context.tds_authentication_method,
);

tracing::info!(
"PyAsyncConnection::connect: attempting connection to datasource: {}",
datasource
);

pyo3_async_runtimes::tokio::future_into_py(py, async move {
tracing::info!(
"PyAsyncConnection::connect: opening TDS connection to {}",
datasource
);
let provider = TdsConnectionProvider {};
let client = provider
.create_client(context, &datasource, None)
.await
.map_err(|e| {
tracing::error!("PyAsyncConnection::connect: failed: {}", e);
// TODO(User Story 47181): map TdsError to a DB-API-compliant exception, preserving SQLSTATE + server error number.
PyRuntimeError::new_err(format!("Failed to connect to SQL Server: {e}"))
})?;
.map_err(|e| map_tds_error("connect", "Failed to connect to SQL Server", e))?;

tracing::info!("PyAsyncConnection::connect: connection established");
Python::attach(|py| {
Py::new(
py,
PyAsyncConnection {
tds_client: Some(Arc::new(Mutex::new(client))),
default_query_timeout: 0,
},
)
})
})
}

/// Close the TDS connection asynchronously.
///
/// ```python
/// await conn.close()
/// ```
///
/// Sends the TDS logout token and tears down the underlying transport.
/// The awaitable is submitted to the shared Tokio runtime so the calling
/// asyncio loop stays unblocked while the graceful shutdown runs.
///
/// Idempotent: awaiting `close()` on an already-closed connection
/// resolves immediately with no I/O. If the graceful shutdown itself
/// errors, the error is logged at `warn` level and the connection is
/// still considered closed — the OS closes the socket on drop either
/// way, so we never leak the resource.
/// Default query timeout (seconds) inherited by cursors created from this
/// connection. `0` means no timeout.
#[getter]
fn timeout(&self) -> u32 {
self.default_query_timeout
}

/// Set the default query timeout (seconds) for future cursors. Existing
/// cursors and in-flight queries are unaffected. Negative values are
/// rejected by PyO3's `u32` extractor (`OverflowError`).
#[setter]
fn set_timeout(&mut self, value: u32) {
tracing::info!(
"PyAsyncConnection::set_timeout: default query timeout set to {}s",
value
);
self.default_query_timeout = value;
}
Comment on lines +143 to +153

@David-Engel David Engel (David-Engel) Aug 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggestion: negative values raise OverflowError here, but mssql_python.Connection.timeout's setter raises ValueError("Timeout cannot be negative") (and TypeError for non-int). test_timeout_setter_rejects_negative locks OverflowError in.

If the Python wrapper will always sit on top, this is fine - but worth a one-line note in the docstring so nobody is surprised when the two surfaces disagree. Otherwise take i64 and raise ValueError explicitly for parity.


/// True after `close()` has been awaited (or if `connect()` never produced a

@saurabh500 Saurabh Singh (saurabh500) Aug 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

connect() returns Err on handshake failure and only ever builds the instance with Some(client), so there's no public path to an instance that starts out closed. The "or if connect() never produced a live client" clause describes a state this API can't reach — trim it.

/// live client). Cheap LBYL check; performs no I/O.
#[getter]
fn closed(&self) -> bool {
self.tds_client.is_none()
}

/// Inverse of `.closed`, provided for sync-path parity with `PyCoreConnection`.
fn is_connected(&self) -> bool {
self.tds_client.is_some()
}

fn __repr__(&self) -> &'static str {
if self.tds_client.is_none() {
"PyAsyncConnection(closed)"
} else {
"PyAsyncConnection(connected)"
}
}

/// Close the connection. Idempotent. Shutdown errors are logged and swallowed.
fn close<'py>(&mut self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
// Detach the client from `self` synchronously (while `&mut self` is
// valid) so the future can own it for `'static + Send`. Subsequent
// method calls on this connection will see `tds_client == None`.
tracing::info!("PyAsyncConnection::close: initiating close");
// `take()` before spawning: gives the future 'static ownership; marks conn closed.
let client_opt = self.tds_client.take();

pyo3_async_runtimes::tokio::future_into_py(py, async move {
Expand All @@ -177,35 +189,38 @@ impl PyAsyncConnection {
);
let mut guard = client.lock().await;
if let Err(e) = guard.close_connection().await {
// Match sync-path semantics: log and swallow. The connection
// is treated as closed regardless — the transport will be
// dropped when the Arc's last reference goes away.
// Log and swallow; connection is closed regardless.
tracing::warn!(
"PyAsyncConnection::close: error during graceful shutdown: {}",
e
);
}
tracing::info!("PyAsyncConnection::close: connection closed");
Python::attach(|py| Ok(py.None()))
})
}

/// Commit the current TDS transaction asynchronously.
///
/// ```python
/// await conn.commit()
/// ```
///
/// Sends a TM_COMMIT (Transaction Manager COMMIT) request over the wire
/// and awaits the server's DONE token. Raises `RuntimeError`
/// synchronously if the connection has already been closed.
///
/// If no transaction is currently open on the server, the commit will
/// fail with the server's own error (SQL Server 3902 — "The COMMIT
/// TRANSACTION request has no corresponding BEGIN TRANSACTION").
/// Async context manager entry. Resolves to `self` with no I/O.
fn __aenter__<'py>(slf: Py<Self>, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
pyo3_async_runtimes::tokio::future_into_py(py, async move { Ok(slf) })
}
Comment on lines +204 to +206

@David-Engel David Engel (David-Engel) Aug 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggestion: cursor(), commit(), and rollback() all raise RuntimeError("Connection is closed"), but async with closed_conn: succeeds and then no-ops on exit. Suggest matching the rest of the surface:

fn __aenter__<'py>(slf: Py<Self>, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
    if slf.borrow(py).tds_client.is_none() {
        return Err(PyRuntimeError::new_err("Connection is closed"));
    }
    pyo3_async_runtimes::tokio::future_into_py(py, async move { Ok(slf) })
}


/// Async context manager exit. Delegates to `close()`; exception info is
/// ignored and never suppressed (`close()` resolves to `None`, which is
/// falsy per Python's `__aexit__` contract).
fn __aexit__<'py>(
&mut self,
py: Python<'py>,
_exc_type: &Bound<'_, PyAny>,
_exc_val: &Bound<'_, PyAny>,
_exc_tb: &Bound<'_, PyAny>,
) -> PyResult<Bound<'py, PyAny>> {
self.close(py)
}

/// Commit the current transaction. If none is open, surfaces SQL Server 3902.
fn commit<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
// Clone the Arc synchronously so the future is `'static + Send`
// without borrowing `self`. Only a shared borrow is required —
// nothing on `self` is mutated here.
// Clone the Arc so the future is `'static + Send`.
let client = self
.tds_client
.as_ref()
Expand All @@ -215,33 +230,18 @@ impl PyAsyncConnection {
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| {
tracing::error!("PyAsyncConnection::commit: failed: {}", e);
// TODO(User Story 47181): map TdsError to a DB-API-compliant exception, preserving SQLSTATE + server error number.
PyRuntimeError::new_err(format!("Commit failed: {e}"))
})?;
guard
.commit_transaction(None, None)
.await
.map_err(|e| map_tds_error("commit", "Commit failed", e))?;
tracing::info!("PyAsyncConnection::commit: transaction committed");
Python::attach(|py| Ok(py.None()))
})
}

/// Roll back the current TDS transaction asynchronously.
///
/// ```python
/// await conn.rollback()
/// ```
///
/// Sends a TM_ROLLBACK (Transaction Manager ROLLBACK) request over the
/// wire and awaits the server's DONE token. Raises `RuntimeError`
/// synchronously if the connection has already been closed.
///
/// If no transaction is currently open on the server, the rollback will
/// fail with the server's own error (SQL Server 3903 — "The ROLLBACK
/// TRANSACTION request has no corresponding BEGIN TRANSACTION").
/// Roll back the current transaction. If none is open, surfaces SQL Server 3903.
fn rollback<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
// Clone the Arc synchronously so the future is `'static + Send`
// without borrowing `self`. Only a shared borrow is required —
// nothing on `self` is mutated here.
// Clone the Arc so the future is `'static + Send`.
let client = self
.tds_client
.as_ref()
Expand All @@ -251,39 +251,23 @@ impl PyAsyncConnection {
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| {
tracing::error!("PyAsyncConnection::rollback: failed: {}", e);
// TODO(User Story 47181): map TdsError to a DB-API-compliant exception, preserving SQLSTATE + server error number.
PyRuntimeError::new_err(format!("Rollback failed: {e}"))
})?;
guard
.rollback_transaction(None, None)
.await
.map_err(|e| map_tds_error("rollback", "Rollback failed", e))?;
tracing::info!("PyAsyncConnection::rollback: transaction rolled back");
Python::attach(|py| Ok(py.None()))
})
}

/// Create an async cursor bound to this connection.
///
/// ```python
/// cur = conn.cursor()
/// await cur.execute("SELECT 1")
/// ```
///
/// This method does not perform I/O — it simply hands out a new
/// [`PyAsyncCursor`] that shares the connection's `TdsClient` via an
/// `Arc<tokio::sync::Mutex<_>>`. Following DB-API 2.0, `cursor()` is a
/// synchronous call; only the cursor's execute/fetch methods will be
/// awaitable.
///
/// Raises `RuntimeError` if the connection has already been closed.
/// A second cursor may be created on the same connection, but both
/// cursors share one TDS wire session and serialize on the same async
/// mutex — matching the non-MARS TDS session model.
/// Sync per DB-API 2.0. A second cursor is allowed; both share the same
/// TDS session and serialize on the same async mutex.
Comment on lines +263 to +264

@David-Engel David Engel (David-Engel) Aug 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggestion: this docstring is now a fragment that says neither what cursor() returns nor that it raises RuntimeError when the connection is closed - both of which are contract that shows up in Python help(). Worth one line back:

/// Create an async cursor. Sync per DB-API 2.0. Raises `RuntimeError` if the
/// connection is closed. A second cursor is allowed; both share the same TDS
/// session and serialize on the same async mutex.

fn cursor(&self) -> PyResult<PyAsyncCursor> {
let client = self
.tds_client
.as_ref()
.ok_or_else(|| PyRuntimeError::new_err("Connection is closed"))?
.clone();
Ok(PyAsyncCursor::new(client))
Ok(PyAsyncCursor::new(client, self.default_query_timeout))
}
}
Loading
Loading