From 35d8955a4ee3d0443af54c8de56b95b46c5c2c0b Mon Sep 17 00:00:00 2001 From: aon <21188659+aon@users.noreply.github.com> Date: Fri, 24 May 2024 15:18:56 -0300 Subject: [PATCH] feat: improve database interactions --- zk_toolbox/Cargo.lock | 2 + zk_toolbox/crates/common/Cargo.toml | 3 +- zk_toolbox/crates/common/src/db.rs | 28 ++++++++++- zk_toolbox/crates/common/src/prerequisites.rs | 6 +-- .../zk_inception/src/configs/secrets.rs | 8 +-- zk_toolbox/crates/zk_supervisor/Cargo.toml | 1 + .../src/commands/database/args/wait.rs | 32 ------------ .../src/commands/database/check_sqlx_data.rs | 2 +- .../src/commands/database/drop.rs | 24 ++++----- .../src/commands/database/migrate.rs | 2 +- .../src/commands/database/mod.rs | 6 +-- .../src/commands/database/prepare.rs | 2 +- .../src/commands/database/reset.rs | 49 +++++-------------- .../src/commands/database/setup.rs | 2 +- zk_toolbox/crates/zk_supervisor/src/dals.rs | 3 +- zk_toolbox/crates/zk_supervisor/src/main.rs | 2 +- 16 files changed, 70 insertions(+), 102 deletions(-) delete mode 100644 zk_toolbox/crates/zk_supervisor/src/commands/database/args/wait.rs diff --git a/zk_toolbox/Cargo.lock b/zk_toolbox/Cargo.lock index ef855b7d5ef4..35fdceb935be 100644 --- a/zk_toolbox/Cargo.lock +++ b/zk_toolbox/Cargo.lock @@ -532,6 +532,7 @@ dependencies = [ "sqlx", "strum 0.26.2", "strum_macros 0.26.2", + "tokio", "toml", "url", "xshell", @@ -4525,6 +4526,7 @@ dependencies = [ "strum 0.26.2", "strum_macros 0.26.2", "tokio", + "url", "xshell", "zk_inception", ] diff --git a/zk_toolbox/crates/common/Cargo.toml b/zk_toolbox/crates/common/Cargo.toml index 588254e445f8..9da29e639d53 100644 --- a/zk_toolbox/crates/common/Cargo.toml +++ b/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 \ No newline at end of file +futures.workspace = true diff --git a/zk_toolbox/crates/common/src/db.rs b/zk_toolbox/crates/common/src/db.rs index 36f134958897..5dfdd1f55c71 100644 --- a/zk_toolbox/crates/common/src/db.rs +++ b/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()) +} diff --git a/zk_toolbox/crates/common/src/prerequisites.rs b/zk_toolbox/crates/common/src/prerequisites.rs index 9003ce4126cf..7551b247c681 100644 --- a/zk_toolbox/crates/common/src/prerequisites.rs +++ b/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 { diff --git a/zk_toolbox/crates/zk_inception/src/configs/secrets.rs b/zk_toolbox/crates/zk_inception/src/configs/secrets.rs index e95dd05df6a2..5c0d93fd2f31 100644 --- a/zk_toolbox/crates/zk_inception/src/configs/secrets.rs +++ b/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() } } diff --git a/zk_toolbox/crates/zk_supervisor/Cargo.toml b/zk_toolbox/crates/zk_supervisor/Cargo.toml index 20543ced1db5..8cca6e9a75e4 100644 --- a/zk_toolbox/crates/zk_supervisor/Cargo.toml +++ b/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 diff --git a/zk_toolbox/crates/zk_supervisor/src/commands/database/args/wait.rs b/zk_toolbox/crates/zk_supervisor/src/commands/database/args/wait.rs deleted file mode 100644 index 6f951bdb37e8..000000000000 --- a/zk_toolbox/crates/zk_supervisor/src/commands/database/args/wait.rs +++ /dev/null @@ -1,32 +0,0 @@ -use super::DatabaseCommonArgs; -use clap::Parser; - -#[derive(Debug, Parser)] -pub struct DatabaseWaitArgs { - #[clap(flatten)] - pub common: DatabaseCommonArgs, - /// Number of tries to wait for the database to be ready - #[clap(long, default_value = "5")] - pub tries: u32, -} - -impl DatabaseWaitArgs { - pub fn fill_values_with_prompt(self) -> DatabaseWaitArgsFinal { - let common = self.common.fill_values_with_prompt("wait for"); - - DatabaseWaitArgsFinal { - prover: common.prover, - core: common.core, - chain: common.chain, - tries: self.tries, - } - } -} - -#[derive(Debug)] -pub struct DatabaseWaitArgsFinal { - pub prover: bool, - pub core: bool, - pub chain: Option, - pub tries: u32, -} diff --git a/zk_toolbox/crates/zk_supervisor/src/commands/database/check_sqlx_data.rs b/zk_toolbox/crates/zk_supervisor/src/commands/database/check_sqlx_data.rs index a30b197187fd..d79a3607802e 100644 --- a/zk_toolbox/crates/zk_supervisor/src/commands/database/check_sqlx_data.rs +++ b/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!( diff --git a/zk_toolbox/crates/zk_supervisor/src/commands/database/drop.rs b/zk_toolbox/crates/zk_supervisor/src/commands/database/drop.rs index 4b7b2bbaebc4..5de56f650904 100644 --- a/zk_toolbox/crates/zk_supervisor/src/commands/database/drop.rs +++ b/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,10 +17,10 @@ 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"); @@ -24,14 +28,10 @@ pub fn run(shell: &Shell, args: DatabaseDropArgs) -> anyhow::Result<()> { 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(()) } diff --git a/zk_toolbox/crates/zk_supervisor/src/commands/database/migrate.rs b/zk_toolbox/crates/zk_supervisor/src/commands/database/migrate.rs index 0c5adb2d9cdb..c1537f643c2a 100644 --- a/zk_toolbox/crates/zk_supervisor/src/commands/database/migrate.rs +++ b/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, 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!( diff --git a/zk_toolbox/crates/zk_supervisor/src/commands/database/mod.rs b/zk_toolbox/crates/zk_supervisor/src/commands/database/mod.rs index 0efaa70b48aa..05147dcdcacb 100644 --- a/zk_toolbox/crates/zk_supervisor/src/commands/database/mod.rs +++ b/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), } } diff --git a/zk_toolbox/crates/zk_supervisor/src/commands/database/prepare.rs b/zk_toolbox/crates/zk_supervisor/src/commands/database/prepare.rs index 1efd515fddc6..cf7d82524330 100644 --- a/zk_toolbox/crates/zk_supervisor/src/commands/database/prepare.rs +++ b/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!( diff --git a/zk_toolbox/crates/zk_supervisor/src/commands/database/reset.rs b/zk_toolbox/crates/zk_supervisor/src/commands/database/reset.rs index df96d394de34..2d177144056f 100644 --- a/zk_toolbox/crates/zk_supervisor/src/commands/database/reset.rs +++ b/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,11 +20,12 @@ 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"); @@ -32,38 +33,12 @@ pub fn run(shell: &Shell, args: DatabaseResetArgs) -> anyhow::Result<()> { Ok(()) } -fn reset_database(shell: &Shell, link_to_code: impl AsRef, 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, + 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 - ); -} diff --git a/zk_toolbox/crates/zk_supervisor/src/commands/database/setup.rs b/zk_toolbox/crates/zk_supervisor/src/commands/database/setup.rs index 606c503d1c2e..48175a8d112c 100644 --- a/zk_toolbox/crates/zk_supervisor/src/commands/database/setup.rs +++ b/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!( diff --git a/zk_toolbox/crates/zk_supervisor/src/dals.rs b/zk_toolbox/crates/zk_supervisor/src/dals.rs index c9ebdcfc042e..99dad59b68f6 100644 --- a/zk_toolbox/crates/zk_supervisor/src/dals.rs +++ b/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 { diff --git a/zk_toolbox/crates/zk_supervisor/src/main.rs b/zk_toolbox/crates/zk_supervisor/src/main.rs index 97869e64fa01..50f4e303e9d0 100644 --- a/zk_toolbox/crates/zk_supervisor/src/main.rs +++ b/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(()) }