From 61b80bd0030fb26ed9cfbb27c8b1c85285273b60 Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Sun, 18 Oct 2026 20:06:23 +0000 Subject: [PATCH] Queue writes in a Rust write transaction and send them together A `WriteTxn` sent each statement and waited for it before the next, so a transaction of forty writes took forty round trips, though nothing read most of their results. Writes can now be queued with `execute_later`, `insert_many_later` and `upsert_many_later`, which take owned parameters. Every statement that is not queued, and the `COMMIT`, first sends the queued writes, in order, in the same round trip: tokio-postgres sends a request when its future is first polled, and `join_in_order` polls them in order. So no read can overtake a queued write, and a transaction ending in queued writes sends them with its `COMMIT`. Statements not yet prepared on the connection are prepared together first. tokio-postgres encodes a request's parameters on its first poll and sends nothing if that fails, while later requests still go out. So the queued writes' parameters are checked against their prepared statements' types before any is sent: otherwise the writes after a bad one, and the COMMIT, would go through. A dry run sends its queue before rolling back, so it still checks the writes. The statement methods must not run concurrently, or a statement could overtake writes queued before it. A queued write that fails surfaces at the next statement, reported as its own failure (with an optional mapping, for callers that turn a constraint violation into a domain error), rather than as the "current transaction is aborted" error of the statements after it. A failed attempt drops its queue unsent. `execute_batch` uses `join_in_order` too: `try_join_all` does not promise to poll long lists in order. Nothing queues writes yet. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01PVLsNh6akt4ifEizY4D1ta --- rust/src/database/postgres/write.rs | 733 +++++++++++++++++++++++----- 1 file changed, 608 insertions(+), 125 deletions(-) diff --git a/rust/src/database/postgres/write.rs b/rust/src/database/postgres/write.rs index 0b0b0097eb..b39464dc6f 100644 --- a/rust/src/database/postgres/write.rs +++ b/rust/src/database/postgres/write.rs @@ -28,18 +28,32 @@ //! column as an array and selecting from `UNNEST` of them. Statements that //! need a per-row sub-select go through [`WriteTxn::execute_batch`], which //! pipelines them on the connection. +//! +//! Round trips are what a write transaction mostly waits on, so it sends as +//! much as it can together: +//! +//! - The `BEGIN` goes out with the first statement. +//! - A write whose result nobody reads can be queued with one of the +//! `*_later` methods. Queued writes go out, in order, in the same round +//! trip as the next statement that is not queued, or as the `COMMIT`. A +//! failure surfaces there, reported as the queued statement's own. +//! - Statements not yet prepared on the connection are prepared together. use std::{ cell::RefCell, - fmt::Write as _, future::Future, - sync::atomic::{AtomicBool, AtomicU8, Ordering}, + pin::Pin, + sync::{ + atomic::{AtomicBool, AtomicU8, Ordering}, + Mutex, + }, + task::Poll, time::{Duration, Instant}, }; use async_trait::async_trait; use futures::{ - future::{join, try_join_all, BoxFuture}, + future::{join, maybe_done, poll_fn, BoxFuture}, FutureExt, }; use log::warn; @@ -217,8 +231,49 @@ impl<'a> Column<'a> { } } +/// A parameter of a queued statement, which owns its value. +pub type OwnedParam = Box; + +/// One column of a batch for [`WriteTxn::insert_many_later`] and +/// [`WriteTxn::upsert_many_later`]: like [`Column`], but owning its values, +/// as the statement is sent after the call returns. +pub struct OwnedColumn { + name: &'static str, + sql_type: &'static str, + values: OwnedParam, + rows: usize, +} + +impl OwnedColumn { + pub fn new(name: &'static str, sql_type: &'static str, values: Vec) -> Self + where + Vec: ToSql + Sync + Send + 'static, + { + OwnedColumn { + name, + sql_type, + rows: values.len(), + values: Box::new(values), + } + } +} + +/// How to turn the error of a queued statement into the one its caller +/// wants, e.g. a constraint violation into a domain error. +pub type MapError = fn(WriteError) -> WriteError; + +/// A statement queued by one of the `*_later` methods. +struct Queued { + sql: String, + params: Vec, + map_error: Option, +} + /// An open transaction. See the module docs. Rolled back if dropped before /// [`WriteTxn::commit`]. +/// +/// Its statement methods must not run concurrently: a statement could then +/// be sent before writes queued ahead of it. pub struct WriteTxn { conn: Option, in_transaction: bool, @@ -228,6 +283,10 @@ pub struct WriteTxn { /// The `BEGIN` failed, so the statement sent behind it may have run /// outside a transaction: the connection is discarded, not reused. begin_failed: AtomicBool, + /// Writes whose results nobody reads, not yet sent. They go out, in + /// order, ahead of the next statement that is not queued, or of the + /// `COMMIT`, in the same round trip. Never locked across an await. + queued: Mutex>, /// A statement failed, so Postgres has aborted the transaction. A /// `COMMIT` would then roll back without an error, so the closure must /// not be able to report success after catching the error. Holds one of @@ -247,6 +306,7 @@ impl WriteTxn { in_transaction: true, begin_pending: AtomicBool::new(true), begin_failed: AtomicBool::new(false), + queued: Mutex::new(Vec::new()), failed: AtomicU8::new(NOT_FAILED), } } @@ -261,34 +321,24 @@ impl WriteTxn { }) } - /// Await `request`, a statement's first request, sending `begin` in front - /// of it. `tokio_postgres` queues a request when its future is first - /// polled and `join` polls the `BEGIN` first, so Postgres runs the - /// statement inside the transaction. - async fn after_begin( - &self, - begin: Option>>, - request: impl Future>, - context: impl FnOnce() -> String, - ) -> Result { - let result = match begin { - None => request.await, - Some(begin) => { - let (begun, result) = join(begin, request).await; - if let Err(e) = begun { - self.begin_failed.store(true, Ordering::Relaxed); - return Err(self.statement_failed("opening a transaction", e)); - } - result - } - }; - result.map_err(|e| self.statement_failed(&context(), e)) - } - fn conn(&self) -> &PooledConnection { self.conn.as_ref().expect("held until dropped") } + fn take_queued(&self) -> Vec { + std::mem::take(&mut *self.queued.lock().expect("not poisoned")) + } + + /// The error for a statement that could not be sent, such as one whose + /// parameters do not encode. The transaction must not be committed after + /// it. + fn client_failed(&self, context: &str, error: anyhow::Error) -> WriteError { + let _ = + self.failed + .compare_exchange(NOT_FAILED, FAILED, Ordering::Relaxed, Ordering::Relaxed); + WriteError::Other(error.context(context.to_owned())) + } + /// The error for a failed statement, noting that the transaction is now /// aborted. fn statement_failed(&self, context: &str, error: tokio_postgres::Error) -> WriteError { @@ -306,36 +356,151 @@ impl WriteTxn { error } - /// Run a statement that returns rows. + /// Send the queued writes, then the request `last` builds for the + /// statement `sql` (prepared, if there is one), all in one round trip, + /// with the `BEGIN` in front if it is still to be sent. Statements not + /// yet prepared on this connection are prepared first, together, in one + /// more round trip. + /// + /// Requests go on the wire in this order, as tokio-postgres sends each + /// when its future is first polled and [`join_in_order`] polls them in + /// order. A failure is reported as the first one on the wire: the + /// requests after it fail only because it aborted the transaction. + async fn send<'s, T, F>( + &'s self, + sql: Option<&str>, + last: impl FnOnce(&'s PooledConnection, Option) -> F, + context: impl FnOnce() -> String, + ) -> Result + where + F: Future> + 's, + { + let conn = self.conn(); + let queued = self.take_queued(); + let mut begin = self.take_begin(); + + let sqls: Vec<&str> = queued.iter().map(|q| q.sql.as_str()).chain(sql).collect(); + let mut statements: Vec> = + sqls.iter().map(|sql| conn.cached(sql)).collect(); + + let mut missing: Vec<&str> = sqls + .iter() + .zip(&statements) + .filter(|(_, statement)| statement.is_none()) + .map(|(sql, _)| *sql) + .collect(); + missing.sort_unstable(); + missing.dedup(); + if !missing.is_empty() { + let prepares = join_in_order(missing.iter().map(|sql| conn.prepare_cached(sql, false))); + let prepared = match begin.take() { + None => prepares.await, + Some(begin) => { + let (begun, prepared) = join(begin, prepares).await; + if let Err(e) = begun { + self.begin_failed.store(true, Ordering::Relaxed); + return Err(self.statement_failed("opening a transaction", e)); + } + prepared + } + }; + for (sql, result) in missing.iter().zip(prepared) { + let statement = result.map_err(|e| { + self.statement_failed(&format!("preparing {}", describe(sql)), e) + })?; + for (slot, slot_sql) in statements.iter_mut().zip(&sqls) { + if slot_sql == sql { + *slot = Some(statement.clone()); + } + } + } + } + + let mut statements = statements.into_iter().map(|s| s.expect("prepared above")); + let queued_statements: Vec = + statements.by_ref().take(queued.len()).collect(); + let last_statement = statements.next(); + + // tokio-postgres encodes a request's parameters when it is first + // polled, and if that fails it sends nothing, while the requests + // after it still go out, `COMMIT` included. So check that every + // queued write's parameters encode before sending any of them. + for (q, statement) in queued.iter().zip(&queued_statements) { + if let Err(error) = check_params(statement, &q.params) { + return Err(self.client_failed(&describe(&q.sql), error)); + } + } + + let writes = join_in_order(queued.iter().zip(&queued_statements).map(|(q, statement)| { + conn.execute_raw( + statement, + q.params.iter().map(|p| p.as_ref() as &(dyn ToSql + Sync)), + ) + })); + let requests = join(writes, last(conn, last_statement)); + let (written, result) = match begin { + None => requests.await, + Some(begin) => { + let (begun, results) = join(begin, requests).await; + if let Err(e) = begun { + self.begin_failed.store(true, Ordering::Relaxed); + return Err(self.statement_failed("opening a transaction", e)); + } + results + } + }; + + for (q, outcome) in queued.iter().zip(written) { + if let Err(e) = outcome { + let error = self.statement_failed(&describe(&q.sql), e); + return Err(match q.map_error { + Some(map_error) => map_error(error), + None => error, + }); + } + } + result.map_err(|e| self.statement_failed(&context(), e)) + } + + /// Run a statement that returns rows, after the queued writes. pub async fn query( &self, sql: &str, params: &[&(dyn ToSql + Sync)], ) -> Result, WriteError> { - let (statement, begin) = self.prepare(sql).await?; - self.after_begin(begin, self.conn().query(&statement, params), || { - describe(sql) - }) + self.send( + Some(sql), + |conn, statement| async move { + conn.query(&statement.expect("has a statement"), params) + .await + }, + || describe(sql), + ) .await } - /// Run a statement, returning the number of rows it changed. + /// Run a statement, after the queued writes, returning the number of rows + /// it changed. pub async fn execute( &self, sql: &str, params: &[&(dyn ToSql + Sync)], ) -> Result { - let (statement, begin) = self.prepare(sql).await?; - self.after_begin(begin, self.conn().execute(&statement, params), || { - describe(sql) - }) + self.send( + Some(sql), + |conn, statement| async move { + conn.execute(&statement.expect("has a statement"), params) + .await + }, + || describe(sql), + ) .await } /// Run one statement once per set of parameters, pipelined on the - /// connection, returning the total number of rows changed. For the - /// statements that need a sub-select per row; the rest can use - /// [`WriteTxn::insert_many`]. + /// connection after the queued writes, returning the total number of + /// rows changed. For the statements that need a sub-select per row; the + /// rest can use [`WriteTxn::insert_many`]. pub async fn execute_batch( &self, sql: &str, @@ -344,16 +509,99 @@ impl WriteTxn { if rows.is_empty() { return Ok(0); } - let (statement, begin) = self.prepare(sql).await?; - let conn = self.conn(); - let counts = self - .after_begin( - begin, - try_join_all(rows.iter().map(|params| conn.execute(&statement, params))), - || describe(sql), - ) - .await?; - Ok(counts.into_iter().sum()) + self.send( + Some(sql), + |conn, statement| async move { + let statement = statement.expect("has a statement"); + let counts = + join_in_order(rows.iter().map(|params| conn.execute(&statement, params))).await; + counts.into_iter().sum::>() + }, + || describe(sql), + ) + .await + } + + /// Send the queued writes now, and wait for them. + pub async fn flush(&self) -> Result<(), WriteError> { + if self.queued.lock().expect("not poisoned").is_empty() { + return Ok(()); + } + self.send(None, |_, _| async { Ok(()) }, String::new).await + } + + /// Queue a write whose result nobody reads. It goes out ahead of the + /// next statement that is not queued, or of the `COMMIT`, and a failure + /// surfaces there. + pub fn execute_later(&self, sql: impl Into, params: Vec) { + self.queue(sql.into(), params, None); + } + + /// As [`WriteTxn::execute_later`], with `map_error` applied to the + /// statement's error. + pub fn execute_later_mapping_error( + &self, + sql: impl Into, + params: Vec, + map_error: MapError, + ) { + self.queue(sql.into(), params, Some(map_error)); + } + + fn queue(&self, sql: String, params: Vec, map_error: Option) { + self.queued.lock().expect("not poisoned").push(Queued { + sql, + params, + map_error, + }); + } + + /// As [`WriteTxn::insert_many`], queued as [`WriteTxn::execute_later`]. + pub fn insert_many_later( + &self, + table: &str, + columns: Vec, + on_conflict: Option<&str>, + ) -> Result<(), WriteError> { + let shapes: Vec<(&str, &str, usize)> = columns + .iter() + .map(|c| (c.name, c.sql_type, c.rows)) + .collect(); + if check_batch(table, &shapes)?.is_none() { + return Ok(()); + } + let mut sql = unnest_insert_sql(table, &shapes); + if let Some(on_conflict) = on_conflict { + sql.push(' '); + sql.push_str(on_conflict); + } + self.execute_later(sql, columns.into_iter().map(|c| c.values).collect()); + Ok(()) + } + + /// As [`WriteTxn::upsert_many`], queued as [`WriteTxn::execute_later`]. + pub fn upsert_many_later( + &self, + table: &str, + keys: Vec, + values: Vec, + ) -> Result<(), WriteError> { + let shapes: Vec<(&str, &str, usize)> = keys + .iter() + .chain(&values) + .map(|c| (c.name, c.sql_type, c.rows)) + .collect(); + if check_batch(table, &shapes)?.is_none() { + return Ok(()); + } + let key_names: Vec<&str> = keys.iter().map(|c| c.name).collect(); + let value_names: Vec<&str> = values.iter().map(|c| c.name).collect(); + let sql = unnest_insert_sql(table, &shapes) + &on_conflict_update(&key_names, &value_names); + self.execute_later( + sql, + keys.into_iter().chain(values).map(|c| c.values).collect(), + ); + Ok(()) } /// Insert a batch of rows in one statement, returning the number @@ -365,11 +613,14 @@ impl WriteTxn { columns: &[Column<'_>], on_conflict: Option<&str>, ) -> Result { - let Some(rows) = batch_rows(table, columns)? else { + let shapes: Vec<(&str, &str, usize)> = columns + .iter() + .map(|c| (c.name, c.sql_type, c.rows)) + .collect(); + if check_batch(table, &shapes)?.is_none() { return Ok(0); - }; - debug_assert!(rows > 0); - let mut sql = unnest_insert(table, columns); + } + let mut sql = unnest_insert_sql(table, &shapes); if let Some(on_conflict) = on_conflict { sql.push(' '); sql.push_str(on_conflict); @@ -390,54 +641,21 @@ impl WriteTxn { values: &[Column<'_>], ) -> Result { let columns: Vec<&Column<'_>> = keys.iter().chain(values).collect(); - let Some(_) = batch_rows_of(table, &columns)? else { + let shapes: Vec<(&str, &str, usize)> = columns + .iter() + .map(|c| (c.name, c.sql_type, c.rows)) + .collect(); + if check_batch(table, &shapes)?.is_none() { return Ok(0); - }; - let mut sql = unnest_insert_of(table, &columns); - let key_names: Vec<&str> = keys.iter().map(|c| c.name).collect(); - if values.is_empty() { - let _ = write!(sql, " ON CONFLICT ({}) DO NOTHING", key_names.join(", ")); - } else { - let updates: Vec = values - .iter() - .map(|c| format!("{0} = EXCLUDED.{0}", c.name)) - .collect(); - let _ = write!( - sql, - " ON CONFLICT ({}) DO UPDATE SET {}", - key_names.join(", "), - updates.join(", ") - ); } + let key_names: Vec<&str> = keys.iter().map(|c| c.name).collect(); + let value_names: Vec<&str> = values.iter().map(|c| c.name).collect(); + let sql = unnest_insert_sql(table, &shapes) + &on_conflict_update(&key_names, &value_names); let params: Vec<&(dyn ToSql + Sync)> = columns.iter().map(|c| c.values).collect(); self.execute(&sql, ¶ms).await } - /// The prepared statement for `sql`, and the transaction's `BEGIN` if - /// it is still to be sent in front of the statement's execution. On a - /// cache miss the `BEGIN` goes with the prepare instead. - async fn prepare( - &self, - sql: &str, - ) -> Result< - ( - tokio_postgres::Statement, - Option>>, - ), - WriteError, - > { - let begin = self.take_begin(); - if let Some(statement) = self.conn().cached(sql) { - return Ok((statement, begin)); - } - let statement = self - .after_begin(begin, self.conn().prepare_cached(sql, false), || { - format!("preparing {}", describe(sql)) - }) - .await?; - Ok((statement, None)) - } - + /// Send the queued writes and commit, in one round trip. async fn commit(mut self) -> Result<(), WriteError> { match self.failed.load(Ordering::Relaxed) { NOT_FAILED => {} @@ -457,22 +675,28 @@ impl WriteTxn { ))); } } - if self.begin_pending.load(Ordering::Relaxed) { + if self.begin_pending.load(Ordering::Relaxed) + && self.queued.lock().expect("not poisoned").is_empty() + { // No statement ran, so there is no transaction to commit. self.in_transaction = false; return Ok(()); } - self.conn() - .batch_execute("COMMIT") - .await - .map_err(|e| WriteError::from_pg("committing", e))?; + self.send( + None, + |conn, _| conn.batch_execute("COMMIT"), + || "committing".to_owned(), + ) + .await?; self.in_transaction = false; Ok(()) } /// Roll back now, rather than on a background task when dropped, so a - /// retry does not race the rollback for the rows. + /// retry does not race the rollback for the rows. Queued writes are + /// dropped unsent. async fn rollback(mut self) { + self.take_queued(); if self.begin_pending.load(Ordering::Relaxed) { // No statement ran, so there is no transaction to roll back. self.in_transaction = false; @@ -495,6 +719,64 @@ impl WriteTxn { } } +/// Check that `params` encode as `statement`'s parameter types, as +/// tokio-postgres would when sending it. +fn check_params( + statement: &tokio_postgres::Statement, + params: &[OwnedParam], +) -> anyhow::Result<()> { + let types = statement.params(); + if types.len() != params.len() { + anyhow::bail!( + "statement takes {} parameters but {} were given", + types.len(), + params.len() + ); + } + let mut buf = bytes::BytesMut::new(); + for (i, (param, ty)) in params.iter().zip(types).enumerate() { + buf.clear(); + param + .to_sql_checked(ty, &mut buf) + .map_err(|e| anyhow::anyhow!("error serializing parameter {i}: {e}"))?; + } + Ok(()) +} + +/// Await `futures` together, polling every one still pending, in order, on +/// each wake-up. tokio-postgres sends a request when its future is first +/// polled, so the requests go out in the order given. (`join_all` and +/// `try_join_all` do not promise an order for long lists.) +pub(crate) fn join_in_order( + futures: impl IntoIterator, +) -> impl Future> { + let mut futures: Vec<_> = futures + .into_iter() + .map(|future| maybe_done(Box::pin(future))) + .collect(); + poll_fn(move |cx| { + let mut done = true; + for future in futures.iter_mut() { + if Pin::new(future).poll(cx).is_pending() { + done = false; + } + } + if !done { + return Poll::Pending; + } + Poll::Ready( + futures + .iter_mut() + .map(|future| { + Pin::new(future) + .take_output() + .expect("polled to completion above") + }) + .collect(), + ) + }) +} + /// For the read helpers shared with the read-only paths, such as the chain /// cover reads in `persist::reader`. #[async_trait] @@ -671,8 +953,10 @@ where Ok(value) } Ok(value) => { + // A dry run still checks its queued writes. + let flushed = txn.flush().await; txn.rollback().await; - Ok(value) + flushed.map(|()| value) } Err(error) => { txn.rollback().await; @@ -683,38 +967,29 @@ where /// The batch's row count, or `None` for an empty batch. All columns must /// have one value per row. -fn batch_rows(table: &str, columns: &[Column<'_>]) -> Result, WriteError> { - batch_rows_of(table, &columns.iter().collect::>()) -} - -fn batch_rows_of(table: &str, columns: &[&Column<'_>]) -> Result, WriteError> { - let Some(first) = columns.first() else { +/// The batch's row count, given each column's `(name, type, rows)`, or +/// `None` for an empty batch. All columns must have the same number of rows. +fn check_batch(table: &str, columns: &[(&str, &str, usize)]) -> Result, WriteError> { + let Some(&(first_name, _, first_rows)) = columns.first() else { return Err(anyhow::anyhow!("a batch for {table} with no columns").into()); }; - if let Some(column) = columns.iter().find(|c| c.rows != first.rows) { + if let Some(&(name, _, rows)) = columns.iter().find(|c| c.2 != first_rows) { return Err(anyhow::anyhow!( - "a batch for {table} with {} values of {} and {} of {}", - first.rows, - first.name, - column.rows, - column.name + "a batch for {table} with {first_rows} values of {first_name} and {rows} of {name}" ) .into()); } - Ok((first.rows > 0).then_some(first.rows)) + Ok((first_rows > 0).then_some(first_rows)) } -/// `INSERT INTO table (a, b) SELECT * FROM UNNEST($1::text[], $2::bigint[])`. -fn unnest_insert(table: &str, columns: &[Column<'_>]) -> String { - unnest_insert_of(table, &columns.iter().collect::>()) -} - -fn unnest_insert_of(table: &str, columns: &[&Column<'_>]) -> String { - let names: Vec<&str> = columns.iter().map(|c| c.name).collect(); +/// `INSERT INTO table (a, b) SELECT * FROM UNNEST($1::text[], $2::bigint[])`, +/// given each column's `(name, type, rows)`. +fn unnest_insert_sql(table: &str, columns: &[(&str, &str, usize)]) -> String { + let names: Vec<&str> = columns.iter().map(|c| c.0).collect(); let arrays: Vec = columns .iter() .enumerate() - .map(|(i, c)| format!("${}::{}[]", i + 1, c.sql_type)) + .map(|(i, c)| format!("${}::{}[]", i + 1, c.1)) .collect(); format!( "INSERT INTO {table} ({}) SELECT * FROM UNNEST({})", @@ -723,6 +998,23 @@ fn unnest_insert_of(table: &str, columns: &[&Column<'_>]) -> String { ) } +/// The `ON CONFLICT` clause of an upsert: do nothing without `values`, +/// otherwise update them from the new row. +fn on_conflict_update(keys: &[&str], values: &[&str]) -> String { + if values.is_empty() { + return format!(" ON CONFLICT ({}) DO NOTHING", keys.join(", ")); + } + let updates: Vec = values + .iter() + .map(|name| format!("{name} = EXCLUDED.{name}")) + .collect(); + format!( + " ON CONFLICT ({}) DO UPDATE SET {}", + keys.join(", "), + updates.join(", ") + ) +} + /// The start of a statement, for error messages. fn describe(sql: &str) -> String { sql.split_whitespace().take(6).collect::>().join(" ") @@ -818,6 +1110,193 @@ mod tests { .collect() } + /// Queued writes go out ahead of the next statement, which sees them, + /// and with the `COMMIT` when nothing follows them. + #[test] + fn queued_writes_are_sent_in_order() { + let Some((rt, pool, table)) = setup() else { + return; + }; + rt.block_on(async { + let seen = run_write_interaction(&pool, "test", |txn, _| { + let table = table.clone(); + async move { + txn.insert_many_later( + &table, + vec![ + OwnedColumn::new("id", "text", vec!["a".to_owned(), "b".to_owned()]), + OwnedColumn::new("n", "bigint", vec![1i64, 2]), + ], + None, + )?; + txn.execute_later( + format!("UPDATE {table} SET n = n + 10 WHERE id = $1"), + vec![Box::new("a".to_owned())], + ); + let rows = txn + .query(&format!("SELECT id, n FROM {table} ORDER BY id"), &[]) + .await?; + // Only queued: sent with the COMMIT. + txn.upsert_many_later( + &table, + vec![OwnedColumn::new("id", "text", vec!["c".to_owned()])], + vec![OwnedColumn::new("n", "bigint", vec![3i64])], + )?; + Ok(rows + .iter() + .map(|row| (row.get::<_, String>(0), row.get::<_, i64>(1))) + .collect::>()) + } + .boxed() + }) + .await + .unwrap(); + assert_eq!(seen, vec![("a".to_owned(), 11), ("b".to_owned(), 2)]); + assert_eq!( + rows(&pool, &table).await, + vec![ + ("a".to_owned(), 11, None, None), + ("b".to_owned(), 2, None, None), + ("c".to_owned(), 3, None, None), + ] + ); + }); + teardown(&rt, &pool, &table); + } + + /// A queued write that fails surfaces at the next statement, as itself + /// (the next statement only fails because it aborted the transaction), + /// mapped as asked. Nothing is committed. + #[test] + fn queued_write_failure_surfaces_at_the_next_statement() { + let Some((rt, pool, table)) = setup() else { + return; + }; + rt.block_on(async { + let error = run_write_interaction(&pool, "test", |txn, _| { + let table = table.clone(); + async move { + txn.execute_later( + format!("INSERT INTO {table} (id, n) VALUES ('a', 1)"), + vec![], + ); + txn.execute_later_mapping_error( + format!("INSERT INTO {table} (id, n) VALUES ('a', 2)"), + vec![], + |error| match error { + WriteError::Constraint { .. } => { + WriteError::Other(anyhow::anyhow!("mapped duplicate")) + } + other => other, + }, + ); + txn.query(&format!("SELECT id FROM {table}"), &[]).await?; + Ok(()) + } + .boxed() + }) + .await + .unwrap_err(); + assert!( + format!("{error}").contains("mapped duplicate"), + "unexpected error: {error}" + ); + assert_eq!(rows(&pool, &table).await, vec![]); + }); + teardown(&rt, &pool, &table); + } + + /// A queued write whose parameters do not encode stops everything queued + /// with it from being sent, the `COMMIT` included, so nothing is + /// committed. + #[test] + fn queued_write_that_does_not_encode_commits_nothing() { + let Some((rt, pool, table)) = setup() else { + return; + }; + rt.block_on(async { + let result: Result<(), WriteError> = run_write_interaction(&pool, "test", |txn, _| { + let table = table.clone(); + async move { + for (id, n) in [("x", "1"), ("y", "2"), ("z", "3")] { + let id: OwnedParam = if id == "y" { + // `id` is text: an integer does not encode. + Box::new(5i64) + } else { + Box::new(id.to_owned()) + }; + txn.execute_later( + format!("INSERT INTO {table} (id, n) VALUES ($1, {n})"), + vec![id], + ); + } + Ok(()) + } + .boxed() + }) + .await; + let error = result.unwrap_err(); + assert!( + format!("{error:#}").contains("error serializing parameter 0"), + "unexpected error: {error:#}" + ); + assert_eq!(rows(&pool, &table).await, vec![]); + }); + teardown(&rt, &pool, &table); + } + + /// A dry run still sends its queued writes, so their failures surface. + #[test] + fn dry_run_checks_queued_writes() { + let Some((rt, pool, table)) = setup() else { + return; + }; + rt.block_on(async { + let result = run_rolled_back_interaction(&pool, "test", |txn, _| { + let table = table.clone(); + async move { + txn.execute_later( + format!("INSERT INTO {table} (id, n) VALUES ('a', NULL)"), + vec![], + ); + Ok(()) + } + .boxed() + }) + .await; + assert!( + matches!(result, Err(WriteError::Constraint { .. })), + "unexpected result: {result:?}" + ); + }); + teardown(&rt, &pool, &table); + } + + /// When the closure fails, its queued writes are dropped unsent. + #[test] + fn queued_writes_are_dropped_on_failure() { + let Some((rt, pool, table)) = setup() else { + return; + }; + rt.block_on(async { + let result: Result<(), WriteError> = run_write_interaction(&pool, "test", |txn, _| { + let table = table.clone(); + async move { + txn.execute_later( + format!("INSERT INTO {table} (id, n) VALUES ('a', 1)"), + vec![], + ); + Err(WriteError::Other(anyhow::anyhow!("closure failed"))) + } + .boxed() + }) + .await; + assert!(result.is_err()); + assert_eq!(rows(&pool, &table).await, vec![]); + }); + teardown(&rt, &pool, &table); + } + /// The `BEGIN` goes out with the first statement, so that statement /// runs inside the transaction, whether it is prepared afresh or comes /// from the cache. @@ -949,7 +1428,11 @@ mod tests { Column::new("id", "text", &ids), Column::new("n", "bigint", &ns), ]; - assert!(batch_rows("t", &columns).is_err()); + let shapes: Vec<(&str, &str, usize)> = columns + .iter() + .map(|c| (c.name, c.sql_type, c.rows)) + .collect(); + assert!(check_batch("t", &shapes).is_err()); } #[test]