feat(Prover CLI): delete cmd

prover/prover_cli/Cargo.toml

@@ -13,6 +13,7 @@ categories.workspace = true
 tokio = { version = "1", features = ["rt-multi-thread", "macros"] }
 env_logger = "0.10"
 log = "0.4"
+dialoguer = "0.11"
 
 clap = { workspace = true, features = ["derive"] }
 tracing.workspace = true
@@ -29,4 +30,3 @@ zksync_types.workspace = true
 zksync_prover_fri_types.workspace = true
 zksync_prover_interface.workspace = true
 prover_dal.workspace = true
-

prover/prover_cli/src/cli.rs

@@ -1,6 +1,6 @@
 use clap::{command, Parser, Subcommand};
 
-use crate::commands::get_file_info;
+use crate::commands::{delete, get_file_info};
 
 pub const VERSION_STRING: &str = env!("CARGO_PKG_VERSION");
 
@@ -14,12 +14,14 @@ struct ProverCLI {
 #[derive(Subcommand)]
 enum ProverCommand {
     FileInfo(get_file_info::Args),
+    Delete(delete::Args),
 }
 
 pub async fn start() -> anyhow::Result<()> {
     let ProverCLI { command } = ProverCLI::parse();
     match command {
         ProverCommand::FileInfo(args) => get_file_info::run(args).await?,
+        ProverCommand::Delete(args) => delete::run(args).await?,
     };
 
     Ok(())

prover/prover_cli/src/commands/delete.rs

@@ -0,0 +1,92 @@
+use anyhow::{bail, Context};
+use clap::Args as ClapArgs;
+use dialoguer::{theme::ColorfulTheme, Input};
+use prover_dal::{Connection, ConnectionPool, Prover, ProverDal};
+use zksync_config::PostgresConfig;
+use zksync_env_config::FromEnv;
+use zksync_types::L1BatchNumber;
+
+#[derive(ClapArgs)]
+pub(crate) struct Args {
+    /// Batch number to delete
+    #[clap(short, long, conflicts_with = "all")]
+    batch: Option<L1BatchNumber>,
+    /// Delete data from all batches
+    #[clap(short, long, conflicts_with = "batch", default_value_t = false)]
+    all: bool,
+}
+
+pub(crate) async fn run(args: Args) -> anyhow::Result<()> {
+    if !args.all && args.batch.is_none() {
+        bail!("Either --all or --batch should be provided");
+    }
+
+    let confirmation = Input::<String>::with_theme(&ColorfulTheme::default())
+        .with_prompt("Are you sure you want to delete the data?")
+        .default("no".to_owned())
+        .interact_text()?;
+
+    if confirmation != "yes" {
+        println!("Aborted");
+        return Ok(());
+    }
+
+    let config = PostgresConfig::from_env()?;
+    let prover_connection_pool =
+        ConnectionPool::<Prover>::builder(config.prover_url()?, config.max_connections()?)
+            .build()
+            .await
+            .context("failed to build a prover_connection_pool")?;
+    let conn = prover_connection_pool.connection().await.unwrap();
+
+    if args.all {
+        delete_prover_db(conn).await?;
+    } else {
+        delete_batch_data(conn, args.batch.unwrap()).await?;
+    }
+
+    Ok(())
+}
+
+async fn delete_prover_db(mut conn: Connection<'_, Prover>) -> anyhow::Result<()> {
+    conn.fri_gpu_prover_queue_dal()
+        .delete()
+        .await
+        .context("failed to delete gpu prover queue")?;
+    conn.fri_prover_jobs_dal()
+        .delete()
+        .await
+        .context("failed to delete prover jobs")?;
+    conn.fri_protocol_versions_dal()
+        .delete()
+        .await
+        .context("failed to delete protocol versions")?;
+    conn.fri_proof_compressor_dal()
+        .delete()
+        .await
+        .context("failed to delete proof compressor")?;
+    conn.fri_witness_generator_dal()
+        .delete()
+        .await
+        .context("failed to delete witness generator")?;
+    Ok(())
+}
+
+async fn delete_batch_data(
+    mut conn: Connection<'_, Prover>,
+    block_number: L1BatchNumber,
+) -> anyhow::Result<()> {
+    conn.fri_proof_compressor_dal()
+        .delete_batch_data(block_number)
+        .await
+        .context("failed to delete proof compressor data")?;
+    conn.fri_prover_jobs_dal()
+        .delete_batch_data(block_number)
+        .await
+        .context("failed to delete prover jobs data")?;
+    conn.fri_witness_generator_dal()
+        .delete_batch_data(block_number)
+        .await
+        .context("failed to delete witness generator data")?;
+    Ok(())
+}

prover/prover_cli/src/commands/mod.rs

@@ -1 +1,2 @@
+pub(crate) mod delete;
 pub(crate) mod get_file_info;

prover/prover_dal/src/fri_gpu_prover_queue_dal.rs

@@ -209,4 +209,25 @@ impl FriGpuProverQueueDal<'_, '_> {
         .unwrap()
         .unwrap_or(0) as usize
     }
+
+    pub async fn delete_gpu_prover_queue_fri(
+        &mut self,
+    ) -> sqlx::Result<sqlx::postgres::PgQueryResult> {
+        sqlx::query!("DELETE FROM gpu_prover_queue_fri")
+            .execute(self.storage.conn())
+            .await
+    }
+
+    pub async fn delete_gpu_prover_queue_fri_archive(
+        &mut self,
+    ) -> sqlx::Result<sqlx::postgres::PgQueryResult> {
+        sqlx::query!("DELETE FROM gpu_prover_queue_fri")
+            .execute(self.storage.conn())
+            .await
+    }
+
+    pub async fn delete(&mut self) -> sqlx::Result<sqlx::postgres::PgQueryResult> {
+        self.delete_gpu_prover_queue_fri().await?;
+        self.delete_gpu_prover_queue_fri_archive().await
+    }
 }

prover/prover_dal/src/fri_proof_compressor_dal.rs

@@ -328,4 +328,26 @@ impl FriProofCompressorDal<'_, '_> {
             .collect()
         }
     }
+
+    pub async fn delete_batch_data(
+        &mut self,
+        block_number: L1BatchNumber,
+    ) -> sqlx::Result<sqlx::postgres::PgQueryResult> {
+        sqlx::query!(
+            r#"
+            DELETE FROM proof_compression_jobs_fri
+            WHERE
+                l1_batch_number = $1
+            "#,
+            i64::from(block_number.0)
+        )
+        .execute(self.storage.conn())
+        .await
+    }
+
+    pub async fn delete(&mut self) -> sqlx::Result<sqlx::postgres::PgQueryResult> {
+        sqlx::query!("DELETE FROM proof_compression_jobs_fri")
+            .execute(self.storage.conn())
+            .await
+    }
 }

prover/prover_dal/src/fri_protocol_versions_dal.rs

@@ -90,4 +90,10 @@ impl FriProtocolVersionsDal<'_, '_> {
         .map(|row| ProtocolVersionId::try_from(row.id as u16).unwrap())
         .collect()
     }
+
+    pub async fn delete(&mut self) -> sqlx::Result<sqlx::postgres::PgQueryResult> {
+        sqlx::query!("DELETE FROM prover_fri_protocol_versions")
+            .execute(self.storage.conn())
+            .await
+    }
 }