From 4e99ce1eb941d55f2d9d96fad37b205cd9170a86 Mon Sep 17 00:00:00 2001 From: Andrei Sidorenko Date: Tue, 11 Aug 2026 16:45:24 +0300 Subject: [PATCH 1/2] Invalided cached prepared statements on `RevalidateCachedQuery` and 26000 --- sqlx-core/src/common/statement_cache.rs | 5 + sqlx-postgres/src/connection/executor.rs | 113 ++++++++++++++++++----- sqlx-postgres/src/connection/mod.rs | 19 ++++ 3 files changed, 114 insertions(+), 23 deletions(-) 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..1254e7dfc9 100644 --- a/sqlx-postgres/src/connection/executor.rs +++ b/sqlx-postgres/src/connection/executor.rs @@ -2,15 +2,9 @@ use crate::error::Error; use crate::executor::{Execute, Executor}; use crate::io::{PortalId, StatementId}; use crate::logger::QueryLogger; -use crate::message::{ - self, BackendMessageFormat, Bind, Close, CommandComplete, DataRow, ParameterDescription, Parse, - ParseComplete, RowDescription, -}; +use crate::message::{self, BackendMessageFormat, Bind, Close, CommandComplete, DataRow, ParameterDescription, Parse, ParseComplete, ReceivedMessage, RowDescription}; use crate::statement::PgStatementMetadata; -use crate::{ - statement::PgStatement, PgArguments, PgConnection, PgQueryResult, PgRow, PgTypeInfo, - PgValueFormat, Postgres, -}; +use crate::{statement::PgStatement, PgArguments, PgConnection, PgDatabaseError, PgQueryResult, PgRow, PgTypeInfo, PgValueFormat, Postgres}; use futures_core::future::BoxFuture; use futures_core::stream::BoxStream; use futures_core::Stream; @@ -19,6 +13,7 @@ use sqlx_core::arguments::Arguments; use sqlx_core::sql_str::SqlStr; use sqlx_core::Either; use std::{pin::pin, sync::Arc}; +use sqlx_core::connection::Connection; async fn prepare( conn: &mut PgConnection, @@ -196,22 +191,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; + metadata_opt: Option> + ) -> 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 +281,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 +410,8 @@ impl PgConnection { )); } } + + message = self.inner.stream.recv().await?; } Ok(()) @@ -485,3 +528,27 @@ 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 Some(db_err) = error.as_database_error() else { + return None; + }; + let Some(pg) = db_err.try_downcast_ref::() else { + return None; + }; + + match (pg.code(), pg.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..3bc92ff7f1 100644 --- a/sqlx-postgres/src/connection/mod.rs +++ b/sqlx-postgres/src/connection/mod.rs @@ -146,6 +146,25 @@ 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 { From d07f554aa182de8231ad77ad00dc675950a40d03 Mon Sep 17 00:00:00 2001 From: Andrei Sidorenko Date: Tue, 11 Aug 2026 18:10:38 +0300 Subject: [PATCH 2/2] Fmt + Clippy --- sqlx-postgres/src/connection/executor.rs | 57 ++++++++++++------------ sqlx-postgres/src/connection/mod.rs | 12 +++-- 2 files changed, 38 insertions(+), 31 deletions(-) diff --git a/sqlx-postgres/src/connection/executor.rs b/sqlx-postgres/src/connection/executor.rs index 1254e7dfc9..96a0ed7b92 100644 --- a/sqlx-postgres/src/connection/executor.rs +++ b/sqlx-postgres/src/connection/executor.rs @@ -2,9 +2,15 @@ use crate::error::Error; use crate::executor::{Execute, Executor}; use crate::io::{PortalId, StatementId}; use crate::logger::QueryLogger; -use crate::message::{self, BackendMessageFormat, Bind, Close, CommandComplete, DataRow, ParameterDescription, Parse, ParseComplete, ReceivedMessage, RowDescription}; +use crate::message::{ + self, BackendMessageFormat, Bind, Close, CommandComplete, DataRow, ParameterDescription, Parse, + ParseComplete, RowDescription, +}; use crate::statement::PgStatementMetadata; -use crate::{statement::PgStatement, PgArguments, PgConnection, PgDatabaseError, PgQueryResult, PgRow, PgTypeInfo, PgValueFormat, Postgres}; +use crate::{ + statement::PgStatement, PgArguments, PgConnection, PgDatabaseError, PgQueryResult, PgRow, + PgTypeInfo, PgValueFormat, Postgres, +}; use futures_core::future::BoxFuture; use futures_core::stream::BoxStream; use futures_core::Stream; @@ -13,7 +19,6 @@ use sqlx_core::arguments::Arguments; use sqlx_core::sql_str::SqlStr; use sqlx_core::Either; use std::{pin::pin, sync::Arc}; -use sqlx_core::connection::Connection; async fn prepare( conn: &mut PgConnection, @@ -196,7 +201,7 @@ impl PgConnection { sql: &str, arguments: Option<&mut PgArguments>, persistent: bool, - metadata_opt: Option> + metadata_opt: Option>, ) -> Result<(PgValueFormat, Arc), Error> { let metadata: Arc; @@ -297,12 +302,9 @@ impl PgConnection { // 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 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, @@ -311,26 +313,29 @@ impl PgConnection { // 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?; + 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) + 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?; + (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) + return Err(err); } } }; @@ -536,14 +541,11 @@ impl<'c> Executor<'c> for &'c mut PgConnection { // transaction pooling mode // - `Some(true)` - if we should invalidate both backend and frontend caches fn check_stale_plan(error: &Error) -> Option { - let Some(db_err) = error.as_database_error() else { - return None; - }; - let Some(pg) = db_err.try_downcast_ref::() else { - return None; - }; + let error = error + .as_database_error()? + .try_downcast_ref::()?; - match (pg.code(), pg.routine()) { + match (error.code(), error.routine()) { // "cached plan must not change result type" ("0A000", Some("RevalidateCachedQuery")) => Some(true), // DISCARD ALL / DEALLOCATE / pgbouncer @@ -551,4 +553,3 @@ fn check_stale_plan(error: &Error) -> Option { _ => None, } } - diff --git a/sqlx-postgres/src/connection/mod.rs b/sqlx-postgres/src/connection/mod.rs index 3bc92ff7f1..76637a8a93 100644 --- a/sqlx-postgres/src/connection/mod.rs +++ b/sqlx-postgres/src/connection/mod.rs @@ -147,15 +147,21 @@ impl PgConnection { } } - pub(crate) async fn invalidate_cached_statement(&mut self, sql: &str, backend: bool) -> Result<(), Error> { + 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(()) + return Ok(()); }; if backend { - self.inner.stream.write_msg(Close::Statement(statement_id))?; + self.inner + .stream + .write_msg(Close::Statement(statement_id))?; self.write_sync(); self.inner.stream.flush().await?;