Skip to content

Commit

Permalink
feat(postgres): add an option to specify extra options (#1539)
Browse files Browse the repository at this point in the history
* feat(postgres): add an option to specify extra options ...

... this allow you to specify stuff like search path and statement
timeouts etc.

* feat(postgres): set options through setting run-time parameters

* feat(postgres): use flat command-list instead of hashmap

* feat(postgres): make the options taking parameters with `Display` trait
  • Loading branch information
liushuyu authored Dec 29, 2021
1 parent 32f1273 commit b3091b0
Show file tree
Hide file tree
Showing 3 changed files with 71 additions and 0 deletions.
4 changes: 4 additions & 0 deletions sqlx-core/src/postgres/connection/establish.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down
32 changes: 32 additions & 0 deletions sqlx-core/src/postgres/options/mod.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use std::env::var;
use std::fmt::Display;
use std::path::{Path, PathBuf};

mod connect;
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -85,6 +87,7 @@ pub struct PgConnectOptions {
pub(crate) statement_cache_capacity: usize,
pub(crate) application_name: Option<String>,
pub(crate) log_settings: LogSettings,
pub(crate) options: Option<String>,
}

impl Default for PgConnectOptions {
Expand Down Expand Up @@ -145,6 +148,7 @@ impl PgConnectOptions {
statement_cache_capacity: 100,
application_name: var("PGAPPNAME").ok(),
log_settings: Default::default(),
options: var("PGOPTIONS").ok(),
}
}

Expand Down Expand Up @@ -332,6 +336,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<K, V, I>(mut self, options: I) -> Self
where
K: Display,
V: Display,
I: IntoIterator<Item = (K, V)>,
{
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<String> {
Expand Down
35 changes: 35 additions & 0 deletions sqlx-core/src/postgres/options/parse.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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),
}
}
Expand Down Expand Up @@ -197,3 +212,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
);
}

0 comments on commit b3091b0

Please sign in to comment.