From 25f5d4a28c2c6e0b7d54770bbb5c1b70a32d8d98 Mon Sep 17 00:00:00 2001 From: Alex Chi Z Date: Fri, 1 Mar 2024 14:37:01 -0500 Subject: [PATCH] apply suggestion + add new errtyp Signed-off-by: Alex Chi Z --- compute_tools/src/compute.rs | 54 +++++++++++++++++------------------- 1 file changed, 25 insertions(+), 29 deletions(-) diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index 768aaedf8ed0..6bf079bbed2e 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -777,36 +777,32 @@ impl ComputeNode { // but we can create a new one and grant it all privileges. let connstr = self.connstr.clone(); let mut client = match Client::connect(connstr.as_str(), NoTls) { - Err(e) => { - if let Some(code) = e.code() { - if code == &SqlState::INVALID_PASSWORD { - // connect with zenith_admin if cloud_admin could not authenticate - info!( - "cannot connect to postgres: {}, retrying with `zenith_admin` username", - e - ); - let mut zenith_admin_connstr = connstr.clone(); - - zenith_admin_connstr - .set_username("zenith_admin") - .map_err(|_| anyhow::anyhow!("invalid connstr"))?; - - let mut client = Client::connect(zenith_admin_connstr.as_str(), NoTls)?; - // Disable forwarding so that users don't get a cloud_admin role - client.simple_query("SET neon.forward_ddl = false")?; - client.simple_query("CREATE USER cloud_admin WITH SUPERUSER")?; - client.simple_query("GRANT zenith_admin TO cloud_admin")?; - drop(client); - - // reconnect with connstring with expected name - Client::connect(connstr.as_str(), NoTls)? - } else { - return Err(e.into()); - } - } else { - return Err(e.into()); + Err(e) => match e.code() { + Some(&SqlState::INVALID_PASSWORD) + | Some(&SqlState::INVALID_AUTHORIZATION_SPECIFICATION) => { + // connect with zenith_admin if cloud_admin could not authenticate + info!( + "cannot connect to postgres: {}, retrying with `zenith_admin` username", + e + ); + let mut zenith_admin_connstr = connstr.clone(); + + zenith_admin_connstr + .set_username("zenith_admin") + .map_err(|_| anyhow::anyhow!("invalid connstr"))?; + + let mut client = Client::connect(zenith_admin_connstr.as_str(), NoTls)?; + // Disable forwarding so that users don't get a cloud_admin role + client.simple_query("SET neon.forward_ddl = false")?; + client.simple_query("CREATE USER cloud_admin WITH SUPERUSER")?; + client.simple_query("GRANT zenith_admin TO cloud_admin")?; + drop(client); + + // reconnect with connstring with expected name + Client::connect(connstr.as_str(), NoTls)? } - } + _ => return Err(e.into()), + }, Ok(client) => client, };