diff --git a/sqlx-core/src/postgres/connection/establish.rs b/sqlx-core/src/postgres/connection/establish.rs index 59e3727c24..2c2515ebbc 100644 --- a/sqlx-core/src/postgres/connection/establish.rs +++ b/sqlx-core/src/postgres/connection/establish.rs @@ -143,6 +143,7 @@ impl PgConnection { cache_type_oid: HashMap::new(), cache_type_info: HashMap::new(), log_settings: options.log_settings.clone(), + override_row_limit: options.override_row_limit, }) } } diff --git a/sqlx-core/src/postgres/connection/executor.rs b/sqlx-core/src/postgres/connection/executor.rs index 3b5ab3ef25..c5e9bb12c5 100644 --- a/sqlx-core/src/postgres/connection/executor.rs +++ b/sqlx-core/src/postgres/connection/executor.rs @@ -232,7 +232,11 @@ impl PgConnection { // the protocol-level limit acts nearly identically to the `LIMIT` in SQL self.stream.write(message::Execute { portal: None, - limit: limit.into(), + limit: if self.override_row_limit { + 0 + } else { + limit.into() + }, }); // finally, [Sync] asks postgres to process the messages that we sent and respond with diff --git a/sqlx-core/src/postgres/connection/mod.rs b/sqlx-core/src/postgres/connection/mod.rs index 2688727ac4..1358725d01 100644 --- a/sqlx-core/src/postgres/connection/mod.rs +++ b/sqlx-core/src/postgres/connection/mod.rs @@ -62,6 +62,8 @@ pub struct PgConnection { pub(crate) transaction_depth: usize, log_settings: LogSettings, + + override_row_limit: bool, } impl PgConnection { diff --git a/sqlx-core/src/postgres/options/mod.rs b/sqlx-core/src/postgres/options/mod.rs index 490b59c75c..50e4d33506 100644 --- a/sqlx-core/src/postgres/options/mod.rs +++ b/sqlx-core/src/postgres/options/mod.rs @@ -86,6 +86,7 @@ pub struct PgConnectOptions { pub(crate) statement_cache_capacity: usize, pub(crate) application_name: Option, pub(crate) log_settings: LogSettings, + pub(crate) override_row_limit: bool, } impl Default for PgConnectOptions { @@ -138,6 +139,7 @@ impl PgConnectOptions { statement_cache_capacity: 100, application_name: var("PGAPPNAME").ok(), log_settings: Default::default(), + override_row_limit: false, } } @@ -296,6 +298,17 @@ impl PgConnectOptions { self } + /// Request all rows from the backend when calling `fetch_one` and + /// `fetch_optional`. This is a workaround for a issue in which + /// CockroachDB reports that sqlx is using multiple portals. + /// See https://github.com/launchbadge/sqlx/issues/933 anad + /// https://github.com/cockroachdb/cockroach/issues/40195?version=v20.1 for + /// more details. + pub fn override_row_limit(&mut self, override_row_limit: bool) -> &mut Self { + self.override_row_limit = override_row_limit; + self + } + /// We try using a socket if hostname starts with `/` or if socket parameter /// is specified. pub(crate) fn fetch_socket(&self) -> Option {