diff --git a/sqlx-macros-core/src/database/impls.rs b/sqlx-macros-core/src/database/impls.rs index 523b85cc14..4d0f9c5fa4 100644 --- a/sqlx-macros-core/src/database/impls.rs +++ b/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 ::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")] diff --git a/sqlx-macros-core/src/database/mod.rs b/sqlx-macros-core/src/database/mod.rs index 0885b3cca8..520b4bae87 100644 --- a/sqlx-macros-core/src/database/mod.rs +++ b/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>; + + /// 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 CachingDescribeBlocking { 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 } }; diff --git a/sqlx-postgres/src/connection/describe.rs b/sqlx-postgres/src/connection/describe.rs index f08e213783..ee9918909d 100644 --- a/sqlx-postgres/src/connection/describe.rs +++ b/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 . + /// + /// 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 . + #[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, diff --git a/tests/postgres/describe.rs b/tests/postgres/describe.rs index 806d8082c8..6ff22b30f5 100644 --- a/tests/postgres/describe.rs +++ b/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::().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::().await?;