diff --git a/rust/src/database/postgres/connection.rs b/rust/src/database/postgres/connection.rs index 9f3845ea37..76e6705456 100644 --- a/rust/src/database/postgres/connection.rs +++ b/rust/src/database/postgres/connection.rs @@ -53,9 +53,7 @@ use pyo3::{ use tokio_postgres::Client; use crate::database::postgres::{ - cursor_state::CursorQueryState, - helpers::BlockingPostgresResult, - value::{pg_row_to_py, PgValue}, + cursor_state::CursorQueryState, helpers::BlockingPostgresResult, value::PgValue, }; /// A single Postgres connection exposed to Python. @@ -221,6 +219,15 @@ impl Cursor { f(inner) } + /// Run `f` against the live query state, or error if the cursor has + /// already been finished/closed (its state taken out, leaving `None`). + pub fn with_cursor_state( + &self, + f: impl FnOnce(&mut CursorQueryState) -> PyResult, + ) -> PyResult { + self.with_inner(|inner| f(&mut inner.query_state)) + } + /// Finish the transaction — `COMMIT` if `commit` is set, otherwise /// `ROLLBACK` — and hand the client back to the connection. /// @@ -273,19 +280,12 @@ impl Cursor { /// Return the next row of the current result set, or `None` if exhausted. fn fetch_one<'py>(&self, py: Python<'py>) -> PyResult>> { - let Some(row) = self.with_inner(|inner| inner.query_state.fetch_one(py))? else { - return Ok(None); - }; - - let py_row = pg_row_to_py(py, &row)?; - Ok(Some(py_row)) + self.with_cursor_state(|state| state.fetch_one(py)) } /// Drain and return all remaining rows of the current result set. fn fetch_all<'py>(&self, py: Python<'py>) -> PyResult>> { - let rows = self.with_inner(|inner| inner.query_state.fetch_all(py))?; - - rows.into_iter().map(|row| pg_row_to_py(py, &row)).collect() + self.with_cursor_state(|state| state.fetch_all(py)) } /// Return the PEP-249 `rowcount` for the last statement. @@ -293,13 +293,7 @@ impl Cursor { /// This is the number of rows affected by a DML statement; for queries /// where it isn't (yet) known it follows PEP-249 and returns `-1`. fn rowcount<'py>(&self, py: Python<'py>) -> PyResult> { - let Some(rowcount) = self.with_inner(|inner| inner.query_state.rowcount(py))? else { - // If we don't have a rowcount yet, PEP-249 says we should return - // -1. - return Ok((-1i64).into_pyobject(py)?); - }; - - Ok(rowcount.into_pyobject(py)?) + self.with_cursor_state(|state| state.rowcount(py)) } } diff --git a/rust/src/database/postgres/cursor_state.rs b/rust/src/database/postgres/cursor_state.rs index 0442c10127..ec63b96abc 100644 --- a/rust/src/database/postgres/cursor_state.rs +++ b/rust/src/database/postgres/cursor_state.rs @@ -8,10 +8,17 @@ use std::pin::Pin; use futures::{StreamExt, TryStreamExt}; -use pyo3::{exceptions::PyRuntimeError, PyResult, Python}; +use pyo3::{ + exceptions::PyRuntimeError, + types::{PyInt, PyTuple}, + Bound, PyResult, Python, +}; use tokio_postgres::{Column, RowStream}; -use crate::database::postgres::helpers::{BlockingPostgres, BlockingPostgresResult}; +use crate::database::postgres::{ + helpers::{BlockingPostgres, BlockingPostgresResult}, + value::pg_row_to_py, +}; /// The state carried over from the cursor's most recent query. #[derive(Default)] @@ -57,7 +64,7 @@ impl CursorQueryState { /// /// On exhaustion the rowcount is captured and the stream dropped; on a /// stream error the state is cleared and the error surfaced to Python. - pub fn fetch_one<'py>(&mut self, py: Python<'py>) -> PyResult> { + pub fn fetch_one<'py>(&mut self, py: Python<'py>) -> PyResult>> { let Some(stream) = self.stream.as_mut() else { return Err(PyRuntimeError::new_err("no active query")); }; @@ -65,7 +72,10 @@ impl CursorQueryState { let next = stream.as_mut().next().block_on(py); match next { - Some(Ok(row)) => Ok(Some(row)), + Some(Ok(row)) => { + let pg_row = pg_row_to_py(py, &row)?; + Ok(Some(pg_row)) + } Some(Err(err)) => { self.stream = None; self.description = None; @@ -82,12 +92,16 @@ impl CursorQueryState { } /// Collect every remaining row into a `Vec`, draining the stream. - pub fn fetch_all<'py>(&mut self, py: Python<'py>) -> PyResult> { + pub fn fetch_all<'py>(&mut self, py: Python<'py>) -> PyResult>> { let Some(stream) = self.stream.as_mut() else { return Err(PyRuntimeError::new_err("no active query")); }; let rows = stream.try_collect::>().block_on_result(py)?; + let rows = rows + .into_iter() + .map(|row| pg_row_to_py(py, &row)) + .collect::>>()?; self.rowcount = stream.rows_affected(); self.stream = None; @@ -96,7 +110,7 @@ impl CursorQueryState { } /// Return the affected-row count, draining the stream first if needed. - pub fn rowcount<'py>(&mut self, py: Python<'py>) -> PyResult> { + pub fn rowcount<'py>(&mut self, py: Python<'py>) -> PyResult> { // `stream.rows_affected()` is only valid after the stream is // drained, so we need to drain it here. This is OK as in Python the // rowcount should only be accessed for queries that DO NOT return @@ -106,7 +120,13 @@ impl CursorQueryState { self.rowcount = stream.rows_affected(); } - Ok(self.rowcount) + let Some(rowcount) = self.rowcount else { + // If we don't have a rowcount yet, PEP-249 says we should return + // -1 + return Ok(PyInt::new(py, -1)); + }; + + Ok(PyInt::new(py, rowcount)) } }