bin/crates-admin: Move spawn_blocking() calls into subcommands

Author: Turbo87

src/admin/default_versions.rs

@@ -1,4 +1,5 @@
 use crate::models::{update_default_version, verify_default_version};
+use crate::tasks::spawn_blocking;
 use crate::{db, schema::crates};
 use anyhow::Context;
 use diesel::prelude::*;
@@ -14,29 +15,32 @@ pub enum Command {
     Verify,
 }
 
-pub fn run(command: Command) -> anyhow::Result<()> {
-    let mut conn = db::oneoff_connection().context("Failed to connect to the database")?;
+pub async fn run(command: Command) -> anyhow::Result<()> {
+    spawn_blocking(move || {
+        let mut conn = db::oneoff_connection().context("Failed to connect to the database")?;
 
-    let crate_ids: Vec<i32> = crates::table
-        .select(crates::id)
-        .load(&mut conn)
-        .context("Failed to load crates")?;
+        let crate_ids: Vec<i32> = crates::table
+            .select(crates::id)
+            .load(&mut conn)
+            .context("Failed to load crates")?;
 
-    let pb = ProgressBar::new(crate_ids.len() as u64);
-    pb.set_style(ProgressStyle::with_template(
-        "{bar:60} ({pos}/{len}, ETA {eta})",
-    )?);
+        let pb = ProgressBar::new(crate_ids.len() as u64);
+        pb.set_style(ProgressStyle::with_template(
+            "{bar:60} ({pos}/{len}, ETA {eta})",
+        )?);
 
-    for crate_id in crate_ids.into_iter().progress_with(pb.clone()) {
-        let func = match command {
-            Command::Update => update_default_version,
-            Command::Verify => verify_default_version,
-        };
+        for crate_id in crate_ids.into_iter().progress_with(pb.clone()) {
+            let func = match command {
+                Command::Update => update_default_version,
+                Command::Verify => verify_default_version,
+            };
 
-        if let Err(error) = func(crate_id, &mut conn) {
-            pb.suspend(|| warn!(%crate_id, %error, "Failed to update the default version"));
+            if let Err(error) = func(crate_id, &mut conn) {
+                pb.suspend(|| warn!(%crate_id, %error, "Failed to update the default version"));
+            }
         }
-    }
 
-    Ok(())
+        Ok(())
+    })
+    .await
 }

src/admin/delete_crate.rs

@@ -1,4 +1,5 @@
 use crate::schema::{crate_owners, teams, users};
+use crate::tasks::spawn_blocking;
 use crate::worker::jobs;
 use crate::{admin::dialoguer, db, schema::crates};
 use anyhow::Context;
@@ -24,13 +25,15 @@ pub struct Opts {
     yes: bool,
 }
 
-pub fn run(opts: Opts) -> anyhow::Result<()> {
-    let conn = &mut db::oneoff_connection().context("Failed to establish database connection")?;
+pub async fn run(opts: Opts) -> anyhow::Result<()> {
+    spawn_blocking(move || {
+        let conn =
+            &mut db::oneoff_connection().context("Failed to establish database connection")?;
 
-    let mut crate_names = opts.crate_names;
-    crate_names.sort();
+        let mut crate_names = opts.crate_names;
+        crate_names.sort();
 
-    let query_result = crates::table
+        let query_result = crates::table
         .select((
             crates::name,
             crates::id,
@@ -45,53 +48,55 @@ pub fn run(opts: Opts) -> anyhow::Result<()> {
         .load::<(String, i32, String)>(conn)
         .context("Failed to look up crate name from the database")?;
 
-    let mut existing_crates: HashMap<String, (i32, Vec<String>)> = HashMap::new();
-    for (name, id, login) in query_result {
-        let entry = existing_crates
-            .entry(name)
-            .or_insert_with(|| (id, Vec::new()));
+        let mut existing_crates: HashMap<String, (i32, Vec<String>)> = HashMap::new();
+        for (name, id, login) in query_result {
+            let entry = existing_crates
+                .entry(name)
+                .or_insert_with(|| (id, Vec::new()));
 
-        entry.1.push(login);
-    }
+            entry.1.push(login);
+        }
 
-    println!("Deleting the following crates:");
-    println!();
-    for name in &crate_names {
-        match existing_crates.get(name) {
-            Some((id, owners)) => {
-                let owners = owners.join(", ");
-                println!(" - {name} (id={id}, owners={owners})");
+        println!("Deleting the following crates:");
+        println!();
+        for name in &crate_names {
+            match existing_crates.get(name) {
+                Some((id, owners)) => {
+                    let owners = owners.join(", ");
+                    println!(" - {name} (id={id}, owners={owners})");
+                }
+                None => println!(" - {name} (⚠️ crate not found)"),
             }
-            None => println!(" - {name} (⚠️ crate not found)"),
         }
-    }
-    println!();
+        println!();
 
-    if !opts.yes && !dialoguer::confirm("Do you want to permanently delete these crates?") {
-        return Ok(());
-    }
+        if !opts.yes && !dialoguer::confirm("Do you want to permanently delete these crates?") {
+            return Ok(());
+        }
 
-    for name in &crate_names {
-        if let Some((id, _)) = existing_crates.get(name) {
-            info!("{name}: Deleting crate from the database…");
-            if let Err(error) = diesel::delete(crates::table.find(id)).execute(conn) {
-                warn!(%id, "{name}: Failed to delete crate from the database: {error}");
-            }
-        } else {
-            info!("{name}: Skipped missing crate");
-        };
+        for name in &crate_names {
+            if let Some((id, _)) = existing_crates.get(name) {
+                info!("{name}: Deleting crate from the database…");
+                if let Err(error) = diesel::delete(crates::table.find(id)).execute(conn) {
+                    warn!(%id, "{name}: Failed to delete crate from the database: {error}");
+                }
+            } else {
+                info!("{name}: Skipped missing crate");
+            };
 
-        info!("{name}: Enqueuing index sync jobs…");
-        if let Err(error) = jobs::enqueue_sync_to_index(name, conn) {
-            warn!("{name}: Failed to enqueue index sync jobs: {error}");
-        }
+            info!("{name}: Enqueuing index sync jobs…");
+            if let Err(error) = jobs::enqueue_sync_to_index(name, conn) {
+                warn!("{name}: Failed to enqueue index sync jobs: {error}");
+            }
 
-        info!("{name}: Enqueuing DeleteCrateFromStorage job…");
-        let job = jobs::DeleteCrateFromStorage::new(name.into());
-        if let Err(error) = job.enqueue(conn) {
-            warn!("{name}: Failed to enqueue DeleteCrateFromStorage job: {error}");
+            info!("{name}: Enqueuing DeleteCrateFromStorage job…");
+            let job = jobs::DeleteCrateFromStorage::new(name.into());
+            if let Err(error) = job.enqueue(conn) {
+                warn!("{name}: Failed to enqueue DeleteCrateFromStorage job: {error}");
+            }
         }
-    }
 
-    Ok(())
+        Ok(())
+    })
+    .await
 }

src/admin/delete_version.rs

@@ -1,6 +1,7 @@
 use crate::models::update_default_version;
 use crate::schema::crates;
 use crate::storage::Storage;
+use crate::tasks::spawn_blocking;
 use crate::worker::jobs;
 use crate::{admin::dialoguer, db, schema::versions};
 use anyhow::Context;
@@ -25,86 +26,88 @@ pub struct Opts {
     yes: bool,
 }
 
-pub fn run(opts: Opts) -> anyhow::Result<()> {
-    let crate_name = &opts.crate_name;
-
-    let conn = &mut db::oneoff_connection().context("Failed to establish database connection")?;
-
-    let store = Storage::from_environment();
-
-    let crate_id: i32 = crates::table
-        .select(crates::id)
-        .filter(crates::name.eq(crate_name))
-        .first(conn)
-        .context("Failed to look up crate id from the database")?;
-
-    println!("Deleting the following versions of the `{crate_name}` crate:");
-    println!();
-    for version in &opts.versions {
-        println!(" - {version}");
-    }
-    println!();
-
-    if !opts.yes && !dialoguer::confirm("Do you want to permanently delete these versions?") {
-        return Ok(());
-    }
-
-    conn.transaction(|conn| {
-        info!(%crate_name, %crate_id, versions = ?opts.versions, "Deleting versions from the database");
-        let result = diesel::delete(
-            versions::table
-                .filter(versions::crate_id.eq(crate_id))
-                .filter(versions::num.eq_any(&opts.versions)),
-        )
-        .execute(conn);
-
-        match result {
-            Ok(num_deleted) if num_deleted == opts.versions.len() => {}
-            Ok(num_deleted) => {
-                warn!(
-                    %crate_name,
-                    "Deleted only {num_deleted} of {num_expected} versions from the database",
-                    num_expected = opts.versions.len()
-                );
-            }
-            Err(error) => {
-                warn!(%crate_name, ?error, "Failed to delete versions from the database")
-            }
+pub async fn run(opts: Opts) -> anyhow::Result<()> {
+    spawn_blocking(move || {
+        let crate_name = &opts.crate_name;
+
+        let conn = &mut db::oneoff_connection().context("Failed to establish database connection")?;
+
+        let store = Storage::from_environment();
+
+        let crate_id: i32 = crates::table
+            .select(crates::id)
+            .filter(crates::name.eq(crate_name))
+            .first(conn)
+            .context("Failed to look up crate id from the database")?;
+
+        println!("Deleting the following versions of the `{crate_name}` crate:");
+        println!();
+        for version in &opts.versions {
+            println!(" - {version}");
         }
+        println!();
 
-        info!(%crate_name, %crate_id, "Updating default version in the database");
-        if let Err(error) = update_default_version(crate_id, conn) {
-            warn!(%crate_name, %crate_id, ?error, "Failed to update default version");
+        if !opts.yes && !dialoguer::confirm("Do you want to permanently delete these versions?") {
+            return Ok(());
         }
 
-        Ok::<_, anyhow::Error>(())
-    })?;
+        conn.transaction(|conn| {
+            info!(%crate_name, %crate_id, versions = ?opts.versions, "Deleting versions from the database");
+            let result = diesel::delete(
+                versions::table
+                    .filter(versions::crate_id.eq(crate_id))
+                    .filter(versions::num.eq_any(&opts.versions)),
+            )
+            .execute(conn);
+
+            match result {
+                Ok(num_deleted) if num_deleted == opts.versions.len() => {}
+                Ok(num_deleted) => {
+                    warn!(
+                        %crate_name,
+                        "Deleted only {num_deleted} of {num_expected} versions from the database",
+                        num_expected = opts.versions.len()
+                    );
+                }
+                Err(error) => {
+                    warn!(%crate_name, ?error, "Failed to delete versions from the database")
+                }
+            }
 
-    info!(%crate_name, "Enqueuing index sync jobs");
-    if let Err(error) = jobs::enqueue_sync_to_index(crate_name, conn) {
-        warn!(%crate_name, ?error, "Failed to enqueue index sync jobs");
-    }
+            info!(%crate_name, %crate_id, "Updating default version in the database");
+            if let Err(error) = update_default_version(crate_id, conn) {
+                warn!(%crate_name, %crate_id, ?error, "Failed to update default version");
+            }
 
-    let rt = tokio::runtime::Builder::new_current_thread()
-        .enable_all()
-        .build()
-        .context("Failed to initialize tokio runtime")?;
+            Ok::<_, anyhow::Error>(())
+        })?;
 
-    for version in &opts.versions {
-        debug!(%crate_name, %version, "Deleting crate file from S3");
-        if let Err(error) = rt.block_on(store.delete_crate_file(crate_name, version)) {
-            warn!(%crate_name, %version, ?error, "Failed to delete crate file from S3");
+        info!(%crate_name, "Enqueuing index sync jobs");
+        if let Err(error) = jobs::enqueue_sync_to_index(crate_name, conn) {
+            warn!(%crate_name, ?error, "Failed to enqueue index sync jobs");
         }
 
-        debug!(%crate_name, %version, "Deleting readme file from S3");
-        match rt.block_on(store.delete_readme(crate_name, version)) {
-            Err(object_store::Error::NotFound { .. }) => {}
-            Err(error) => {
-                warn!(%crate_name, %version, ?error, "Failed to delete readme file from S3")
+        let rt = tokio::runtime::Builder::new_current_thread()
+            .enable_all()
+            .build()
+            .context("Failed to initialize tokio runtime")?;
+
+        for version in &opts.versions {
+            debug!(%crate_name, %version, "Deleting crate file from S3");
+            if let Err(error) = rt.block_on(store.delete_crate_file(crate_name, version)) {
+                warn!(%crate_name, %version, ?error, "Failed to delete crate file from S3");
+            }
+
+            debug!(%crate_name, %version, "Deleting readme file from S3");
+            match rt.block_on(store.delete_readme(crate_name, version)) {
+                Err(object_store::Error::NotFound { .. }) => {}
+                Err(error) => {
+                    warn!(%crate_name, %version, ?error, "Failed to delete readme file from S3")
+                }
+                Ok(_) => {}
             }
-            Ok(_) => {}
         }
-    }
 
-    Ok(())
+        Ok(())
+    }).await
 }