diff --git a/sqlx-core/src/common/statement_cache.rs b/sqlx-core/src/common/statement_cache.rs index eb800ca623..1b3dbae8f7 100644 --- a/sqlx-core/src/common/statement_cache.rs +++ b/sqlx-core/src/common/statement_cache.rs @@ -21,6 +21,11 @@ impl StatementCache { self.inner.get_mut(k) } + /// Removes statement by key, returning removed statement if removed + pub fn remove(&mut self, k: &str) -> Option { + self.inner.remove(k) + } + /// Inserts a new statement to the cache, returning the least recently used /// statement id if the cache is full, or if inserting with an existing key, /// the replaced existing statement. diff --git a/sqlx-postgres/src/connection/executor.rs b/sqlx-postgres/src/connection/executor.rs index e0f4c3d44a..96a0ed7b92 100644 --- a/sqlx-postgres/src/connection/executor.rs +++ b/sqlx-postgres/src/connection/executor.rs @@ -8,8 +8,8 @@ use crate::message::{ }; use crate::statement::PgStatementMetadata; use crate::{ - statement::PgStatement, PgArguments, PgConnection, PgQueryResult, PgRow, PgTypeInfo, - PgValueFormat, Postgres, + statement::PgStatement, PgArguments, PgConnection, PgDatabaseError, PgQueryResult, PgRow, + PgTypeInfo, PgValueFormat, Postgres, }; use futures_core::future::BoxFuture; use futures_core::stream::BoxStream; @@ -196,22 +196,16 @@ impl PgConnection { Ok(statement) } - pub(crate) async fn run<'e, 'c: 'e, 'q: 'e>( + async fn try_get_or_prepare<'e, 'c: 'e, 'q: 'e>( &'c mut self, - query: SqlStr, - arguments: Option, + sql: &str, + arguments: Option<&mut PgArguments>, persistent: bool, metadata_opt: Option>, - ) -> Result, Error>> + 'e, Error> { - let mut logger = QueryLogger::new(query, self.inner.log_settings.clone()); - let sql = logger.sql().as_str(); - - // before we continue, wait until we are "ready" to accept more queries - self.wait_until_ready().await?; - - let mut metadata: Arc; + ) -> Result<(PgValueFormat, Arc), Error> { + let metadata: Arc; - let format = if let Some(mut arguments) = arguments { + let format = if let Some(arguments) = arguments { // Check this before we write anything to the stream. // // Note: Postgres actually interprets this value as unsigned, @@ -292,10 +286,62 @@ impl PgConnection { self.inner.stream.flush().await?; + Ok((format, metadata)) + } + + pub(crate) async fn run<'e, 'c: 'e, 'q: 'e>( + &'c mut self, + query: SqlStr, + mut arguments: Option, + persistent: bool, + metadata_opt: Option>, + ) -> Result, Error>> + 'e, Error> { + let mut logger = QueryLogger::new(query, self.inner.log_settings.clone()); + let sql = logger.sql().as_str(); + + // before we continue, wait until we are "ready" to accept more queries + self.wait_until_ready().await?; + + let (mut format, mut metadata) = self + .try_get_or_prepare(sql, arguments.as_mut(), persistent, metadata_opt.clone()) + .await?; + + let mut message = match self.inner.stream.recv().await { + Ok(msg) => msg, + Err(err) => { + if let Some(clear_backend_cache) = check_stale_plan(&err) { + // Save transaction mode. It will be lost after invalidating + let is_in_tx = self.in_transaction(); + + self.invalidate_cached_statement(sql, clear_backend_cache) + .await?; + + // If we were in transaction mode we can't retry statement, + // so we can immediately return err + if is_in_tx { + return Err(err); + } + + // Otherwise we can retry statement in hope everything is ok. + (format, metadata) = self + .try_get_or_prepare( + sql, + // It should be safe to retry `patch` on the same arguments + arguments.as_mut(), + persistent, + metadata_opt.clone(), + ) + .await?; + + self.inner.stream.recv().await? + } else { + return Err(err); + } + } + }; + Ok(try_stream! { loop { - let message = self.inner.stream.recv().await?; - match message.format { BackendMessageFormat::BindComplete | BackendMessageFormat::ParseComplete @@ -369,6 +415,8 @@ impl PgConnection { )); } } + + message = self.inner.stream.recv().await?; } Ok(()) @@ -485,3 +533,23 @@ impl<'c> Executor<'c> for &'c mut PgConnection { }) } } + +// Returns: +// - `None` - if not 'stale query plan' +// - `Some(false)` - if it is stale plan, but we don't need to deallocate objects on backend. +// It can happen because of `DISCARD ALL`, `DEALLOCATE` or due to pgbouncer in +// transaction pooling mode +// - `Some(true)` - if we should invalidate both backend and frontend caches +fn check_stale_plan(error: &Error) -> Option { + let error = error + .as_database_error()? + .try_downcast_ref::()?; + + match (error.code(), error.routine()) { + // "cached plan must not change result type" + ("0A000", Some("RevalidateCachedQuery")) => Some(true), + // DISCARD ALL / DEALLOCATE / pgbouncer + ("26000", _) => Some(false), + _ => None, + } +} diff --git a/sqlx-postgres/src/connection/mod.rs b/sqlx-postgres/src/connection/mod.rs index d594585b6c..76637a8a93 100644 --- a/sqlx-postgres/src/connection/mod.rs +++ b/sqlx-postgres/src/connection/mod.rs @@ -146,6 +146,31 @@ impl PgConnection { TransactionStatus::Error | TransactionStatus::Idle => false, } } + + pub(crate) async fn invalidate_cached_statement( + &mut self, + sql: &str, + backend: bool, + ) -> Result<(), Error> { + self.wait_until_ready().await?; + + let Some((statement_id, _)) = self.inner.cache_statement.remove(sql) else { + return Ok(()); + }; + + if backend { + self.inner + .stream + .write_msg(Close::Statement(statement_id))?; + self.write_sync(); + self.inner.stream.flush().await?; + + self.wait_for_close_complete(1).await?; + self.recv_ready_for_query().await?; + } + + Ok(()) + } } impl Debug for PgConnection {