Add fetch_next_batch

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<Vec<Bound<'py, PyTuple>>> {
+        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

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<Vec<Bound<'py, PyTuple>>> {
+        // 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<Result<tokio_postgres::Row, tokio_postgres::Error>>,
+    ) -> PyResult<Option<Bound<'py, PyTuple>>> {
+        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<Pin<&mut RowStream>> {
+        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

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<Self::Item> {
+        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<Option<Self::Item>> {
+        self.next().now_or_never()
+    }
+}
+
+impl BlockingPostgresStream for Pin<&mut tokio_postgres::RowStream> {}

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)
     # ------------------------------------------------------------------