diff --git a/sqlx-core/src/postgres/connection/establish.rs b/sqlx-core/src/postgres/connection/establish.rs index 59e3727c24..bd23780026 100644 --- a/sqlx-core/src/postgres/connection/establish.rs +++ b/sqlx-core/src/postgres/connection/establish.rs @@ -40,6 +40,10 @@ impl PgConnection { params.push(("application_name", application_name)); } + if let Some(ref options) = options.options { + params.push(("options", options)); + } + stream .send(Startup { username: Some(&options.username), diff --git a/sqlx-core/src/postgres/options/mod.rs b/sqlx-core/src/postgres/options/mod.rs index 1770959fdf..35e4897943 100644 --- a/sqlx-core/src/postgres/options/mod.rs +++ b/sqlx-core/src/postgres/options/mod.rs @@ -1,4 +1,5 @@ use std::env::var; +use std::fmt::Display; use std::path::{Path, PathBuf}; mod connect; @@ -33,6 +34,7 @@ pub use ssl_mode::PgSslMode; /// | `password` | `None` | Password to be used if the server demands password authentication. | /// | `port` | `5432` | Port number to connect to at the server host, or socket file name extension for Unix-domain connections. | /// | `dbname` | `None` | The database name. | +/// | `options` | `None` | The runtime parameters to send to the server at connection start. | /// /// The URI scheme designator can be either `postgresql://` or `postgres://`. /// Each of the URI parts is optional. @@ -85,6 +87,7 @@ pub struct PgConnectOptions { pub(crate) statement_cache_capacity: usize, pub(crate) application_name: Option, pub(crate) log_settings: LogSettings, + pub(crate) options: Option, } impl Default for PgConnectOptions { @@ -145,6 +148,7 @@ impl PgConnectOptions { statement_cache_capacity: 100, application_name: var("PGAPPNAME").ok(), log_settings: Default::default(), + options: var("PGOPTIONS").ok(), } } @@ -319,6 +323,34 @@ impl PgConnectOptions { self } + /// Set additional startup options for the connection as a list of key-value pairs. + /// + /// # Example + /// + /// ```rust + /// # use sqlx_core::postgres::PgConnectOptions; + /// let options = PgConnectOptions::new() + /// .options([("geqo", "off"), ("statement_timeout", "5min")]); + /// ``` + pub fn options(mut self, options: I) -> Self + where + K: Display, + V: Display, + I: IntoIterator, + { + let mut options_str = String::new(); + for (k, v) in options { + options_str += &format!("-c {}={}", k, v); + } + if let Some(ref mut v) = self.options { + v.push(' '); + v.push_str(&options_str); + } else { + self.options = Some(options_str); + } + self + } + /// We try using a socket if hostname starts with `/` or if socket parameter /// is specified. pub(crate) fn fetch_socket(&self) -> Option { diff --git a/sqlx-core/src/postgres/options/parse.rs b/sqlx-core/src/postgres/options/parse.rs index 5c5cd71ee8..d65d303cf0 100644 --- a/sqlx-core/src/postgres/options/parse.rs +++ b/sqlx-core/src/postgres/options/parse.rs @@ -85,6 +85,21 @@ impl FromStr for PgConnectOptions { "application_name" => options = options.application_name(&*value), + "options" => { + if let Some(options) = options.options.as_mut() { + options.push(' '); + options.push_str(&*value); + } else { + options.options = Some(value.to_string()); + } + } + + k if k.starts_with("options[") => { + if let Some(key) = k.strip_prefix("options[").unwrap().strip_suffix(']') { + options = options.options([(key, &*value)]); + } + } + _ => log::warn!("ignoring unrecognized connect parameter: {}={}", key, value), } } @@ -195,3 +210,23 @@ fn it_parses_socket_correctly_with_username_percent_encoded() { assert_eq!(Some("/var/lib/postgres/".into()), opts.socket); assert_eq!(Some("database"), opts.database.as_deref()); } +#[test] +fn it_parses_libpq_options_correctly() { + let uri = "postgres:///?options=-c%20synchronous_commit%3Doff%20--search_path%3Dpostgres"; + let opts = PgConnectOptions::from_str(uri).unwrap(); + + assert_eq!( + Some("-c synchronous_commit=off --search_path=postgres".into()), + opts.options + ); +} +#[test] +fn it_parses_sqlx_options_correctly() { + let uri = "postgres:///?options[synchronous_commit]=off&options[search_path]=postgres"; + let opts = PgConnectOptions::from_str(uri).unwrap(); + + assert_eq!( + Some("-c synchronous_commit=off -c search_path=postgres".into()), + opts.options + ); +}