feat(toolbox): add zk supervisor database commands

zk_toolbox/Cargo.toml

@@ -32,7 +32,9 @@ clap = { version = "4.4", features = ["derive", "wrap_help"] }
 cliclack = "0.2.5"
 console = "0.15.8"
 ethers = "2.0"
+futures = "0.3.30"
 human-panic = "2.0"
+lazy_static = "1.4.0"
 once_cell = "1.19.0"
 rand = "0.8.5"
 serde = { version = "1.0", features = ["derive"] }
@@ -41,9 +43,8 @@ serde_yaml = "0.9"
 sqlx = { version = "0.7.4", features = ["runtime-tokio", "migrate", "postgres"] }
 strum = "0.26.2"
 strum_macros = "0.26.2"
+thiserror = "1.0.57"
 tokio = { version = "1.37", features = ["full"] }
 toml = "0.8.12"
 url = { version = "2.5.0", features = ["serde"] }
 xshell = "0.2.6"
-futures = "0.3.30"
-thiserror = "1.0.57"

zk_toolbox/crates/common/Cargo.toml

@@ -16,13 +16,14 @@ clap.workspace = true
 cliclack.workspace = true
 console.workspace = true
 ethers.workspace = true
+futures.workspace = true
 once_cell.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 serde_yaml.workspace = true
 sqlx.workspace = true
 strum_macros.workspace = true
+tokio.workspace = true
 toml.workspace = true
 url.workspace = true
 xshell.workspace = true
-futures.workspace = true

zk_toolbox/crates/common/src/cmd.rs

@@ -1,3 +1,5 @@
+use std::process::Output;
+
 use anyhow::bail;
 use console::style;
 
