compute_ctl: only try zenith_admin if could not auth

Signed-off-by: Alex Chi Z <chi@neon.tech>
neondatabase · Feb 28, 2024 · 37b631e · 37b631e
1 parent 54586d6
commit 37b631e
Showing 1 changed file with 29 additions and 19 deletions.
diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs
@@ -17,6 +17,7 @@ use chrono::{DateTime, Utc};
 use futures::future::join_all;
 use futures::stream::FuturesUnordered;
 use futures::StreamExt;
+use postgres::error::SqlState;
 use postgres::{Client, NoTls};
 use tokio;
 use tokio_postgres;
@@ -777,25 +778,34 @@ impl ComputeNode {
         let connstr = self.connstr.clone();
         let mut client = match Client::connect(connstr.as_str(), NoTls) {
             Err(e) => {
-                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)?
+                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());
+                }
             }
             Ok(client) => client,
         };