From c087ee127a3fc90dd2a9f235b6d74305f2fea513 Mon Sep 17 00:00:00 2001 From: liushuyu Date: Sun, 14 Nov 2021 01:46:20 -0700 Subject: [PATCH 1/4] feat(postgres): add an option to specify extra options ... ... this allow you to specify stuff like search path and statement timeouts etc. --- sqlx-core/src/postgres/connection/establish.rs | 4 ++++ sqlx-core/src/postgres/options/mod.rs | 17 +++++++++++++++++ sqlx-core/src/postgres/options/parse.rs | 9 +++++++++ 3 files changed, 30 insertions(+) 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..b6c0628c42 100644 --- a/sqlx-core/src/postgres/options/mod.rs +++ b/sqlx-core/src/postgres/options/mod.rs @@ -33,6 +33,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 command-line options 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 +86,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 +147,7 @@ impl PgConnectOptions { statement_cache_capacity: 100, application_name: var("PGAPPNAME").ok(), log_settings: Default::default(), + options: var("PGOPTIONS").ok(), } } @@ -319,6 +322,20 @@ impl PgConnectOptions { self } + /// Sets the extra options. Defaults to None + /// + /// # Example + /// + /// ```rust + /// # use sqlx_core::postgres::PgConnectOptions; + /// let options = PgConnectOptions::new() + /// .options("-c geqo=off -c statement_timeout=5min"); + /// ``` + pub fn options(mut self, options: &str) -> Self { + self.options = Some(options.to_owned()); + 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..f51b1de1fe 100644 --- a/sqlx-core/src/postgres/options/parse.rs +++ b/sqlx-core/src/postgres/options/parse.rs @@ -85,6 +85,8 @@ impl FromStr for PgConnectOptions { "application_name" => options = options.application_name(&*value), + "options" => options = options.options(&*value), + _ => log::warn!("ignoring unrecognized connect parameter: {}={}", key, value), } } @@ -195,3 +197,10 @@ 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_options() { + 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"), opts.options.as_deref()); +} From 40c2c349538743650932a8026c4d259fcc4afdb7 Mon Sep 17 00:00:00 2001 From: liushuyu Date: Tue, 23 Nov 2021 17:12:32 -0700 Subject: [PATCH 2/4] feat(postgres): set options through setting run-time parameters --- .../src/postgres/connection/establish.rs | 4 +- sqlx-core/src/postgres/mod.rs | 1 + sqlx-core/src/postgres/options/mod.rs | 40 +++++++++++++++--- sqlx-core/src/postgres/options/parse.rs | 41 ++++++++++++++++--- 4 files changed, 73 insertions(+), 13 deletions(-) diff --git a/sqlx-core/src/postgres/connection/establish.rs b/sqlx-core/src/postgres/connection/establish.rs index bd23780026..0fc5dc716e 100644 --- a/sqlx-core/src/postgres/connection/establish.rs +++ b/sqlx-core/src/postgres/connection/establish.rs @@ -40,8 +40,8 @@ impl PgConnection { params.push(("application_name", application_name)); } - if let Some(ref options) = options.options { - params.push(("options", options)); + for (k, v) in options.options.iter() { + params.push((k, v)); } stream diff --git a/sqlx-core/src/postgres/mod.rs b/sqlx-core/src/postgres/mod.rs index 85f8d2457b..6d3a8e71d4 100644 --- a/sqlx-core/src/postgres/mod.rs +++ b/sqlx-core/src/postgres/mod.rs @@ -32,6 +32,7 @@ pub use error::{PgDatabaseError, PgErrorPosition}; pub use listener::{PgListener, PgNotification}; pub use message::PgSeverity; pub use options::{PgConnectOptions, PgSslMode}; +pub(crate) use options::parse_options; pub use query_result::PgQueryResult; pub use row::PgRow; pub use statement::PgStatement; diff --git a/sqlx-core/src/postgres/options/mod.rs b/sqlx-core/src/postgres/options/mod.rs index b6c0628c42..3b7303489d 100644 --- a/sqlx-core/src/postgres/options/mod.rs +++ b/sqlx-core/src/postgres/options/mod.rs @@ -1,3 +1,4 @@ +use std::collections::HashMap; use std::env::var; use std::path::{Path, PathBuf}; @@ -33,7 +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 command-line options to send to the server at connection start. | +/// | `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. @@ -86,7 +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, + pub(crate) options: HashMap, } impl Default for PgConnectOptions { @@ -147,7 +148,9 @@ impl PgConnectOptions { statement_cache_capacity: 100, application_name: var("PGAPPNAME").ok(), log_settings: Default::default(), - options: var("PGOPTIONS").ok(), + options: var("PGOPTIONS") + .and_then(|v| Ok(parse_options(&v).unwrap_or_default().into_iter().collect())) + .unwrap_or_default(), } } @@ -329,10 +332,17 @@ impl PgConnectOptions { /// ```rust /// # use sqlx_core::postgres::PgConnectOptions; /// let options = PgConnectOptions::new() - /// .options("-c geqo=off -c statement_timeout=5min"); + /// .options(&[("geqo", "off"), ("statement_timeout", "5min")]); /// ``` - pub fn options(mut self, options: &str) -> Self { - self.options = Some(options.to_owned()); + pub fn options(mut self, options: I) -> Self + where + K: ToString, + V: ToString, + I: IntoIterator, + { + for (k, v) in options { + self.options.insert(k.to_string(), v.to_string()); + } self } @@ -353,6 +363,24 @@ impl PgConnectOptions { } } +/// Parse a libpq style options string +pub(crate) fn parse_options(input: &str) -> Option> { + let mut options = Vec::new(); + for part in input.split(' ') { + let part = part.trim(); + if part.is_empty() || part == "-c" { + continue; + } + let pair = part.splitn(2, '=').collect::>(); + if pair.len() != 2 { + return None; + } + options.push((pair[0].to_string(), pair[1].to_string())); + } + + Some(options) +} + fn default_host(port: u16) -> String { // try to check for the existence of a unix socket and uses that let socket = format!(".s.PGSQL.{}", port); diff --git a/sqlx-core/src/postgres/options/parse.rs b/sqlx-core/src/postgres/options/parse.rs index f51b1de1fe..e9948fc2ea 100644 --- a/sqlx-core/src/postgres/options/parse.rs +++ b/sqlx-core/src/postgres/options/parse.rs @@ -1,5 +1,5 @@ use crate::error::Error; -use crate::postgres::PgConnectOptions; +use crate::postgres::{parse_options, PgConnectOptions}; use percent_encoding::percent_decode_str; use std::net::IpAddr; use std::str::FromStr; @@ -85,7 +85,17 @@ impl FromStr for PgConnectOptions { "application_name" => options = options.application_name(&*value), - "options" => options = options.options(&*value), + "options" => { + if let Some(value) = parse_options(&value) { + options = options.options(value); + } + } + + 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), } @@ -198,9 +208,30 @@ fn it_parses_socket_correctly_with_username_percent_encoded() { assert_eq!(Some("database"), opts.database.as_deref()); } #[test] -fn it_parses_options() { - let uri = "postgres:///?options=-c%20synchronous_commit%3Doff%20--search_path%3Dpostgres"; +fn it_parses_libpq_options_correctly() { + let uri = "postgres:///?options=-c%20synchronous_commit%3Doff%20-c%20search_path%3Dpostgres"; + let opts = PgConnectOptions::from_str(uri).unwrap(); + + assert_eq!( + Some(&"off".to_string()), + opts.options.get("synchronous_commit") + ); + assert_eq!( + Some(&"postgres".to_string()), + opts.options.get("search_path") + ); +} +#[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 --search_path=postgres"), opts.options.as_deref()); + assert_eq!( + Some(&"off".to_string()), + opts.options.get("synchronous_commit") + ); + assert_eq!( + Some(&"postgres".to_string()), + opts.options.get("search_path") + ); } From 28462fda27dc294a3b4002b73410f23353c3a7e0 Mon Sep 17 00:00:00 2001 From: liushuyu Date: Wed, 24 Nov 2021 22:11:28 -0700 Subject: [PATCH 3/4] feat(postgres): use flat command-list instead of hashmap --- .../src/postgres/connection/establish.rs | 4 +-- sqlx-core/src/postgres/mod.rs | 1 - sqlx-core/src/postgres/options/mod.rs | 36 ++++++------------- sqlx-core/src/postgres/options/parse.rs | 27 ++++++-------- 4 files changed, 24 insertions(+), 44 deletions(-) diff --git a/sqlx-core/src/postgres/connection/establish.rs b/sqlx-core/src/postgres/connection/establish.rs index 0fc5dc716e..bd23780026 100644 --- a/sqlx-core/src/postgres/connection/establish.rs +++ b/sqlx-core/src/postgres/connection/establish.rs @@ -40,8 +40,8 @@ impl PgConnection { params.push(("application_name", application_name)); } - for (k, v) in options.options.iter() { - params.push((k, v)); + if let Some(ref options) = options.options { + params.push(("options", options)); } stream diff --git a/sqlx-core/src/postgres/mod.rs b/sqlx-core/src/postgres/mod.rs index 6d3a8e71d4..85f8d2457b 100644 --- a/sqlx-core/src/postgres/mod.rs +++ b/sqlx-core/src/postgres/mod.rs @@ -32,7 +32,6 @@ pub use error::{PgDatabaseError, PgErrorPosition}; pub use listener::{PgListener, PgNotification}; pub use message::PgSeverity; pub use options::{PgConnectOptions, PgSslMode}; -pub(crate) use options::parse_options; pub use query_result::PgQueryResult; pub use row::PgRow; pub use statement::PgStatement; diff --git a/sqlx-core/src/postgres/options/mod.rs b/sqlx-core/src/postgres/options/mod.rs index 3b7303489d..8f22786319 100644 --- a/sqlx-core/src/postgres/options/mod.rs +++ b/sqlx-core/src/postgres/options/mod.rs @@ -1,4 +1,3 @@ -use std::collections::HashMap; use std::env::var; use std::path::{Path, PathBuf}; @@ -87,7 +86,7 @@ pub struct PgConnectOptions { pub(crate) statement_cache_capacity: usize, pub(crate) application_name: Option, pub(crate) log_settings: LogSettings, - pub(crate) options: HashMap, + pub(crate) options: Option, } impl Default for PgConnectOptions { @@ -148,9 +147,7 @@ impl PgConnectOptions { statement_cache_capacity: 100, application_name: var("PGAPPNAME").ok(), log_settings: Default::default(), - options: var("PGOPTIONS") - .and_then(|v| Ok(parse_options(&v).unwrap_or_default().into_iter().collect())) - .unwrap_or_default(), + options: var("PGOPTIONS").ok(), } } @@ -332,7 +329,7 @@ impl PgConnectOptions { /// ```rust /// # use sqlx_core::postgres::PgConnectOptions; /// let options = PgConnectOptions::new() - /// .options(&[("geqo", "off"), ("statement_timeout", "5min")]); + /// .options([("geqo", "off"), ("statement_timeout", "5min")]); /// ``` pub fn options(mut self, options: I) -> Self where @@ -340,8 +337,15 @@ impl PgConnectOptions { V: ToString, I: IntoIterator, { + let mut options_str = String::new(); for (k, v) in options { - self.options.insert(k.to_string(), v.to_string()); + options_str += &format!("-c {}={}", k.to_string(), v.to_string()); + } + if let Some(ref mut v) = self.options { + v.push(' '); + v.push_str(&options_str); + } else { + self.options = Some(options_str); } self } @@ -363,24 +367,6 @@ impl PgConnectOptions { } } -/// Parse a libpq style options string -pub(crate) fn parse_options(input: &str) -> Option> { - let mut options = Vec::new(); - for part in input.split(' ') { - let part = part.trim(); - if part.is_empty() || part == "-c" { - continue; - } - let pair = part.splitn(2, '=').collect::>(); - if pair.len() != 2 { - return None; - } - options.push((pair[0].to_string(), pair[1].to_string())); - } - - Some(options) -} - fn default_host(port: u16) -> String { // try to check for the existence of a unix socket and uses that let socket = format!(".s.PGSQL.{}", port); diff --git a/sqlx-core/src/postgres/options/parse.rs b/sqlx-core/src/postgres/options/parse.rs index e9948fc2ea..d65d303cf0 100644 --- a/sqlx-core/src/postgres/options/parse.rs +++ b/sqlx-core/src/postgres/options/parse.rs @@ -1,5 +1,5 @@ use crate::error::Error; -use crate::postgres::{parse_options, PgConnectOptions}; +use crate::postgres::PgConnectOptions; use percent_encoding::percent_decode_str; use std::net::IpAddr; use std::str::FromStr; @@ -86,8 +86,11 @@ impl FromStr for PgConnectOptions { "application_name" => options = options.application_name(&*value), "options" => { - if let Some(value) = parse_options(&value) { - options = options.options(value); + if let Some(options) = options.options.as_mut() { + options.push(' '); + options.push_str(&*value); + } else { + options.options = Some(value.to_string()); } } @@ -209,16 +212,12 @@ fn it_parses_socket_correctly_with_username_percent_encoded() { } #[test] fn it_parses_libpq_options_correctly() { - let uri = "postgres:///?options=-c%20synchronous_commit%3Doff%20-c%20search_path%3Dpostgres"; + let uri = "postgres:///?options=-c%20synchronous_commit%3Doff%20--search_path%3Dpostgres"; let opts = PgConnectOptions::from_str(uri).unwrap(); assert_eq!( - Some(&"off".to_string()), - opts.options.get("synchronous_commit") - ); - assert_eq!( - Some(&"postgres".to_string()), - opts.options.get("search_path") + Some("-c synchronous_commit=off --search_path=postgres".into()), + opts.options ); } #[test] @@ -227,11 +226,7 @@ fn it_parses_sqlx_options_correctly() { let opts = PgConnectOptions::from_str(uri).unwrap(); assert_eq!( - Some(&"off".to_string()), - opts.options.get("synchronous_commit") - ); - assert_eq!( - Some(&"postgres".to_string()), - opts.options.get("search_path") + Some("-c synchronous_commit=off -c search_path=postgres".into()), + opts.options ); } From 47acea75e54473f0c2f5177ed2443bd27156cd01 Mon Sep 17 00:00:00 2001 From: liushuyu Date: Tue, 28 Dec 2021 21:07:09 -0700 Subject: [PATCH 4/4] feat(postgres): make the options taking parameters with `Display` trait --- sqlx-core/src/postgres/options/mod.rs | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/sqlx-core/src/postgres/options/mod.rs b/sqlx-core/src/postgres/options/mod.rs index 8f22786319..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; @@ -322,7 +323,7 @@ impl PgConnectOptions { self } - /// Sets the extra options. Defaults to None + /// Set additional startup options for the connection as a list of key-value pairs. /// /// # Example /// @@ -333,13 +334,13 @@ impl PgConnectOptions { /// ``` pub fn options(mut self, options: I) -> Self where - K: ToString, - V: ToString, + K: Display, + V: Display, I: IntoIterator, { let mut options_str = String::new(); for (k, v) in options { - options_str += &format!("-c {}={}", k.to_string(), v.to_string()); + options_str += &format!("-c {}={}", k, v); } if let Some(ref mut v) = self.options { v.push(' ');