postgres: only force a generic query plan where EXPLAIN is available (#4286)

The query macros set `plan_cache_mode = 'force_generic_plan'` on their
describe connection (#3541) to keep nullability inference from being skewed
by the NULL placeholder arguments. It was issued inside a `DO` block, which
CockroachDB rejects at parse time -- `SET` inside a function body, SQLSTATE
0A000 -- breaking the query macros against CockroachDB since 0.9.0 (#4274).

The setting only benefits the EXPLAIN-based inference, which already runs
only when `is_explain_available()` is true (false for CockroachDB,
Materialize and QuestDB). Gate the setting on that same check so it is never
issued on a database that cannot use it, replacing the `pg_settings` probe
with a `server_version_num() >= 12` guard.

Author: BastienClement

sqlx-macros-core/src/database/impls.rs

@@ -3,11 +3,13 @@ macro_rules! impl_database_ext {
         $database:path,
         row: $row:path,
         $(describe-blocking: $describe:path,)?
+        $(prepare-connection: $prepare:path,)?
     ) => {
         impl $crate::database::DatabaseExt for $database {
             const DATABASE_PATH: &'static str = stringify!($database);
             const ROW_PATH: &'static str = stringify!($row);
             impl_describe_blocking!($database, $($describe)?);
+            impl_prepare_describe_connection!($database, $($prepare)?);
         }
     }
 }
@@ -38,6 +40,19 @@ macro_rules! impl_describe_blocking {
     };
 }
 
+macro_rules! impl_prepare_describe_connection {
+    ($database:path $(,)?) => {
+        // No override: use the `DatabaseExt::prepare_describe_connection` default (no-op).
+    };
+    ($database:path, $prepare:path) => {
+        async fn prepare_describe_connection(
+            conn: &mut <Self as sqlx_core::database::Database>::Connection,
+        ) -> sqlx_core::Result<()> {
+            $prepare(conn).await
+        }
+    };
+}
+
 // The paths below will also be emitted from the macros, so they need to match the final facade.
 mod sqlx {
     #[cfg(feature = "mysql")]
@@ -61,6 +76,7 @@ impl_database_ext! {
 impl_database_ext! {
     sqlx::postgres::Postgres,
     row: sqlx::postgres::PgRow,
+    prepare-connection: sqlx::postgres::PgConnection::force_generic_plan_for_describe,
 }
 
 #[cfg(feature = "_sqlite")]

sqlx-macros-core/src/database/mod.rs

@@ -30,6 +30,18 @@ pub trait DatabaseExt: Database + TypeChecking {
         database_url: &str,
         driver_config: &config::drivers::Config,
     ) -> sqlx_core::Result<Describe<Self>>;
+
+    /// Prepare a freshly-opened connection used by the query macros for `describe`.
+    ///
+    /// Defaults to a no-op. Postgres overrides this to force a generic query plan,
+    /// which gives more accurate nullability inference for parameterized queries
+    /// (see launchbadge/sqlx#3541). The override is gated so it is skipped where it
+    /// doesn't apply -- e.g. CockroachDB, which rejected the previous implementation
+    /// (see launchbadge/sqlx#4274).
+    #[allow(async_fn_in_trait)]
+    async fn prepare_describe_connection(_conn: &mut Self::Connection) -> sqlx_core::Result<()> {
+        Ok(())
+    }
 }
 
 #[allow(dead_code)]
@@ -65,25 +77,7 @@ impl<DB: DatabaseExt> CachingDescribeBlocking<DB> {
                 hash_map::Entry::Occupied(hit) => hit.into_mut(),
                 hash_map::Entry::Vacant(miss) => {
                     let conn = miss.insert(DB::Connection::connect(database_url).await?);
-
-                    #[cfg(feature = "postgres")]
-                    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?;
-                    }
+                    DB::prepare_describe_connection(conn).await?;
                     conn
                 }
             };

sqlx-postgres/src/connection/describe.rs

@@ -1,4 +1,5 @@
 use crate::error::Error;
+use crate::executor::Executor;
 use crate::io::StatementId;
 use crate::query_as::query_as;
 use crate::statement::PgStatementMetadata;
@@ -18,6 +19,29 @@ impl PgConnection {
         !is_cockroachdb && !is_materialize && !is_questdb
     }
 
+    /// Prepare a freshly-opened connection that the query macros use for `describe`.
+    ///
+    /// Forces a generic query plan so that nullability inference via `EXPLAIN` reflects
+    /// the real query shape, rather than a plan specialized for the `NULL` placeholder
+    /// arguments bound while describing a statement.
+    /// See <https://github.com/launchbadge/sqlx/pull/3541>.
+    ///
+    /// Gated on `is_explain_available()`: the setting only matters for the `EXPLAIN`-based
+    /// inference, which is skipped on databases that don't support it -- CockroachDB,
+    /// Materialize and QuestDB -- and on the server version, as `plan_cache_mode` was only
+    /// introduced in PostgreSQL 12. Issuing it unconditionally previously broke the macros
+    /// against CockroachDB, which rejects `SET` inside the `DO` block this used to run
+    /// (`SQLSTATE 0A000`). See <https://github.com/launchbadge/sqlx/issues/4274>.
+    #[doc(hidden)]
+    pub async fn force_generic_plan_for_describe(&mut self) -> Result<(), Error> {
+        if self.is_explain_available() && self.server_version_num().is_some_and(|v| v >= 120_000) {
+            self.execute("SET plan_cache_mode = 'force_generic_plan'")
+                .await?;
+        }
+
+        Ok(())
+    }
+
     pub(crate) async fn get_nullable_for_columns(
         &mut self,
         stmt_id: StatementId,

tests/postgres/describe.rs

@@ -42,6 +42,28 @@ async fn it_describes_expression() -> anyhow::Result<()> {
     Ok(())
 }
 
+// Regression test for launchbadge/sqlx#3541 and #4274: the query macros force a
+// generic query plan on their describe connection so that nullability inference via
+// `EXPLAIN` isn't skewed by the `NULL` placeholder arguments bound during `describe`.
+// This checks the mechanism still applies on PostgreSQL.
+//
+// The complementary half of #4274 -- skipping this on databases without `EXPLAIN`
+// support (CockroachDB, Materialize, QuestDB), where the previous `DO` block failed
+// to even parse -- cannot be exercised here, as CI runs no such database.
+#[sqlx_macros::test]
+async fn it_forces_generic_plan_for_describe() -> anyhow::Result<()> {
+    let mut conn = new::<Postgres>().await?;
+
+    conn.force_generic_plan_for_describe().await?;
+
+    let mode: String = sqlx::query_scalar("SHOW plan_cache_mode")
+        .fetch_one(&mut conn)
+        .await?;
+    assert_eq!(mode, "force_generic_plan");
+
+    Ok(())
+}
+
 #[sqlx_macros::test]
 async fn it_describes_enum() -> anyhow::Result<()> {
     let mut conn = new::<Postgres>().await?;