From a2de1c141e2110290c811478c2b2c6eb67596e7c Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Mon, 29 Jun 2026 16:55:54 +0100 Subject: [PATCH] Add fetch_next_batch --- rust/src/database/postgres/connection.rs | 15 +++ rust/src/database/postgres/cursor_state.rs | 125 ++++++++++++++++--- rust/src/database/postgres/helpers.rs | 33 ++++- tests/synapse_rust/test_database_postgres.py | 115 ++++++++++++++++- 4 files changed, 268 insertions(+), 20 deletions(-) diff --git a/rust/src/database/postgres/connection.rs b/rust/src/database/postgres/connection.rs index 76e6705456..91526c1ba0 100644 --- a/rust/src/database/postgres/connection.rs +++ b/rust/src/database/postgres/connection.rs @@ -288,6 +288,21 @@ impl Cursor { self.with_cursor_state(|state| state.fetch_all(py)) } + /// Fetch the next batch of rows from the current result set. + /// + /// Blocks for the first row, then returns any further rows that are already + /// available without blocking. Returns an empty list only once the result + /// set is exhausted. `capacity` is a hint for the size of the returned + /// buffer, not a limit on the number of rows returned. + #[pyo3(signature = (capacity = 100))] + fn fetch_next_batch<'py>( + &self, + py: Python<'py>, + capacity: usize, + ) -> PyResult>> { + self.with_cursor_state(|state| state.fetch_next_batch(py, capacity)) + } + /// Return the PEP-249 `rowcount` for the last statement. /// /// This is the number of rows affected by a DML statement; for queries diff --git a/rust/src/database/postgres/cursor_state.rs b/rust/src/database/postgres/cursor_state.rs index ec63b96abc..009c32ce1a 100644 --- a/rust/src/database/postgres/cursor_state.rs +++ b/rust/src/database/postgres/cursor_state.rs @@ -16,7 +16,7 @@ use pyo3::{ use tokio_postgres::{Column, RowStream}; use crate::database::postgres::{ - helpers::{BlockingPostgres, BlockingPostgresResult}, + helpers::{BlockingPostgres, BlockingPostgresResult, BlockingPostgresStream as _}, value::pg_row_to_py, }; @@ -71,24 +71,75 @@ impl CursorQueryState { let next = stream.as_mut().next().block_on(py); - match next { - 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; - Err(PyRuntimeError::new_err(format!( - "error fetching row from postgres: {err}" - ))) - } - None => { - self.rowcount = stream.rows_affected(); - self.stream = None; - Ok(None) + self.parse_stream_row(py, next) + } + + /// Fetch the next batch of rows. + /// + /// This method will block on the first row if it's not immediately + /// available, but will return any additional rows that are also ready + /// without blocking. + /// + /// This is a convenience for Python code that wants to avoid the overhead + /// of calling `fetch_one` repeatedly, but still wants to avoid blocking on + /// the entire result set. + /// + /// Will only return an empty batch if the stream is exhausted. + pub fn fetch_next_batch<'py>( + &mut self, + py: Python<'py>, + capacity: usize, + ) -> PyResult>> { + // If there's no live stream, either the previous result set has already + // been fully drained — in which case we return an empty batch, per the + // contract — or nothing has been executed yet, which is an error. We + // tell the two apart by `description`, which is set for the duration of + // a query and only cleared by the next `execute`. + if self.stream.is_none() { + return if self.description.is_some() { + Ok(Vec::new()) + } else { + Err(PyRuntimeError::new_err("no active query")) + }; + } + + let mut stream = self.get_stream_mut()?; + + // Wait for at least one row. + let first_row = stream.block_on_next(py); + + let first_pg_row = self.parse_stream_row(py, first_row)?; + let Some(first_pg_row) = first_pg_row else { + // The stream is exhausted, so we return an empty batch. + return Ok(Vec::new()); + }; + + let mut buffer = Vec::with_capacity(capacity); + buffer.push(first_pg_row); + + loop { + // `parse_stream_row` requires a mutable reference to self, and so + // we can't hold onto a mutable reference to the `stream` across the + // call (as it `stream` is a mutable reference of `self`). So we + // have to get a new mutable reference to the stream after each row + // is parsed. Hopefully the compiler can mostly optimise this away. + stream = self.get_stream_mut()?; + + let Some(next) = stream.get_next_if_ready() else { + // The stream isn't ready yet, so we return what we have (which + // we know is non-empty because we pushed the first row above). + break; + }; + + let row = self.parse_stream_row(py, next)?; + match row { + Some(pg_row) => buffer.push(pg_row), + // End of stream, so we return what we have. + None => break, } } + + Ok(buffer) } /// Collect every remaining row into a `Vec`, draining the stream. @@ -128,6 +179,46 @@ impl CursorQueryState { Ok(PyInt::new(py, rowcount)) } + + /// Parse a row from the stream, handling errors and end-of-stream. + /// + /// On a stream error the state is cleared and the error surfaced to Python. + /// On end-of-stream the rowcount is captured and the stream dropped. + fn parse_stream_row<'py>( + &'_ mut self, + py: Python<'py>, + row: Option>, + ) -> PyResult>> { + match row { + Some(Ok(row)) => { + let pg_row = pg_row_to_py(py, &row)?; + Ok(Some(pg_row)) + } + Some(Err(err)) => { + self.stream = None; + Err(PyRuntimeError::new_err(format!( + "error fetching row from postgres: {err}" + ))) + } + None => { + // The stream is exhausted: capture the rowcount and drop the + // stream so we never poll a completed stream again (doing so + // surfaces as a spurious "connection closed" error). + self.rowcount = self.get_stream_mut()?.rows_affected(); + self.stream = None; + Ok(None) + } + } + } + + /// Get a mutable reference to the row stream, or an error if there is no + /// active query. + fn get_stream_mut(&mut self) -> PyResult> { + let Some(stream) = self.stream.as_mut() else { + return Err(PyRuntimeError::new_err("no active query")); + }; + Ok(stream.as_mut()) + } } /// Consume and discard every row of a stream, propagating any error. Used to diff --git a/rust/src/database/postgres/helpers.rs b/rust/src/database/postgres/helpers.rs index 73b7ed850e..90f71ac0d2 100644 --- a/rust/src/database/postgres/helpers.rs +++ b/rust/src/database/postgres/helpers.rs @@ -6,8 +6,9 @@ //! runtime's own connection task can run) while we wait. The [`Ungil`] bounds //! are what let us hand the future across the `detach` boundary. -use std::future::Future; +use std::{future::Future, pin::Pin}; +use futures::{FutureExt, StreamExt}; use pyo3::{marker::Ungil, PyResult, Python}; use crate::database::{postgres::pg_err_to_py, runtime::runtime}; @@ -52,3 +53,33 @@ where F::Output: Ungil + Send, { } + +pub trait BlockingPostgresStream +where + Self: futures::Stream + Sized + Send + Ungil + Unpin, + Self::Item: Ungil + Send, +{ + /// Get the next item from the stream, blocking on the shared runtime if + /// necessary. + /// + /// If the stream is not ready to yield an item, this will release the GIL + /// and block until the next item is available. + /// + /// This method will return `None` if the stream is exhausted. + fn block_on_next(&mut self, py: Python<'_>) -> Option { + match self.get_next_if_ready() { + Some(row) => row, + None => self.next().block_on(py), + } + } + + /// Get the next item from the stream if it's ready, without blocking. + /// + /// Returns `None` if the stream is not ready to yield an item. Returns + /// `Some(None)` if the stream is exhausted. + fn get_next_if_ready(&mut self) -> Option> { + self.next().now_or_never() + } +} + +impl BlockingPostgresStream for Pin<&mut tokio_postgres::RowStream> {} diff --git a/tests/synapse_rust/test_database_postgres.py b/tests/synapse_rust/test_database_postgres.py index 6f5ed6058b..48f971d0b1 100644 --- a/tests/synapse_rust/test_database_postgres.py +++ b/tests/synapse_rust/test_database_postgres.py @@ -36,8 +36,7 @@ from tests.utils import ( def _build_dsn() -> str: - """Build a libpq keyword/value connection string from the test config. - """ + """Build a libpq keyword/value connection string from the test config.""" parts = [f"dbname={POSTGRES_BASE_DB}"] if POSTGRES_USER is not None: @@ -177,6 +176,118 @@ class PostgresConnectionTestCase(unittest.TestCase): self.assertEqual(self.conn.run_interaction(interaction), [1, 2, 3]) + # ------------------------------------------------------------------ + # fetch_next_batch() + # ------------------------------------------------------------------ + + def test_fetch_next_batch_returns_first_row(self) -> None: + """A non-empty result set yields a batch containing at least the first + row, which blocks until available.""" + + def interaction(cursor: Any) -> list[Any]: + cursor.execute("SELECT 42::int, 'hello'::text") + return cursor.fetch_next_batch() + + # The batch must be non-empty and start with the first row. We don't + # assert the exact length: how many further rows are already buffered + # (and so returned without blocking) is timing-dependent. + batch = self.conn.run_interaction(interaction) + self.assertEqual(batch[0], (42, "hello")) + + def test_fetch_next_batch_empty_when_no_rows(self) -> None: + """An empty result set yields an empty batch.""" + + def interaction(cursor: Any) -> list[Any]: + cursor.execute("SELECT 1 WHERE false") + return cursor.fetch_next_batch() + + self.assertEqual(self.conn.run_interaction(interaction), []) + + def test_fetch_next_batch_collects_all_rows_across_batches(self) -> None: + """Looping until an empty batch is returned yields every row exactly + once, in order, regardless of how the rows are split across batches.""" + + def interaction(cursor: Any) -> list[Any]: + cursor.execute( + """ + SELECT id FROM generate_series(1, 1000) AS s(id) ORDER BY id + """ + ) + rows = [] + while True: + batch = cursor.fetch_next_batch() + if not batch: + break + rows.extend(batch) + return rows + + self.assertEqual( + self.conn.run_interaction(interaction), + [(n,) for n in range(1, 1001)], + ) + + def test_fetch_next_batch_empty_after_exhaustion(self) -> None: + """Once the result set is drained, further batches are empty.""" + + def interaction(cursor: Any) -> list[Any]: + cursor.execute("SELECT 1") + first = cursor.fetch_next_batch() + second = cursor.fetch_next_batch() + return [first, second] + + first, second = self.conn.run_interaction(interaction) + self.assertEqual(first, [(1,)]) + self.assertEqual(second, []) + + def test_fetch_next_batch_capacity_is_not_a_limit(self) -> None: + """`capacity` is only a buffer hint; a batch may exceed it.""" + + def interaction(cursor: Any) -> list[Any]: + cursor.execute( + "SELECT id FROM generate_series(1, 100) AS s(id) ORDER BY id" + ) + rows = [] + while True: + batch = cursor.fetch_next_batch(1) + if not batch: + break + rows.extend(batch) + return rows + + self.assertEqual( + self.conn.run_interaction(interaction), + [(n,) for n in range(1, 101)], + ) + + def test_fetch_next_batch_without_query_raises(self) -> None: + """Calling fetch_next_batch before execute is an error.""" + + def interaction(cursor: Any) -> None: + cursor.fetch_next_batch() + + with self.assertRaises(RuntimeError): + self.conn.run_interaction(interaction) + + def test_fetch_next_batch_interleaves_with_fetch_one(self) -> None: + """fetch_one and fetch_next_batch share the same underlying stream, so + rows already consumed by one are not seen by the other.""" + + def interaction(cursor: Any) -> list[Any]: + cursor.execute("SELECT id FROM generate_series(1, 10) AS s(id) ORDER BY id") + first = cursor.fetch_one() + rows = [first] + while True: + batch = cursor.fetch_next_batch() + if not batch: + break + rows.extend(batch) + return rows + + self.assertEqual( + self.conn.run_interaction(interaction), + [(n,) for n in range(1, 11)], + ) + # ------------------------------------------------------------------ # Value round-tripping (ToSql + FromSql for each supported type) # ------------------------------------------------------------------