mirror of
https://github.com/element-hq/synapse.git
synced 2026-08-14 18:01:13 +00:00
Add Cursor.executemany to the Rust Postgres backend
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) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.8
parent
020bfbaf82
commit
6910464b8a
@@ -54,6 +54,7 @@
|
||||
|
||||
use std::sync::{Arc, Mutex, MutexGuard, TryLockError};
|
||||
|
||||
use futures::future::try_join_all;
|
||||
use pyo3::{
|
||||
exceptions::PyRuntimeError,
|
||||
prelude::*,
|
||||
@@ -389,6 +390,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<Vec<PgValue>>,
|
||||
) -> 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::<u64>())
|
||||
})?;
|
||||
|
||||
// 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
|
||||
|
||||
@@ -366,6 +366,19 @@ impl<S: CursorRowStream> CursorQueryState<S> {
|
||||
}
|
||||
}
|
||||
|
||||
/// 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<u64>) {
|
||||
*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.
|
||||
@@ -850,6 +863,39 @@ mod tests {
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn on_command_complete_retains_rowcount_without_a_result_set() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let mut state = CursorQueryState::<FakeStream>::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::<i64>().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 mut state = CursorQueryState::<FakeStream>::new();
|
||||
// An empty `executemany` batch: nothing ran, so no count.
|
||||
state.on_command_complete(None);
|
||||
assert_eq!(state.rowcount(py).unwrap().extract::<i64>().unwrap(), -1);
|
||||
assert!(state.description().is_none());
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_query_resets_an_active_cursor() {
|
||||
Python::initialize();
|
||||
|
||||
@@ -167,6 +167,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()
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user