Skip to content

Commit

Permalink
Use URL crate, add comments
Browse files Browse the repository at this point in the history
  • Loading branch information
smklein committed Jun 18, 2024
1 parent dce0276 commit e4d6c78
Show file tree
Hide file tree
Showing 4 changed files with 34 additions and 19 deletions.
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 2 additions & 1 deletion nexus/db-queries/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
8 changes: 5 additions & 3 deletions nexus/db-queries/src/db/pool.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 },
Expand Down Expand Up @@ -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,
Expand Down
41 changes: 26 additions & 15 deletions nexus/db-queries/src/db/pool_connection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<PgConnection>;

Expand All @@ -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 {
Expand All @@ -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<String, anyhow::Error> {
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())
}
}

Expand All @@ -72,7 +83,7 @@ impl backend::Connector for DieselPgConnector {
&self,
backend: &Backend,
) -> Result<Self::Connection, Error> {
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)
Expand Down

0 comments on commit e4d6c78

Please sign in to comment.