From f3e1947e0ea72b7933cccb4b0b84089f2aeeb175 Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Mon, 29 Jun 2026 19:41:10 +0000 Subject: [PATCH] Add CursorQueryState result-set state machine Model the lifecycle of a cursor's most recent query as an explicit state machine (Idle / Active / Closed) so that illegal field combinations are unrepresentable and fetching past the end of a result set is a clean, specific error rather than the spurious "connection closed" you get from re-polling a finished stream. The row stream is fused so an exhausted-but-not-yet-reported stream can sit safely in `Active`. This change adds `fetch_one`, `fetch_all` and `rowcount` (the batched `fetch_next_batch` follows separately). On a stream error a cursor resets to `Idle`; on normal exhaustion it moves to `Closed`, retaining the PEP-249 rowcount from the command tag. The state machine is generic over the stream type, defaulting to `RowStream`, via a small `CursorRowStream` trait that abstracts the three things the logic needs (the affected-row count, row->PyTuple conversion, and error rendering). This is what lets the state transitions, exhaustion handling and error recovery be unit-tested against an in-memory fake stream, with no live Postgres server. `cursor_state` is `pub` for now so its not-yet-consumed items don't trip clippy's dead_code lint; the connection code wires it up and tightens that later. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/src/database/postgres/cursor_state.rs | 608 +++++++++++++++++++++ rust/src/database/postgres/mod.rs | 1 + 2 files changed, 609 insertions(+) create mode 100644 rust/src/database/postgres/cursor_state.rs diff --git a/rust/src/database/postgres/cursor_state.rs b/rust/src/database/postgres/cursor_state.rs new file mode 100644 index 0000000000..3e0adc24c3 --- /dev/null +++ b/rust/src/database/postgres/cursor_state.rs @@ -0,0 +1,608 @@ +//! Tracks the result of the cursor's most recent `execute`. +//! +//! [`tokio_postgres`] returns rows as a [`RowStream`] that is consumed lazily, +//! so a cursor only ever holds onto the *current* query's stream plus the +//! metadata (column names, rowcount) derived from it. +//! +//! The lifecycle is modelled as an explicit state machine ([`CursorQueryState`]) +//! so that illegal field combinations are unrepresentable, and so that fetching +//! after the result set has been reported exhausted is a clean error rather +//! than undefined behaviour (re-polling a finished stream surfaces as a +//! spurious "connection closed"). The row stream is [`Fuse`]d, which lets an +//! exhausted-but-not-yet-reported stream sit safely in the `Active` state: a +//! fused stream keeps yielding `None` instead of re-polling its completed +//! inner stream. +//! +//! ## Generic over the stream +//! +//! The state machine is generic over the underlying stream type `S`, defaulting +//! to [`RowStream`] (the only stream used in production). The +//! [`CursorRowStream`] trait abstracts the three things the logic needs from a +//! stream beyond [`futures::Stream`] itself — the affected-row count, how to +//! turn a row into a Python tuple, and how to render a stream error as a +//! `PyErr`. This is what lets the state-machine logic (exhaustion handling, +//! state transitions, error recovery) be unit-tested against an in-memory fake +//! stream with no live Postgres server; see the tests at the bottom of the +//! module. + +use std::{mem, pin::Pin}; + +use futures::{stream::Fuse, StreamExt, TryStreamExt}; +use pyo3::{ + exceptions::PyRuntimeError, + marker::Ungil, + types::{PyInt, PyTuple}, + Bound, PyErr, PyResult, Python, +}; +use tokio_postgres::RowStream; + +use crate::database::postgres::{helpers::BlockingPostgres, value::pg_row_to_py}; + +/// The capabilities the cursor state machine needs from the underlying row +/// stream, beyond [`futures::Stream`] itself. +/// +/// Implemented for [`tokio_postgres::RowStream`] in production; the unit tests +/// provide an in-memory implementation so the state-machine logic can be +/// exercised without a live database. +/// +/// The bounds (`Send` + [`Ungil`] on the stream and its row/error types) are +/// what let the rows and errors cross the GIL-release boundary in the +/// `block_on*` helpers. +pub trait CursorRowStream: + futures::Stream> + Send + Ungil + Sized +{ + /// A single result row. + type Row: Send + Ungil; + /// The stream's error type. + type Error: Send + Ungil; + + /// The command tag's affected-row count, valid only once the stream has + /// been fully drained. Maps to [`RowStream::rows_affected`]. + fn rows_affected(&self) -> Option; + + /// Convert a yielded row into a Python tuple, one element per column. + fn row_to_py<'py>(py: Python<'py>, row: &Self::Row) -> PyResult>; + + /// Build the Python error raised when the stream yields an error. + fn stream_err(err: &Self::Error) -> PyErr; +} + +impl CursorRowStream for RowStream { + type Row = tokio_postgres::Row; + type Error = tokio_postgres::Error; + + fn rows_affected(&self) -> Option { + RowStream::rows_affected(self) + } + + fn row_to_py<'py>(py: Python<'py>, row: &Self::Row) -> PyResult> { + pg_row_to_py(py, row) + } + + fn stream_err(err: &Self::Error) -> PyErr { + PyRuntimeError::new_err(format!("error fetching row from postgres: {err}")) + } +} + +/// A live row stream, fused so that polling it after completion returns `None` +/// rather than re-polling the finished inner stream. +type FusedStream = Pin>>; + +/// The lifecycle of the cursor's most recent query. +pub enum CursorQueryState { + /// No query has been executed yet, or the previous result set was reset by + /// the next `execute`. Fetching is an error. + Idle, + /// A query is in flight; rows can be fetched. Exhaustion has not yet been + /// reported to the caller. + Active { + /// Live row stream for the current query. + stream: FusedStream, + /// Column names for the result set (empty for a DML statement). + /// + /// TODO: currently write-only; kept to back a future PEP-249 + /// `Cursor.description` accessor. + description: Vec, + }, + /// The result set has been fully consumed and exhaustion reported — a fetch + /// returned `None`/`[]`, or `fetch_all`/`rowcount` drained it. Any further + /// `fetch_*` is a programming error; `rowcount` still returns the count. + Closed { + /// Column names for the result set, carried over from `Active` so they + /// survive once the rows are gone. + /// + /// TODO: currently write-only; like `Active::description` it is kept to + /// back a future PEP-249 `Cursor.description` accessor. `#[allow]`d + /// until that reader lands rather than dropped, so the column metadata + /// isn't silently lost at the `Active` -> `Closed` transition. + #[allow(dead_code)] + description: Vec, + /// PEP-249 `rowcount` from the command tag, if it was captured. + rowcount: Option, + }, +} + +// Implemented by hand rather than derived: a derived `Default` would add an +// `S: Default` bound (which `RowStream` doesn't satisfy), even though the +// default state (`Idle`) doesn't depend on `S` at all. The hand-written impl is +// therefore *not* equivalent to the derive clippy suggests, so we silence it. +#[allow(clippy::derivable_impls)] +impl Default for CursorQueryState { + fn default() -> Self { + Self::Idle + } +} + +impl CursorQueryState { + /// A fresh state with no query yet run. + pub fn new() -> Self { + Self::Idle + } + + /// Reset to `Idle`, discarding any previous result set. Called at the + /// start of every `execute`. + pub fn new_query(&mut self) { + *self = Self::Idle; + } + + /// Record the stream and column metadata for a newly-started query. + /// + /// `description` is the list of column names (empty for a DML statement); + /// taking it as a plain `Vec` rather than the driver's column + /// metadata keeps this layer decoupled from `tokio_postgres::Column` and so + /// testable. + pub fn on_query_start(&mut self, stream: S, description: Vec) { + *self = Self::Active { + stream: Box::pin(stream.fuse()), + description, + }; + } + + /// Pull the next row from the stream, or `None` once it's exhausted. + /// + /// Returning `None` reports exhaustion and moves the cursor to `Closed`, so + /// a subsequent fetch is an error. On a stream error the cursor is reset to + /// `Idle` and the error surfaced to Python. + pub fn fetch_one<'py>(&mut self, py: Python<'py>) -> PyResult>> { + let Self::Active { + stream, + description, + } = self + else { + return Err(self.fetch_after_end_err()); + }; + + // Unlike `fetch_next_batch` (which uses `block_on_next` to grab any + // already-buffered rows without releasing the GIL), a single fetch has + // to wait for the one row either way, so we block directly rather than + // bothering with the non-blocking fast path. + match stream.as_mut().next().block_on(py) { + Some(Ok(row)) => Ok(Some(S::row_to_py(py, &row)?)), + Some(Err(err)) => { + *self = Self::Idle; + Err(S::stream_err(&err)) + } + None => { + let rowcount = rows_affected(stream); + *self = Self::Closed { + description: mem::take(description), + rowcount, + }; + Ok(None) + } + } + } + + /// Collect every remaining row into a `Vec`, draining the stream. + /// + /// Draining reports exhaustion and moves the cursor to `Closed`, so a + /// subsequent fetch is an error. On a stream error the cursor is reset to + /// `Idle` and the error surfaced to Python. + pub fn fetch_all<'py>(&mut self, py: Python<'py>) -> PyResult>> { + let Self::Active { + stream, + description, + } = self + else { + return Err(self.fetch_after_end_err()); + }; + + let rows = match stream.as_mut().try_collect::>().block_on(py) { + Ok(rows) => rows, + Err(err) => { + *self = Self::Idle; + return Err(S::stream_err(&err)); + } + }; + let rows = rows + .into_iter() + .map(|row| S::row_to_py(py, &row)) + .collect::>>()?; + + let rowcount = rows_affected(stream); + *self = Self::Closed { + description: mem::take(description), + rowcount, + }; + + Ok(rows) + } + + /// Return the affected-row count, draining the stream first if needed. + /// + /// Unlike the `fetch_*` methods this is always valid: reading the rowcount + /// of an already-exhausted (`Closed`) cursor returns the captured count + /// rather than erroring, per PEP-249. + pub fn rowcount<'py>(&mut self, py: Python<'py>) -> PyResult> { + // `rows_affected()` is only valid after the stream is drained, so we + // drain it here. This is OK as in Python the rowcount should only be + // accessed for queries that DO NOT return rows, e.g. INSERT, UPDATE, + // DELETE. + if let Self::Active { + stream, + description, + } = self + { + if let Err(err) = drain_stream(stream.as_mut()).block_on(py) { + *self = Self::Idle; + return Err(S::stream_err(&err)); + } + let rowcount = rows_affected(stream); + *self = Self::Closed { + description: mem::take(description), + rowcount, + }; + } + + match self { + Self::Closed { + rowcount: Some(rowcount), + .. + } => Ok(PyInt::new(py, *rowcount)), + // No rowcount known: either we're `Idle` (no query run, or reset by + // an earlier error), or `Closed` with a command tag that carried no + // count. The stream is always drained by the block above before we + // get here, so "not yet drained" isn't a case. PEP-249 says -1. + _ => Ok(PyInt::new(py, -1)), + } + } + + /// The error to raise when a `fetch_*` method finds no rows available, + /// distinguishing "no query was ever run" from "the result set has already + /// been exhausted". Only meaningful for the non-`Active` states. + fn fetch_after_end_err(&self) -> PyErr { + match self { + Self::Closed { .. } => { + PyRuntimeError::new_err("cannot fetch: the result set is already exhausted") + } + _ => PyRuntimeError::new_err("no active query"), + } + } +} + +/// The command tag's affected-row count, valid only once the (fused) stream has +/// been fully drained. +fn rows_affected(stream: &FusedStream) -> Option { + // The first `get_ref` unwraps the `Pin`, the second the `Fuse`. + stream.as_ref().get_ref().get_ref().rows_affected() +} + +/// Consume and discard every row of a stream, propagating any error. Used to +/// reach the trailing command-complete message that carries the rowcount. +async fn drain_stream(mut stream: Pin<&mut Fuse>) -> Result<(), S::Error> { + while let Some(row) = stream.next().await { + row?; + } + Ok(()) +} + +#[cfg(test)] +mod tests { + //! The state machine is driven here by an in-memory [`FakeStream`] instead + //! of a live [`RowStream`], so these tests cover the exhaustion handling, + //! the state transitions, and the error recovery without any Postgres + //! server. Rows are simple integer tuples; that's enough to assert the + //! plumbing carries values through `row_to_py`. + //! + //! Note the one thing they deliberately *don't* cover: [`FakeStream`]'s + //! `poll_next` is always `Poll::Ready`, so the GIL-release / blocking path + //! in the `block_on*` helpers is never exercised here (it has its own tests + //! in `helpers`). These tests are purely about the state-machine logic. + + use std::{ + collections::VecDeque, + task::{Context, Poll}, + }; + + use pyo3::prelude::*; + + use super::*; + + /// A result row: a list of integer cells. + #[derive(Clone, Debug, PartialEq)] + struct FakeRow(Vec); + + /// A stream error carrying a message we can assert on. + #[derive(Debug)] + struct FakeError(&'static str); + + /// An in-memory stream of pre-canned items, plus the `rows_affected` value + /// the command tag would report once it's drained. + struct FakeStream { + items: VecDeque>, + rows_affected: Option, + } + + impl futures::Stream for FakeStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + // Every item is immediately available; `None` once drained. + Poll::Ready(self.items.pop_front()) + } + } + + impl CursorRowStream for FakeStream { + type Row = FakeRow; + type Error = FakeError; + + fn rows_affected(&self) -> Option { + self.rows_affected + } + + fn row_to_py<'py>(py: Python<'py>, row: &FakeRow) -> PyResult> { + PyTuple::new(py, &row.0) + } + + fn stream_err(err: &FakeError) -> PyErr { + PyRuntimeError::new_err(format!("fake stream error: {}", err.0)) + } + } + + /// Build an `Active` state over the given rows (all `Ok`) with the given + /// command-tag rowcount. + fn active_with( + rows: Vec>, + rows_affected: Option, + ) -> CursorQueryState { + let mut state = CursorQueryState::::new(); + state.on_query_start( + FakeStream { + items: rows.into_iter().map(|r| Ok(FakeRow(r))).collect(), + rows_affected, + }, + vec!["col".to_string()], + ); + state + } + + /// Assert a `Bound` decodes to the given integer cells. + fn assert_tuple(tuple: &Bound<'_, PyTuple>, expected: &[i64]) { + let got: Vec = tuple.extract().unwrap(); + assert_eq!(got, expected); + } + + #[test] + fn fetch_one_yields_rows_then_reports_and_closes() { + Python::initialize(); + Python::attach(|py| { + let mut state = active_with(vec![vec![1], vec![2]], Some(0)); + + assert_tuple(&state.fetch_one(py).unwrap().unwrap(), &[1]); + assert_tuple(&state.fetch_one(py).unwrap().unwrap(), &[2]); + // The end of the stream is reported as `None`, closing the cursor. + assert!(state.fetch_one(py).unwrap().is_none()); + assert!(matches!(state, CursorQueryState::Closed { .. })); + + // A fetch after exhaustion is a clean, specific error. + let err = state.fetch_one(py).unwrap_err(); + assert!(err.to_string().contains("already exhausted"), "{err}"); + }); + } + + #[test] + fn fetch_one_on_idle_errors_with_no_active_query() { + Python::initialize(); + Python::attach(|py| { + let mut state = CursorQueryState::::new(); + let err = state.fetch_one(py).unwrap_err(); + assert!(err.to_string().contains("no active query"), "{err}"); + }); + } + + #[test] + fn fetch_all_drains_and_closes() { + Python::initialize(); + Python::attach(|py| { + let mut state = active_with(vec![vec![1, 2], vec![3, 4]], Some(0)); + + let rows = state.fetch_all(py).unwrap(); + assert_eq!(rows.len(), 2); + assert_tuple(&rows[0], &[1, 2]); + assert_tuple(&rows[1], &[3, 4]); + assert!(matches!(state, CursorQueryState::Closed { .. })); + + // Draining again is an error: the result set is gone. + assert!(state + .fetch_all(py) + .unwrap_err() + .to_string() + .contains("exhausted")); + }); + } + + #[test] + fn fetch_all_on_empty_stream_returns_empty_and_closes() { + Python::initialize(); + Python::attach(|py| { + let mut state = active_with(vec![], Some(0)); + assert!(state.fetch_all(py).unwrap().is_empty()); + assert!(matches!(state, CursorQueryState::Closed { .. })); + }); + } + + #[test] + fn fetch_one_after_fetch_all_errors() { + Python::initialize(); + Python::attach(|py| { + let mut state = active_with(vec![vec![1]], Some(0)); + let _ = state.fetch_all(py).unwrap(); + assert!(state + .fetch_one(py) + .unwrap_err() + .to_string() + .contains("exhausted")); + }); + } + + #[test] + fn fetch_one_then_fetch_all_continues_partially_consumed_stream() { + Python::initialize(); + Python::attach(|py| { + let mut state = active_with(vec![vec![1], vec![2], vec![3]], Some(0)); + + // Consume one row with `fetch_one`... + assert_tuple(&state.fetch_one(py).unwrap().unwrap(), &[1]); + // ...then `fetch_all` should pick up the *rest*, not restart. + let rest = state.fetch_all(py).unwrap(); + assert_eq!(rest.len(), 2); + assert_tuple(&rest[0], &[2]); + assert_tuple(&rest[1], &[3]); + assert!(matches!(state, CursorQueryState::Closed { .. })); + }); + } + + #[test] + fn fetch_all_error_resets_to_idle() { + Python::initialize(); + Python::attach(|py| { + // First row ok, then an error part-way through the collection. + let mut state = CursorQueryState::::new(); + state.on_query_start( + FakeStream { + items: VecDeque::from(vec![Ok(FakeRow(vec![1])), Err(FakeError("kaboom"))]), + rows_affected: Some(0), + }, + vec!["col".to_string()], + ); + + let err = state.fetch_all(py).unwrap_err(); + assert!(err.to_string().contains("kaboom"), "{err}"); + // A mid-stream error resets to `Idle`, exactly like `fetch_one`. + assert!(matches!(state, CursorQueryState::Idle)); + assert!(state + .fetch_all(py) + .unwrap_err() + .to_string() + .contains("no active query")); + }); + } + + #[test] + fn rowcount_after_drain_error_resets_to_idle() { + Python::initialize(); + Python::attach(|py| { + let mut state = CursorQueryState::::new(); + state.on_query_start( + FakeStream { + items: VecDeque::from(vec![Err(FakeError("drain failed"))]), + rows_affected: Some(7), + }, + vec!["col".to_string()], + ); + + // Draining for the rowcount hits the error and surfaces it. + let err = state.rowcount(py).unwrap_err(); + assert!(err.to_string().contains("drain failed"), "{err}"); + // The cursor is reset to `Idle`, so no stale count is retained: a + // subsequent rowcount is the PEP-249 "unknown" sentinel, -1. + assert!(matches!(state, CursorQueryState::Idle)); + assert_eq!(state.rowcount(py).unwrap().extract::().unwrap(), -1); + }); + } + + #[test] + fn stream_error_resets_to_idle_and_surfaces() { + Python::initialize(); + Python::attach(|py| { + // A stream whose second item is an error. + let mut state = CursorQueryState::::new(); + state.on_query_start( + FakeStream { + items: VecDeque::from(vec![Ok(FakeRow(vec![1])), Err(FakeError("boom"))]), + rows_affected: None, + }, + vec!["col".to_string()], + ); + + assert_tuple(&state.fetch_one(py).unwrap().unwrap(), &[1]); + let err = state.fetch_one(py).unwrap_err(); + assert!(err.to_string().contains("boom"), "{err}"); + + // After an error the cursor is reset to `Idle`, so the next fetch + // reports "no active query" rather than "exhausted". + assert!(matches!(state, CursorQueryState::Idle)); + assert!(state + .fetch_one(py) + .unwrap_err() + .to_string() + .contains("no active query")); + }); + } + + #[test] + fn rowcount_is_minus_one_before_any_query() { + Python::initialize(); + Python::attach(|py| { + let mut state = CursorQueryState::::new(); + assert_eq!(state.rowcount(py).unwrap().extract::().unwrap(), -1); + }); + } + + #[test] + fn rowcount_drains_and_reports_command_tag() { + Python::initialize(); + Python::attach(|py| { + // A DML-style result: no rows, but a command tag of 5 affected rows. + let mut state = active_with(vec![], Some(5)); + assert_eq!(state.rowcount(py).unwrap().extract::().unwrap(), 5); + assert!(matches!(state, CursorQueryState::Closed { .. })); + + // Reading it again on a `Closed` cursor returns the retained count, + // rather than erroring like a `fetch_*` would. + assert_eq!(state.rowcount(py).unwrap().extract::().unwrap(), 5); + }); + } + + #[test] + fn rowcount_drains_remaining_rows_of_a_select() { + Python::initialize(); + Python::attach(|py| { + // Even with rows pending, asking for the rowcount drains them and + // then reports the command tag. + let mut state = active_with(vec![vec![1], vec![2]], Some(2)); + assert_eq!(state.rowcount(py).unwrap().extract::().unwrap(), 2); + // Stream is drained, so fetching is now an "exhausted" error. + assert!(state + .fetch_one(py) + .unwrap_err() + .to_string() + .contains("exhausted")); + }); + } + + #[test] + fn new_query_resets_an_active_cursor() { + Python::initialize(); + Python::attach(|py| { + let mut state = active_with(vec![vec![1]], Some(0)); + state.new_query(); + assert!(matches!(state, CursorQueryState::Idle)); + assert!(state + .fetch_one(py) + .unwrap_err() + .to_string() + .contains("no active query")); + }); + } +} diff --git a/rust/src/database/postgres/mod.rs b/rust/src/database/postgres/mod.rs index 3710f8b972..1513743cb3 100644 --- a/rust/src/database/postgres/mod.rs +++ b/rust/src/database/postgres/mod.rs @@ -16,6 +16,7 @@ use pyo3::types::PyModule; // stops clippy's `dead_code` lint from firing on them before the // cursor/connection code (added in later changes) wires them up; the visibility // is tightened back to private once that happens. +pub mod cursor_state; pub mod helpers; pub mod value;