[Draft] Async Query Execution - POC

[subrata-ms]

Description

Related Issues

Checklist

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

[github-actions]

📊 Code Coverage Report

🔥 Diff Coverage 9%

🎯 Overall Coverage 90.5%

📦 Project: mssql-tds + 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 (0.0%): Missing lines 33,37-41,44-45,47,54-60,62-65,67-72,75-80,86-101,103-105,107-109,111,113
• mssql-py-core/src/async_cursor.rs (0.0%): Missing lines 38-43,53-58,61,70-75,80-87,91-93,95-97,101-107,109-116,118-121,125-126,130-134,136,138-139,143-144,146-147,152-154,156-160,162-168,170-175,180-186,188-190,198-200
• mssql-py-core/src/async_runtime.rs (100%)
• mssql-py-core/src/connection.rs (100%)
• mssql-py-core/src/lib.rs (100%)

Summary

• Total: 172 lines
• Missing: 156 lines
• Coverage: 9%

mssql-py-core/src/async_connection.rs

  29     tds_client: Arc<Mutex<TdsClient>>,
  30     is_closed: Arc<AtomicBool>,
  31 }
  32 
! 33 #[pymethods]
  34 impl PyCoreAsyncConnection {
  35     #[new]
  36     #[pyo3(signature = (client_context_dict, _python_logger=None))]
! 37     fn new(
! 38         py: Python<'_>,
! 39         client_context_dict: &Bound<'_, PyDict>,
! 40         _python_logger: Option<&Bound<'_, PyAny>>,
! 41     ) -> PyResult<Self> {
  42         // Reuse the existing sync connection's dict→ClientContext logic so
  43         // the async path stays byte-compatible with `bulkcopy`'s inputs.
! 44         let ctx = PyCoreConnection::dict_to_client_context(client_context_dict)?;
! 45         let datasource = ctx.data_source.clone();
  46 
! 47         tracing::info!(
  48             "PyCoreAsyncConnection::new — connecting to datasource={}",
  49             datasource
  50         );

---

  50         );
  51 
  52         // Drive the async connect on the shared runtime with the GIL released,
  53         // so other Python threads keep making progress during the handshake.
! 54         let handle = shared_runtime().handle().clone();
! 55         let client_result = py.detach(|| {
! 56             handle.block_on(async move {
! 57                 let provider = TdsConnectionProvider {};
! 58                 provider.create_client(ctx, &datasource, None).await
! 59             })
! 60         });
  61 
! 62         let client = client_result.map_err(|e| {
! 63             tracing::error!("PyCoreAsyncConnection::new — connect failed: {}", e);
! 64             PyRuntimeError::new_err(format!("Failed to connect to SQL Server: {e}"))
! 65         })?;
  66 
! 67         tracing::info!("PyCoreAsyncConnection::new — connected");
! 68         Ok(Self {
! 69             tds_client: Arc::new(Mutex::new(client)),
! 70             is_closed: Arc::new(AtomicBool::new(false)),
! 71         })
! 72     }
  73 
  74     /// Return a new async cursor sharing this connection's `TdsClient`.
! 75     fn cursor(&self) -> PyResult<PyCoreAsyncCursor> {
! 76         if self.is_closed.load(Ordering::Acquire) {
! 77             return Err(PyRuntimeError::new_err("Connection is closed"));
! 78         }
! 79         Ok(PyCoreAsyncCursor::new(self.tds_client.clone()))
! 80     }
  81 
  82     /// Best-effort synchronous close.
  83     ///
  84     /// Sends the TDS close packet and tears down the TCP connection on the

---

   82     /// Best-effort synchronous close.
   83     ///
   84     /// Sends the TDS close packet and tears down the TCP connection on the
   85     /// shared runtime. Idempotent: subsequent calls are no-ops.
