Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(postgres): add an option to specify extra options #1539

Merged
merged 4 commits into from
Dec 29, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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 @@ -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<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);
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This appears to have a bug where if multiple options are provided the -c is tacked onto the previous one. This can be worked around by adding the space manually.

		.options([
			("default_transaction_isolation", "serializable "),
			("idle_in_transaction_session_timeout", "5min "),
			("statement_timeout", "2s"),
		]);

Or by calling this multiple times.

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bug filed: #1730

}
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 @@ -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
);
}