Invalided cached prepared statements on RevalidateCachedQuery and 26000

Author: Flassie

sqlx-core/src/common/statement_cache.rs

@@ -21,6 +21,11 @@ impl<T> StatementCache<T> {
         self.inner.get_mut(k)
     }
 
+    /// Removes statement by key, returning removed statement if removed
+    pub fn remove(&mut self, k: &str) -> Option<T> {
+        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.

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<PgArguments>,
+        sql: &str,
+        arguments: Option<&mut PgArguments>,
         persistent: bool,
-        metadata_opt: Option<Arc<PgStatementMetadata>>,
-    ) -> Result<impl Stream<Item = Result<Either<PgQueryResult, PgRow>, 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<PgStatementMetadata>;
+        metadata_opt: Option<Arc<PgStatementMetadata>>
+    ) -> Result<(PgValueFormat, Arc<PgStatementMetadata>), Error> {
+        let metadata: Arc<PgStatementMetadata>;
 
-        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<PgArguments>,
+        persistent: bool,
+        metadata_opt: Option<Arc<PgStatementMetadata>>,
+    ) -> Result<impl Stream<Item = Result<Either<PgQueryResult, PgRow>, 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<bool> {
+    let Some(db_err) = error.as_database_error() else {
+        return None;
+    };
+    let Some(pg) = db_err.try_downcast_ref::<PgDatabaseError>() 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,
+    }
+}
+

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 {