diff --git a/Cargo.lock b/Cargo.lock index b76b6e6436..5efa665903 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4807,6 +4807,7 @@ dependencies = [ "term", "thiserror", "tokio", + "url", "usdt", "uuid", ] diff --git a/nexus/db-queries/Cargo.toml b/nexus/db-queries/Cargo.toml index 0b62dbdfd9..865869a5ae 100644 --- a/nexus/db-queries/Cargo.toml +++ b/nexus/db-queries/Cargo.toml @@ -46,8 +46,9 @@ strum.workspace = true swrite.workspace = true thiserror.workspace = true tokio = { workspace = true, features = ["full"] } -uuid.workspace = true +url.workspace = true usdt.workspace = true +uuid.workspace = true db-macros.workspace = true nexus-auth.workspace = true diff --git a/nexus/db-queries/src/db/pool.rs b/nexus/db-queries/src/db/pool.rs index 3c350dbe79..dccee6fa3f 100644 --- a/nexus/db-queries/src/db/pool.rs +++ b/nexus/db-queries/src/db/pool.rs @@ -83,7 +83,7 @@ fn make_postgres_connector( // - Creating async_bb8_diesel connections that also wrap DTraceConnections. let user = "root"; let db = "omicron"; - let args = Some("sslmode=disable"); + let args = vec![("sslmode", "disable")]; Arc::new(DieselPgConnector::new( log, DieselPgConnectorArgs { user, db, args }, @@ -125,9 +125,11 @@ impl Pool { } /// Creates a new qorb-backed connection pool which returns an error - /// if claims are not quickly available. + /// if claims are not available within one millisecond. /// - /// This is intended for test-only usage. + /// This is intended for test-only usage, in particular for tests where + /// claim requests should rapidly return errors when a backend has been + /// intentionally disabled. #[cfg(any(test, feature = "testing"))] pub fn new_single_host_failfast( log: &Logger, diff --git a/nexus/db-queries/src/db/pool_connection.rs b/nexus/db-queries/src/db/pool_connection.rs index e195643008..9a33370a5a 100644 --- a/nexus/db-queries/src/db/pool_connection.rs +++ b/nexus/db-queries/src/db/pool_connection.rs @@ -13,6 +13,7 @@ use diesel::PgConnection; use diesel_dtrace::DTraceConnection; use qorb::backend::{self, Backend, Error}; use slog::Logger; +use url::Url; pub type DbConnection = DTraceConnection; @@ -22,14 +23,15 @@ pub const DISALLOW_FULL_TABLE_SCAN_SQL: &str = /// A [backend::Connector] which provides access to [PgConnection]. pub(crate) struct DieselPgConnector { log: Logger, - prefix: String, - suffix: String, + user: String, + db: String, + args: Vec<(String, String)>, } pub(crate) struct DieselPgConnectorArgs<'a> { pub(crate) user: &'a str, pub(crate) db: &'a str, - pub(crate) args: Option<&'a str>, + pub(crate) args: Vec<(&'a str, &'a str)>, } impl DieselPgConnector { @@ -47,20 +49,29 @@ impl DieselPgConnector { let DieselPgConnectorArgs { user, db, args } = args; Self { log: log.clone(), - prefix: format!("postgresql://{user}@"), - suffix: format!( - "/{db}{}", - args.map(|args| format!("?{args}")).unwrap_or("".to_string()) - ), + user: user.to_string(), + db: db.to_string(), + args: args + .into_iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect(), } } - fn to_url(&self, address: std::net::SocketAddr) -> String { - format!( - "{prefix}{address}{suffix}", - prefix = self.prefix, - suffix = self.suffix, - ) + fn to_url( + &self, + address: std::net::SocketAddr, + ) -> Result { + let user = &self.user; + let db = &self.db; + let mut url = + Url::parse(&format!("postgresql://{user}@{address}/{db}"))?; + + for (k, v) in &self.args { + url.query_pairs_mut().append_pair(k, v); + } + + Ok(url.as_str().to_string()) } } @@ -72,7 +83,7 @@ impl backend::Connector for DieselPgConnector { &self, backend: &Backend, ) -> Result { - let url = self.to_url(backend.address); + let url = self.to_url(backend.address).map_err(Error::Other)?; let conn = tokio::task::spawn_blocking(move || { let pg_conn = DbConnection::establish(&url)