[PostgreSQL] Invalided stale prepared statements

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

@@ -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<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>;
+    ) -> 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 +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<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 +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<bool> {
+    let error = error
+        .as_database_error()?
+        .try_downcast_ref::<PgDatabaseError>()?;
+
+    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,
+    }
+}

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 {