@@ -31,13 +33,6 @@ impl<'a> Cmd<'a> {
 
     /// Run the command without capturing its output.
     pub fn run(&mut self) -> anyhow::Result<()> {
-        self.run_cmd()?;
-        Ok(())
-    }
-
-    /// Run the command and capture its output, logging the command
-    /// and its output if verbose selected.
-    fn run_cmd(&mut self) -> anyhow::Result<()> {
         if global_config().verbose || self.force_run {
             logger::debug(format!("Running: {}", self.inner));
             logger::new_empty_line();
@@ -60,6 +55,25 @@ impl<'a> Cmd<'a> {
         Ok(())
     }
 
+    /// Run the command and return its output.
+    pub fn run_with_output(&mut self) -> anyhow::Result<Output> {
+        if global_config().verbose || self.force_run {
+            logger::debug(format!("Running: {}", self.inner));
+            logger::new_empty_line();
+        }
+
+        self.inner.set_ignore_status(true);
+        let output = self.inner.output()?;
+
+        if global_config().verbose || self.force_run {
+            logger::raw(log_output(&output));
+            logger::new_empty_line();
+            logger::new_line();
+        }
+
+        Ok(output)
+    }
+
     fn check_output_status(&self, output: &std::process::Output) -> anyhow::Result<()> {
         if !output.status.success() {
             logger::new_line();

zk_toolbox/crates/common/src/db.rs

@@ -1,5 +1,7 @@
 use std::{collections::HashMap, path::PathBuf};
 
+use anyhow::anyhow;
+use serde::{Deserialize, Serialize};
 use sqlx::{
     migrate::{Migrate, MigrateError, Migrator},
     Connection, PgConnection,
@@ -9,22 +11,63 @@ use xshell::Shell;
 
 use crate::{config::global_config, logger};
 
-pub async fn init_db(db_url: &Url, name: &str) -> anyhow::Result<()> {
+/// Database configuration.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct DatabaseConfig {
+    /// Database URL.
+    pub url: Url,
+    /// Database name.
+    pub name: String,
+}
+
+impl DatabaseConfig {
+    /// Create a new `Db` instance.
+    pub fn new(url: Url, name: String) -> Self {
+        Self { url, name }
+    }
+
+    /// Create a new `Db` instance from a URL.
+    pub fn from_url(url: Url) -> anyhow::Result<Self> {
+        let name = url
+            .path_segments()
+            .ok_or(anyhow!("Failed to parse database name from URL"))?
+            .last()
+            .ok_or(anyhow!("Failed to parse database name from URL"))?;
+        let url_without_db_name = {
+            let mut url = url.clone();
+            url.set_path("");
+            url
+        };
+        Ok(Self {
+            url: url_without_db_name,
+            name: name.to_string(),
+        })
+    }
+
+    /// Get the full URL of the database.
+    pub fn full_url(&self) -> Url {
+        let mut url = self.url.clone();
+        url.set_path(&self.name);
+        url
+    }
+}
+
+pub async fn init_db(db: &DatabaseConfig) -> anyhow::Result<()> {
     // Connect to the database.
-    let mut connection = PgConnection::connect(db_url.as_ref()).await?;
+    let mut connection = PgConnection::connect(db.url.as_str()).await?;
 
-    let query = format!("CREATE DATABASE {}", name);
+    let query = format!("CREATE DATABASE {}", db.name);
     // Create DB.
     sqlx::query(&query).execute(&mut connection).await?;
 
     Ok(())
 }
 
-pub async fn drop_db_if_exists(db_url: &Url, name: &str) -> anyhow::Result<()> {
+pub async fn drop_db_if_exists(db: &DatabaseConfig) -> anyhow::Result<()> {
     // Connect to the database.
-    let mut connection = PgConnection::connect(db_url.as_ref()).await?;
+    let mut connection = PgConnection::connect(db.url.as_str()).await?;
 
-    let query = format!("DROP DATABASE IF EXISTS {}", name);
+    let query = format!("DROP DATABASE IF EXISTS {}", db.name);
     // DROP DB.
     sqlx::query(&query).execute(&mut connection).await?;
 
@@ -34,7 +77,7 @@ pub async fn drop_db_if_exists(db_url: &Url, name: &str) -> anyhow::Result<()> {
 pub async fn migrate_db(
     shell: &Shell,
     migrations_folder: PathBuf,
-    db_url: &str,
+    db_url: &Url,
 ) -> anyhow::Result<()> {
     // Most of this file is copy-pasted from SQLx CLI:
     // https://github.com/launchbadge/sqlx/blob/main/sqlx-cli/src/migrate.rs
@@ -45,7 +88,7 @@ pub async fn migrate_db(
     }
     let migrator = Migrator::new(migrations_folder).await?;
 
-    let mut conn = PgConnection::connect(db_url).await?;
+    let mut conn = PgConnection::connect(db_url.as_str()).await?;
     conn.ensure_migrations_table().await?;
 
     let version = conn.dirty_version().await?;
@@ -83,7 +126,7 @@ pub async fn migrate_db(
                 let text = if skip { "Skipped" } else { "Applied" };
 
                 if global_config().verbose {
-                    logger::raw(&format!(
+                    logger::step(&format!(
                         "    {} {}/{} {} ({elapsed:?})",
                         text,
                         migration.version,
@@ -104,3 +147,15 @@ pub async fn migrate_db(
 
     Ok(())
 }
+
+pub async fn wait_for_db(db_url: &Url, tries: u32) -> anyhow::Result<()> {
+    for i in 0..tries {
+        if PgConnection::connect(db_url.as_str()).await.is_ok() {
+            return Ok(());
+        }
+        if i < tries - 1 {
+            tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+        }
+    }
+    anyhow::bail!("Unable to connect to Postgres, connection cannot be established");
+}

zk_toolbox/crates/common/src/term/logger.rs

@@ -43,10 +43,14 @@ pub fn success(msg: impl Display) {
     log::success(msg).unwrap();
 }
 
-pub fn raw(msg: impl Display) {
+pub fn step(msg: impl Display) {
     log::step(msg).unwrap();
 }
 
+pub fn raw(msg: impl Display) {
+    term_write(msg);
+}
+
 pub fn note(msg: impl Display, content: impl Display) {
     cliclack::note(msg, content).unwrap();
 }

zk_toolbox/crates/common/src/term/spinner.rs

@@ -34,4 +34,13 @@ impl Spinner {
             self.time.elapsed().as_secs_f64()
         ));
     }
+
+    /// Interrupt the spinner with a failed message.
+    pub fn fail(self) {
+        self.pb.error(format!(
+            "{} failed in {} secs",
+            self.msg,
+            self.time.elapsed().as_secs_f64()
+        ));
+    }
 }

zk_toolbox/crates/config/src/secrets.rs

@@ -5,8 +5,8 @@ use crate::{consts::SECRETS_FILE, traits::FileConfigWithDefaultName};
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct DatabaseSecrets {
-    pub server_url: String,
-    pub prover_url: String,
+    pub server_url: Url,
+    pub prover_url: Url,
     #[serde(flatten)]
     pub other: serde_json::Value,
 }
@@ -29,28 +29,3 @@ pub struct SecretsConfig {
 impl FileConfigWithDefaultName for SecretsConfig {
     const FILE_NAME: &'static str = SECRETS_FILE;
 }
-
-#[derive(Debug, Serialize)]
-pub struct DatabaseConfig {
-    pub base_url: Url,
-    pub database_name: String,
-}
-
-impl DatabaseConfig {
-    pub fn new(base_url: Url, database_name: String) -> Self {
-        Self {
-            base_url,
-            database_name,
-        }
-    }
-
-    pub fn full_url(&self) -> String {
-        format!("{}/{}", self.base_url, self.database_name)
-    }
-}
-
-#[derive(Debug, Serialize)]
-pub struct DatabasesConfig {
-    pub server: DatabaseConfig,
-    pub prover: DatabaseConfig,
-}

zk_toolbox/crates/zk_inception/Cargo.toml

@@ -17,6 +17,7 @@ cliclack.workspace = true
 config.workspace = true
 console.workspace = true
 human-panic.workspace = true
+lazy_static.workspace = true
 serde_yaml.workspace = true
 serde.workspace = true
 serde_json.workspace = true