From c159abd149e7cb89e96f2c46c8da03881d52b3eb Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Thu, 2 Jul 2026 11:11:19 +0000 Subject: [PATCH] Add Cursor.executemany to the Rust Postgres backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implement DBAPI2's `executemany`, which Synapse uses for batched writes. The statement is prepared once and run for each parameter set inside the connection's (lazily opened) transaction, so a failure part-way aborts the whole batch. The per-set executions are pipelined — their futures are driven concurrently so tokio_postgres streams the batch onto the connection in one round-trip rather than one per statement. Afterwards `rowcount` reports the total rows affected across all executions (as psycopg2 does) via a new fetchless "command complete" cursor state; there is no result set to fetch or describe. An empty parameter sequence is a no-op that leaves rowcount at -1. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/src/database/postgres/connection.rs | 60 +++++++++++ rust/src/database/postgres/cursor_state.rs | 48 +++++++++ tests/synapse_rust/test_database_postgres.py | 101 +++++++++++++++++++ 3 files changed, 209 insertions(+) diff --git a/rust/src/database/postgres/connection.rs b/rust/src/database/postgres/connection.rs index a9184a969e..61b3713ef8 100644 --- a/rust/src/database/postgres/connection.rs +++ b/rust/src/database/postgres/connection.rs @@ -54,6 +54,7 @@ use std::sync::{Arc, Mutex, MutexGuard, TryLockError}; +use futures::future::try_join_all; use pyo3::{ exceptions::PyRuntimeError, prelude::*, @@ -394,6 +395,65 @@ impl Cursor { Ok(()) } + /// Execute `query` once for each parameter set in `params_seq`. + /// + /// This is DBAPI2's `executemany`, used by Synapse for batched writes. The + /// statement is `prepare`d once and then run for each parameter set; it + /// produces no fetchable rows. Like a single `execute`, the whole batch runs + /// inside the connection's (lazily opened) transaction, so a failure + /// part-way through aborts it and leaves the caller to roll back. + /// + /// The per-parameter-set executions are *pipelined*: their futures are + /// driven concurrently, so `tokio_postgres` streams the whole batch onto the + /// connection without waiting for a round-trip between each — one round-trip + /// for the batch rather than one per statement. Results are matched to + /// requests in order, and on the first error the rest are abandoned (the + /// transaction is aborted anyway). + /// + /// After it returns `rowcount` reports the total number of rows affected + /// across all executions (as psycopg2 does), `description` is `None`, and + /// any `fetch_*` is an error. An empty `params_seq` runs nothing at all — + /// no statement is sent and no transaction is opened — and leaves + /// `rowcount` at the PEP-249 "unknown" sentinel (`-1`). + #[pyo3(signature = (query, params_seq))] + fn executemany( + &self, + py: Python<'_>, + query: &str, + params_seq: Vec>, + ) -> PyResult<()> { + // Drop any previous result set before starting the new statement. + self.lock_state()?.new_query(); + + // An empty batch is a no-op (matching psycopg2): don't send anything or + // open a transaction, and leave `rowcount` reporting "unknown". + if params_seq.is_empty() { + self.lock_state()?.on_command_complete(None); + return Ok(()); + } + + let total = self.connection.with_client(py, |client| { + // Prepare once, then build a future per parameter set. Driving them + // concurrently is what makes `tokio_postgres` pipeline them onto the + // connection; blocking on the joined future runs the whole batch. + let statement = client.prepare(query).block_on_result(py)?; + + let counts = try_join_all( + params_seq + .into_iter() + .map(|params| client.execute_raw(&statement, params)), + ) + .block_on_result(py)?; + + Ok(counts.into_iter().sum::()) + })?; + + // Retain the summed affected-row count for `rowcount`, as psycopg2 does. + self.lock_state()?.on_command_complete(Some(total)); + + Ok(()) + } + /// Execute a multi-statement SQL script (statements separated by `;`). /// /// Unlike [`Cursor::execute`], which `prepare`s a single statement, this diff --git a/rust/src/database/postgres/cursor_state.rs b/rust/src/database/postgres/cursor_state.rs index 6a17207f8f..537788641b 100644 --- a/rust/src/database/postgres/cursor_state.rs +++ b/rust/src/database/postgres/cursor_state.rs @@ -366,6 +366,19 @@ impl CursorQueryState { } } + /// Record a completed command that produced no fetchable rows, retaining + /// its affected-row count for `rowcount`. + /// + /// Used by `executemany`, which runs a statement repeatedly for its side + /// effects: there is no result set to fetch (`description` is `None` and + /// any `fetch_*` errors as exhausted), but the (summed) rowcount is kept. + pub fn on_command_complete(&mut self, rowcount: Option) { + *self = Self::Closed { + description: Vec::new(), + rowcount, + }; + } + /// 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. @@ -890,6 +903,41 @@ mod tests { }); } + #[test] + fn on_command_complete_retains_rowcount_without_a_result_set() { + Python::initialize(); + Python::attach(|py| { + let _guard = enter_runtime(); + let mut state = CursorQueryState::::new(); + // Simulate a completed `executemany` that affected 5 rows. + state.on_command_complete(Some(5)); + + assert!(matches!(state, CursorQueryState::Closed { .. })); + // The rowcount is retained, but there is nothing to describe or + // fetch. + assert_eq!(state.rowcount(py).unwrap().extract::().unwrap(), 5); + assert!(state.description().is_none()); + assert!(state + .fetch_one(py) + .unwrap_err() + .to_string() + .contains("exhausted")); + }); + } + + #[test] + fn on_command_complete_with_no_count_reports_minus_one() { + Python::initialize(); + Python::attach(|py| { + let _guard = enter_runtime(); + let mut state = CursorQueryState::::new(); + // An empty `executemany` batch: nothing ran, so no count. + state.on_command_complete(None); + assert_eq!(state.rowcount(py).unwrap().extract::().unwrap(), -1); + assert!(state.description().is_none()); + }); + } + #[test] fn new_query_resets_an_active_cursor() { Python::initialize(); diff --git a/tests/synapse_rust/test_database_postgres.py b/tests/synapse_rust/test_database_postgres.py index 0b4aa331fc..de962b3b3f 100644 --- a/tests/synapse_rust/test_database_postgres.py +++ b/tests/synapse_rust/test_database_postgres.py @@ -180,6 +180,107 @@ class PostgresConnectionTestCase(unittest.TestCase): self.assertEqual(run_interaction(self.conn, interaction), [1, 2, 3]) + # ------------------------------------------------------------------ + # executemany() + # ------------------------------------------------------------------ + + def test_executemany_runs_once_per_param_set(self) -> None: + """executemany applies the statement for each parameter set.""" + + def interaction(cursor: Any) -> list[Any]: + cursor.execute("CREATE TEMP TABLE em (id int, name text)") + cursor.executemany( + "INSERT INTO em (id, name) VALUES ($1, $2)", + [[1, "a"], [2, "b"], [3, "c"]], + ) + cursor.execute("SELECT id, name FROM em ORDER BY id") + return cursor.fetch_all() + + self.assertEqual( + run_interaction(self.conn, interaction), + [(1, "a"), (2, "b"), (3, "c")], + ) + + def test_executemany_rowcount_is_total_affected(self) -> None: + """rowcount after executemany is the sum across all executions.""" + + def interaction(cursor: Any) -> int: + cursor.execute("CREATE TEMP TABLE em (id int)") + cursor.executemany("INSERT INTO em VALUES ($1)", [[1], [2], [3]]) + return cursor.rowcount() + + self.assertEqual(run_interaction(self.conn, interaction), 3) + + def test_executemany_leaves_no_result_set(self) -> None: + """executemany produces nothing to describe.""" + + def describe(cursor: Any) -> Any: + cursor.execute("CREATE TEMP TABLE em (id int)") + cursor.executemany("INSERT INTO em VALUES ($1)", [[1], [2]]) + return cursor.description() + + self.assertIsNone(run_interaction(self.conn, describe)) + + def test_executemany_fetch_after_raises(self) -> None: + """There is no result set after executemany, so fetching is an error.""" + + def interaction(cursor: Any) -> Any: + cursor.execute("CREATE TEMP TABLE em (id int)") + cursor.executemany("INSERT INTO em VALUES ($1)", [[1], [2]]) + cursor.fetch_one() + + with self.assertRaises(RuntimeError): + run_interaction(self.conn, interaction) + + def test_executemany_empty_is_noop(self) -> None: + """An empty parameter sequence runs nothing and affects no rows.""" + + def interaction(cursor: Any) -> int: + cursor.execute("CREATE TEMP TABLE em (id int)") + cursor.executemany("INSERT INTO em VALUES ($1)", []) + # Nothing ran, so rowcount is the "unknown" sentinel... + self.assertEqual(cursor.rowcount(), -1) + # ...and the table is untouched. + cursor.execute("SELECT count(*) FROM em") + row = cursor.fetch_one() + assert row is not None + return row[0] + + self.assertEqual(run_interaction(self.conn, interaction), 0) + + def test_executemany_rolls_back_on_error(self) -> None: + """A failure part-way through executemany aborts the transaction, so + no rows from the batch survive.""" + + table = "rust_pg_test_executemany_rollback" + try: + + def interaction(cursor: Any) -> None: + cursor.execute(f"CREATE TABLE {table} (id int PRIMARY KEY)") + # The third set duplicates the first, violating the primary key. + cursor.executemany( + f"INSERT INTO {table} VALUES ($1)", + [[1], [2], [1]], + ) + + with self.assertRaises(postgres.IntegrityError): + run_interaction(self.conn, interaction) + + # The CREATE TABLE and the whole batch were in one transaction that + # rolled back, so the table should not exist. + def table_exists(cursor: Any) -> Any: + cursor.execute("SELECT to_regclass($1)::text", [table]) + row = cursor.fetch_one() + assert row is not None + return row[0] + + self.assertIsNone(run_interaction(self.conn, table_exists)) + finally: + run_interaction( + self.conn, + lambda cursor: cursor.execute(f"DROP TABLE IF EXISTS {table}"), + ) + # ------------------------------------------------------------------ # fetch_next_batch() # ------------------------------------------------------------------