!  86     fn close(&self, py: Python<'_>) -> PyResult<()> {
!  87         if self.is_closed.swap(true, Ordering::AcqRel) {
!  88             return Ok(());
!  89         }
!  90         let handle = shared_runtime().handle().clone();
!  91         let client = self.tds_client.clone();
!  92         py.detach(|| {
!  93             handle.block_on(async move {
!  94                 let mut guard = client.lock().await;
!  95                 if let Err(e) = guard.close_connection().await {
!  96                     tracing::warn!("PyCoreAsyncConnection::close — {}", e);
!  97                 }
!  98             })
!  99         });
! 100         Ok(())
! 101     }
  102 
! 103     fn is_connected(&self) -> bool {
! 104         !self.is_closed.load(Ordering::Acquire)
! 105     }
  106 
! 107     fn __repr__(&self) -> String {
! 108         if self.is_closed.load(Ordering::Acquire) {
! 109             "PyCoreAsyncConnection(closed)".to_string()
  110         } else {
! 111             "PyCoreAsyncConnection(connected)".to_string()
  112         }
! 113     }
  114 }

---

mssql-py-core/src/async_cursor.rs

  34     active_cancel: SyncMutex<Option<CancelHandle>>,
  35 }
  36 
  37 impl AsyncCursorState {
! 38     fn new() -> Self {
! 39         Self {
! 40             has_resultset: AtomicBool::new(false),
! 41             active_cancel: SyncMutex::new(None),
! 42         }
! 43     }
  44 }
  45 
  46 #[pyclass]
  47 pub struct PyCoreAsyncCursor {

---

  49     state: Arc<AsyncCursorState>,
  50 }
  51 
  52 impl PyCoreAsyncCursor {
! 53     pub(crate) fn new(tds_client: Arc<Mutex<TdsClient>>) -> Self {
! 54         Self {
! 55             tds_client,
! 56             state: Arc::new(AsyncCursorState::new()),
! 57         }
! 58     }
  59 }
  60 
! 61 #[pymethods]
  62 impl PyCoreAsyncCursor {
  63     /// Cooperative cancel — sends a TDS attention for the currently in-flight
  64     /// operation, if any. Safe to call from another asyncio task while an
  65     /// `execute_async` or `fetchone_async` future is being awaited.

---

   66     ///
   67     /// After the future resumes and returns, the server-side cancellation
   68     /// surfaces as an `OperationCancelledError` mapped to a Python exception
   69     /// by [`convert_tds_error`].
!  70     fn cancel(&self) {
!  71         if let Some(handle) = self.state.active_cancel.lock().take() {
!  72             info!("PyCoreAsyncCursor::cancel — dispatching TDS attention");
!  73             handle.cancel();
!  74         }
!  75     }
   76 
   77     /// Execute a T-SQL batch. Returns a Python awaitable that resolves to
   78     /// `None` on success.
   79     #[pyo3(signature = (query, timeout_sec=30))]
!  80     fn execute_async<'py>(
!  81         &self,
!  82         py: Python<'py>,
!  83         query: String,
!  84         timeout_sec: Option<u32>,
!  85     ) -> PyResult<Bound<'py, PyAny>> {
!  86         let tds_client = self.tds_client.clone();
!  87         let state = self.state.clone();
   88 
   89         // Install a fresh cancel handle BEFORE spawning the future so that a
   90         // concurrent `.cancel()` call can find it immediately.
!  91         let parent = CancelHandle::new();
!  92         let child = parent.child_handle();
!  93         *state.active_cancel.lock() = Some(parent);
   94 
!  95         pyo3_async_runtimes::tokio::future_into_py(py, async move {
!  96             let _clear_on_return = ClearActiveCancelOnDrop(state.clone());
!  97             let mut client = tds_client.lock().await;
   98 
   99             // Drain any lingering result set from a prior call before running
  100             // the next batch — matches the sync cursor's contract.
! 101             if let Some(rs) = client.get_current_resultset() {
! 102                 info!("execute_async: closing stale result set");
! 103                 if let Err(e) = rs.close().await {
! 104                     error!("execute_async: closing stale result set failed: {}", e);
! 105                     return Err(convert_tds_error(e));
! 106                 }
! 107             }
  108 
! 109             info!("execute_async: executing query ({} chars)", query.len());
! 110             client
! 111                 .execute(query, timeout_sec, Some(&child))
! 112                 .await
! 113                 .map_err(|e| {
! 114                     error!("execute_async: execute failed: {}", e);
! 115                     convert_tds_error(e)
! 116                 })?;
  117 
! 118             state.has_resultset.store(true, Ordering::Release);
! 119             Python::attach(|py| Ok(py.None()))
! 120         })
! 121     }
  122 
  123     /// Fetch the next row as a Python tuple, or `None` when the result set
  124     /// is exhausted. Idempotent once exhausted.
