feat: improve database interactions

zk_toolbox/crates/common/Cargo.toml

@@ -23,7 +23,8 @@ serde_yaml.workspace = true
 sqlx.workspace = true
 strum.workspace = true
 strum_macros.workspace = true
+tokio.workspace = true
 toml.workspace = true
 url.workspace = true
 xshell.workspace = true
-futures.workspace = true
+futures.workspace = true

zk_toolbox/crates/common/src/db.rs

@@ -33,7 +33,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
@@ -44,7 +44,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?;
@@ -103,3 +103,27 @@ pub async fn migrate_db(
 
     Ok(())
 }
+
+pub async fn wait_for_db(url: &Url, tries: u32) -> anyhow::Result<()> {
+    dbg!(url.as_str());
+    for i in 0..tries {
+        if PgConnection::connect(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");
+}
+
+/// Splits the URL into the base URL and the database name.
+pub fn split_db_url(url: &Url) -> (Url, String) {
+    let db_name = url.path_segments().unwrap().last().unwrap();
+    let url_without_db_name = {
+        let mut url = url.clone();
+        url.set_path("");
+        url
+    };
+    (url_without_db_name, db_name.to_string())
+}

zk_toolbox/crates/common/src/prerequisites.rs

@@ -1,7 +1,7 @@
 use crate::{cmd::Cmd, logger};
 use xshell::{cmd, Shell};
 
-const PREREQUISITES: [Prerequisite; 7] = [
+const PREREQUISITES: [Prerequisite; 6] = [
     Prerequisite {
         name: "git",
         download_link: "https://git-scm.com/book/en/v2/Getting-Started-Installing-Git",
@@ -26,10 +26,6 @@ const PREREQUISITES: [Prerequisite; 7] = [
         name: "yarn",
         download_link: "https://yarnpkg.com/getting-started/install",
     },
-    Prerequisite {
-        name: "pg_isready",
-        download_link: "https://www.postgresql.org/download",
-    },
 ];
 
 struct Prerequisite {

zk_toolbox/crates/zk_inception/src/configs/secrets.rs

@@ -5,8 +5,8 @@ use crate::configs::{ReadConfig, SaveConfig};
 
 #[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,
 }
@@ -40,8 +40,8 @@ impl DatabaseConfig {
         }
     }
 
-    pub fn full_url(&self) -> String {
-        format!("{}/{}", self.base_url, self.database_name)
+    pub fn full_url(&self) -> Url {
+        self.base_url.join(&self.database_name).unwrap()
     }
 }
 

zk_toolbox/crates/zk_supervisor/Cargo.toml

@@ -25,5 +25,6 @@ human-panic.workspace = true
 strum.workspace = true
 strum_macros.workspace = true
 tokio.workspace = true
+url.workspace = true
 xshell.workspace = true
 zk_inception.workspace = true

zk_toolbox/crates/zk_supervisor/src/commands/database/check_sqlx_data.rs

@@ -39,7 +39,7 @@ pub fn check_sqlx_data(
 ) -> anyhow::Result<()> {
     let dir = link_to_code.as_ref().join(&dal.path);
     let _dir_guard = shell.push_dir(dir);
-    let url = dal.url;
+    let url = dal.url.as_str();
 
     let spinner = Spinner::new(&format!("Checking sqlx data for dal {}...", dal.path));
     Cmd::new(cmd!(

zk_toolbox/crates/zk_supervisor/src/commands/database/drop.rs

@@ -1,9 +1,13 @@
 use super::args::drop::DatabaseDropArgs;
 use crate::dals::{get_core_dal, get_prover_dal, Dal};
-use common::{cmd::Cmd, logger, spinner::Spinner};
-use xshell::{cmd, Shell};
+use common::{
+    db::{drop_db_if_exists, split_db_url},
+    logger,
+    spinner::Spinner,
+};
+use xshell::Shell;
 
-pub fn run(shell: &Shell, args: DatabaseDropArgs) -> anyhow::Result<()> {
+pub async fn run(shell: &Shell, args: DatabaseDropArgs) -> anyhow::Result<()> {
     let args = args.fill_values_with_prompt();
     if !args.common.prover && !args.common.core {
         logger::outro("No databases selected to drop");
@@ -13,25 +17,21 @@ pub fn run(shell: &Shell, args: DatabaseDropArgs) -> anyhow::Result<()> {
     logger::info("Dropping databases");
 
     if args.common.prover {
-        drop_database(shell, get_prover_dal(shell)?)?;
+        drop_database(get_prover_dal(shell)?).await?;
     }
     if args.common.core {
-        drop_database(shell, get_core_dal(shell)?)?;
+        drop_database(get_core_dal(shell)?).await?;
     }
 
     logger::outro("Databases dropped successfully");
 
     Ok(())
 }
 
-pub fn drop_database(shell: &Shell, dal: Dal) -> anyhow::Result<()> {
+pub async fn drop_database(dal: Dal) -> anyhow::Result<()> {
     let spinner = Spinner::new(&format!("Dropping DB for dal {}...", dal.path));
-    let url = dal.url;
-    Cmd::new(cmd!(
-        shell,
-        "cargo sqlx database drop -y --database-url {url}"
-    ))
-    .run()?;
+    let (url, db_name) = split_db_url(&dal.url);
+    drop_db_if_exists(&url, &db_name).await?;
     spinner.finish();
     Ok(())
 }

zk_toolbox/crates/zk_supervisor/src/commands/database/migrate.rs

@@ -34,7 +34,7 @@ pub fn run(shell: &Shell, args: DatabaseMigrateArgs) -> anyhow::Result<()> {
 fn migrate_database(shell: &Shell, link_to_code: impl AsRef<Path>, dal: Dal) -> anyhow::Result<()> {
     let dir = link_to_code.as_ref().join(&dal.path);
     let _dir_guard = shell.push_dir(dir);
-    let url = dal.url;
+    let url = dal.url.as_str();
 
     let spinner = Spinner::new(&format!("Migrating DB for dal {}...", dal.path));
     Cmd::new(cmd!(

zk_toolbox/crates/zk_supervisor/src/commands/database/mod.rs

@@ -33,14 +33,14 @@ pub enum DatabaseCommands {
     Setup(DatabaseSetupArgs),
 }
 
-pub fn run(shell: &Shell, args: DatabaseCommands) -> anyhow::Result<()> {
+pub async fn run(shell: &Shell, args: DatabaseCommands) -> anyhow::Result<()> {
     match args {
         DatabaseCommands::CheckSqlxData(args) => check_sqlx_data::run(shell, args),
-        DatabaseCommands::Drop(args) => drop::run(shell, args),
+        DatabaseCommands::Drop(args) => drop::run(shell, args).await,
         DatabaseCommands::Migrate(args) => migrate::run(shell, args),
         DatabaseCommands::NewMigration(args) => new_migration::run(shell, args),
         DatabaseCommands::Prepare(args) => prepare::run(shell, args),
-        DatabaseCommands::Reset(args) => reset::run(shell, args),
+        DatabaseCommands::Reset(args) => reset::run(shell, args).await,
         DatabaseCommands::Setup(args) => setup::run(shell, args),
     }
 }

zk_toolbox/crates/zk_supervisor/src/commands/database/prepare.rs

@@ -39,7 +39,7 @@ pub fn prepare_sqlx_data(
 ) -> anyhow::Result<()> {
     let dir = link_to_code.as_ref().join(&dal.path);
     let _dir_guard = shell.push_dir(dir);
-    let url = dal.url;
+    let url = dal.url.as_str();
 
     let spinner = Spinner::new(&format!("Preparing sqlx data for dal {}...", dal.path));
     Cmd::new(cmd!(

zk_toolbox/crates/zk_supervisor/src/commands/database/reset.rs

@@ -1,11 +1,11 @@
 use super::{args::reset::DatabaseResetArgs, drop::drop_database, setup::setup_database};
 use crate::dals::{get_core_dal, get_prover_dal, Dal};
-use common::{cmd::Cmd, logger, spinner::Spinner};
+use common::logger;
 use std::path::Path;
-use xshell::{cmd, Shell};
+use xshell::Shell;
 use zk_inception::configs::EcosystemConfig;
 
-pub fn run(shell: &Shell, args: DatabaseResetArgs) -> anyhow::Result<()> {
+pub async fn run(shell: &Shell, args: DatabaseResetArgs) -> anyhow::Result<()> {
     let args = args.fill_values_with_prompt();
     if !args.common.prover && !args.common.core {
         logger::outro("No databases selected");
@@ -20,50 +20,25 @@ pub fn run(shell: &Shell, args: DatabaseResetArgs) -> anyhow::Result<()> {
             shell,
             ecoseystem_config.link_to_code.clone(),
             get_prover_dal(shell)?,
-        )?;
+        )
+        .await?;
     }
     if args.common.core {
         logger::info("Resetting core database");
-        reset_database(shell, ecoseystem_config.link_to_code, get_core_dal(shell)?)?;
+        reset_database(shell, ecoseystem_config.link_to_code, get_core_dal(shell)?).await?;
     }
 
     logger::outro("Databases resetted");
 
     Ok(())
 }
 
-fn reset_database(shell: &Shell, link_to_code: impl AsRef<Path>, dal: Dal) -> anyhow::Result<()> {
-    wait_database(shell, dal.clone(), 5)?;
-    drop_database(shell, dal.clone())?;
+async fn reset_database(
+    shell: &Shell,
+    link_to_code: impl AsRef<Path>,
+    dal: Dal,
+) -> anyhow::Result<()> {
+    drop_database(dal.clone()).await?;
     setup_database(shell, link_to_code, dal)?;
     Ok(())
 }
-
-fn wait_database(shell: &Shell, dal: Dal, tries: u32) -> anyhow::Result<()> {
-    let url = dal.url;
-    let spinner = Spinner::new(&format!(
-        "Waiting until DB for dal {} is ready...",
-        dal.path
-    ));
-
-    for i in 0..tries {
-        let output = Cmd::new(cmd!(shell, "pg_isready -d {url}")).run_with_output()?;
-        if output.status.success() {
-            spinner.finish();
-            return Ok(());
-        }
-
-        // Only sleep if there are more tries left
-        if i < tries - 1 {
-            std::thread::sleep(std::time::Duration::from_secs(1))
-        }
-    }
-
-    // If we reach here, it means the database is not ready
-    spinner.fail();
-    anyhow::bail!(
-        "DB for dal {} is not ready after {} attempts",
-        dal.path,
-        tries
-    );
-}

zk_toolbox/crates/zk_supervisor/src/commands/database/setup.rs

@@ -39,7 +39,7 @@ pub fn setup_database(
 ) -> anyhow::Result<()> {
     let dir = link_to_code.as_ref().join(&dal.path);
     let _dir_guard = shell.push_dir(dir);
-    let url = dal.url;
+    let url = dal.url.as_str();
 
     let spinner = Spinner::new(&format!("Setting up DB for dal {}...", dal.path));
     Cmd::new(cmd!(

zk_toolbox/crates/zk_supervisor/src/dals.rs

@@ -1,5 +1,6 @@
 use anyhow::anyhow;
 use common::config::global_config;
+use url::Url;
 use xshell::Shell;
 use zk_inception::configs::{EcosystemConfig, Secrets};
 
@@ -9,7 +10,7 @@ const PROVER_DAL_PATH: &str = "prover/prover_dal";
 #[derive(Debug, Clone)]
 pub struct Dal {
     pub path: String,
-    pub url: String,
+    pub url: Url,
 }
 
 pub fn get_prover_dal(shell: &Shell) -> anyhow::Result<Dal> {

zk_toolbox/crates/zk_supervisor/src/main.rs

@@ -83,7 +83,7 @@ async fn main() -> anyhow::Result<()> {
 
 async fn run_subcommand(args: Supervisor, shell: &Shell) -> anyhow::Result<()> {
     match args.command {
-        SupervisorSubcommands::Database(command) => commands::database::run(shell, command)?,
+        SupervisorSubcommands::Database(command) => commands::database::run(shell, command).await?,
     }
     Ok(())
 }