apply suggestion + add new errtyp

Signed-off-by: Alex Chi Z <chi@neon.tech>
neondatabase · Mar 1, 2024 · 25f5d4a · 25f5d4a
1 parent 37b631e
commit 25f5d4a
Showing 1 changed file with 25 additions and 29 deletions.
diff --git 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,
         };