Postgres: set generic execution plan only in query macro

launchbadge · Oct 10, 2024 · f252003 · f252003
1 parent 12c2986
commit f252003
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 28 deletions.
diff --git a/sqlx-macros-core/src/database/mod.rs b/sqlx-macros-core/src/database/mod.rs
@@ -54,7 +54,26 @@ impl<DB: DatabaseExt> CachingDescribeBlocking<DB> {
             let conn = match cache.entry(database_url.to_string()) {
                 hash_map::Entry::Occupied(hit) => hit.into_mut(),
                 hash_map::Entry::Vacant(miss) => {
-                    miss.insert(DB::Connection::connect(database_url).await?)
+                    let conn = miss.insert(DB::Connection::connect(database_url).await?);
+
+                    if DB::NAME == sqlx_postgres::Postgres::NAME {
+                        conn.execute(
+                            "
+                            DO $$
+                            BEGIN
+                                IF EXISTS (
+                                    SELECT 1
+                                    FROM pg_settings
+                                    WHERE name = 'plan_cache_mode'
+                                ) THEN
+                                    SET SESSION plan_cache_mode = 'force_generic_plan';
+                                END IF;
+                            END $$;
+                        ",
+                        )
+                        .await?;
+                    }
+                    conn
                 }
             };
 

diff --git a/sqlx-postgres/src/connection/describe.rs b/sqlx-postgres/src/connection/describe.rs
@@ -11,12 +11,8 @@ use crate::types::Oid;
 use crate::HashMap;
 use crate::{PgColumn, PgConnection, PgTypeInfo};
 use futures_core::future::BoxFuture;
-use futures_util::TryFutureExt;
 use smallvec::SmallVec;
-use sqlx_core::executor::Executor;
-use sqlx_core::from_row::FromRow;
 use sqlx_core::query_builder::QueryBuilder;
-use sqlx_core::raw_sql::raw_sql;
 use std::sync::Arc;
 
 /// Describes the type of the `pg_type.typtype` column
@@ -523,22 +519,7 @@ WHERE rngtypid = $1
             .display()
             .ok_or_else(|| err_protocol!("cannot EXPLAIN unnamed statement: {stmt_id:?}"))?;
 
-        let mut explain = format!(
-            "
-            BEGIN;
-            DO $$
-            BEGIN
-              IF EXISTS (
-                SELECT 1
-                FROM pg_settings
-                WHERE name = 'plan_cache_mode'
-              ) THEN
-                SET LOCAL plan_cache_mode = 'force_generic_plan';
-              END IF;
-            END $$;
-            EXPLAIN (VERBOSE, FORMAT JSON) EXECUTE {stmt_id_display}
-        "
-        );
+        let mut explain = format!("EXPLAIN (VERBOSE, FORMAT JSON) EXECUTE {stmt_id_display}");
         let mut comma = false;
 
         if params_len > 0 {
@@ -554,15 +535,11 @@ WHERE rngtypid = $1
                 comma = true;
             }
 
-            explain += ");
-            ROLLBACK;
-            ";
+            explain += ")";
         }
 
-        let (Json(explains),): (Json<SmallVec<[Explain; 1]>>,) = self
-            .fetch_one(raw_sql(&explain))
-            .map_ok(|row| FromRow::from_row(&row))
-            .await??;
+        let (Json(explains),): (Json<SmallVec<[Explain; 1]>>,) =
+            query_as(&explain).fetch_one(self).await?;
 
         let mut nullables = Vec::new();