! 125     fn fetchone_async<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
! 126         let state = self.state.clone();
  127 
  128         // Fast path — no active result set, resolve to None without acquiring
  129         // the connection lock.
! 130         if !state.has_resultset.load(Ordering::Acquire) {
! 131             return pyo3_async_runtimes::tokio::future_into_py(py, async {
! 132                 Python::attach(|py| Ok(py.None()))
! 133             });
! 134         }
  135 
! 136         let tds_client = self.tds_client.clone();
  137 
! 138         pyo3_async_runtimes::tokio::future_into_py(py, async move {
! 139             let mut client = tds_client.lock().await;
  140 
  141             // If the batch was closed between `has_resultset` check and lock
  142             // acquisition, resolve to None.
! 143             let col_count = match client.get_current_resultset() {
! 144                 Some(rs) => rs.get_metadata().len(),
  145                 None => {
! 146                     state.has_resultset.store(false, Ordering::Release);
! 147                     return Python::attach(|py| Ok(py.None()));
  148                 }
  149             };
  150 
  151             // Rebuild the result-set borrow (dropped after `get_metadata` above).
! 152             let rs = client
! 153                 .get_current_resultset()
! 154                 .expect("resultset present under held lock");
  155 
! 156             let mut writer = PyRowWriter::new(col_count);
! 157             let has_row = rs.next_row_into(&mut writer).await.map_err(|e| {
! 158                 error!("fetchone_async: next_row_into failed: {}", e);
! 159                 convert_tds_error(e)
! 160             })?;
  161 
! 162             if !has_row {
! 163                 if let Err(e) = rs.close().await {
! 164                     error!("fetchone_async: close result set failed: {}", e);
! 165                 }
! 166                 state.has_resultset.store(false, Ordering::Release);
! 167                 return Python::attach(|py| Ok(py.None()));
! 168             }
  169 
! 170             Python::attach(|py| {
! 171                 let tup = writer.to_py_tuple(py)?;
! 172                 Ok::<Py<PyAny>, PyErr>(tup.into())
! 173             })
! 174         })
! 175     }
  176 
  177     /// Best-effort close: clears the local result-set flag and dispatches a
  178     /// cancel if any operation is still in flight. Does not drop the
  179     /// connection.
! 180     fn close(&self) -> PyResult<()> {
! 181         self.state.has_resultset.store(false, Ordering::Release);
! 182         if let Some(handle) = self.state.active_cancel.lock().take() {
! 183             handle.cancel();
! 184         }
! 185         Ok(())
! 186     }
  187 
! 188     fn __repr__(&self) -> &'static str {
! 189         "PyCoreAsyncCursor(ready)"
! 190     }
  191 }
  192 
  193 /// Drop guard: on natural (non-panic) completion of an `execute_async` future,
  194 /// clear the active cancel slot so a late `.cancel()` call becomes a no-op

---

  194 /// clear the active cancel slot so a late `.cancel()` call becomes a no-op
  195 /// instead of triggering a spurious attention against the next batch.
  196 struct ClearActiveCancelOnDrop(Arc<AsyncCursorState>);
  197 impl Drop for ClearActiveCancelOnDrop {
! 198     fn drop(&mut self) {
! 199         let _ = self.0.active_cancel.lock().take();
! 200     }
  201 }

---

---

🔗 Quick Links

View Azure DevOps Build · Coverage Report