Skip to content

Commit 8a9de13

Browse files
committed
node: Make sure PGAPPNAME has an effect
tokio-postgres does not respect PGAPPNAME, so that its connections are not named properly in pg_stat_activity etc. We manually add it to the URL to help tracking and debugging connection issues
1 parent e62f6de commit 8a9de13

File tree

1 file changed

+30
-2
lines changed

1 file changed

+30
-2
lines changed

node/src/config.rs

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -252,7 +252,7 @@ impl Shard {
252252
fn validate(&mut self, name: &str) -> Result<()> {
253253
ShardName::new(name.to_string()).map_err(|e| anyhow!(e))?;
254254

255-
self.connection = shellexpand::env(&self.connection)?.into_owned();
255+
self.expand_connection()?;
256256

257257
if matches!(self.pool_size, PoolSize::None) {
258258
return Err(anyhow!("missing pool size definition for shard `{}`", name));
@@ -301,6 +301,25 @@ impl Shard {
301301
replicas,
302302
})
303303
}
304+
305+
fn expand_connection(&mut self) -> Result<()> {
306+
let mut url = Url::parse(shellexpand::env(&self.connection)?.as_ref())?;
307+
// Put the PGAPPNAME into the URL since tokio-postgres ignores this
308+
// environment variable
309+
if let Some(app_name) = std::env::var("PGAPPNAME").ok() {
310+
let query = match url.query() {
311+
Some(query) => {
312+
format!("{query}&application_name={app_name}")
313+
}
314+
None => {
315+
format!("application_name={app_name}")
316+
}
317+
};
318+
url.set_query(Some(&query));
319+
}
320+
self.connection = url.to_string();
321+
Ok(())
322+
}
304323
}
305324

306325
#[derive(Clone, Debug, Deserialize, Serialize)]
@@ -1944,6 +1963,10 @@ mod tests {
19441963
let query = NodeId::new("query_node_1").unwrap();
19451964
let other = NodeId::new("other_node_1").unwrap();
19461965

1966+
let appname = std::env::var("PGAPPNAME").ok();
1967+
unsafe {
1968+
std::env::set_var("PGAPPNAME", "config-test");
1969+
}
19471970
let shard = {
19481971
let mut shard = toml::from_str::<Shard>(
19491972
r#"
@@ -1961,10 +1984,15 @@ fdw_pool_size = [
19611984
shard.validate("index_node_1").unwrap();
19621985
shard
19631986
};
1987+
if let Some(appname) = appname {
1988+
unsafe {
1989+
std::env::set_var("PGAPPNAME", appname);
1990+
}
1991+
}
19641992

19651993
assert_eq!(
19661994
shard.connection,
1967-
"postgresql://postgres:postgres@postgres/graph"
1995+
"postgresql://postgres:postgres@postgres/graph?application_name=config-test"
19681996
);
19691997

19701998
assert_eq!(shard.pool_size.size_for(&index, "ashard").unwrap(), 20);

0 commit comments

Comments
 (0)