models/default_versions: Add async fn variants

src/bin/crates-admin/default_versions.rs

@@ -1,10 +1,8 @@
 use anyhow::Context;
-use crates_io::models::{update_default_version, verify_default_version};
-use crates_io::tasks::spawn_blocking;
+use crates_io::models::{async_update_default_version, async_verify_default_version};
 use crates_io::{db, schema::crates};
 use diesel::prelude::*;
-use diesel_async::async_connection_wrapper::AsyncConnectionWrapper;
-use diesel_async::{AsyncPgConnection, RunQueryDsl};
+use diesel_async::RunQueryDsl;
 use indicatif::{ProgressBar, ProgressIterator, ProgressStyle};
 
 #[derive(clap::Parser, Debug, Eq, PartialEq)]
@@ -28,25 +26,21 @@ pub async fn run(command: Command) -> anyhow::Result<()> {
         .await
         .context("Failed to load crates")?;
 
-    let mut conn = AsyncConnectionWrapper::<AsyncPgConnection>::from(conn);
-    spawn_blocking(move || {
-        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 result = match command {
+            Command::Update => async_update_default_version(crate_id, &mut conn).await,
+            Command::Verify => async_verify_default_version(crate_id, &mut conn).await,
+        };
 
-            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) = result {
+            pb.suspend(|| warn!(%crate_id, %error, "Failed to update the default version"));
         }
+    }
 
-        Ok(())
-    })
-    .await?
+    Ok(())
 }

src/models.rs

@@ -1,7 +1,10 @@
 pub use self::action::{NewVersionOwnerAction, VersionAction, VersionOwnerAction};
 pub use self::category::{Category, CrateCategory, NewCategory};
 pub use self::crate_owner_invitation::{CrateOwnerInvitation, NewCrateOwnerInvitationOutcome};
-pub use self::default_versions::{update_default_version, verify_default_version};
+pub use self::default_versions::{
+    async_update_default_version, async_verify_default_version, update_default_version,
+    verify_default_version,
+};
 pub use self::deleted_crate::NewDeletedCrate;
 pub use self::dependency::{Dependency, DependencyKind, ReverseDependency};
 pub use self::download::VersionDownload;

src/models/default_versions.rs

@@ -2,6 +2,7 @@ use crate::schema::{default_versions, versions};
 use crate::sql::SemverVersion;
 use crate::util::diesel::prelude::*;
 use crate::util::diesel::Conn;
+use diesel_async::AsyncPgConnection;
 
 /// A subset of the columns of the `versions` table.
 ///
@@ -45,6 +46,44 @@ impl Ord for Version {
     }
 }
 
+/// Updates the `default_versions` table entry for the specified crate.
+///
+/// This function first loads all versions of the crate from the database,
+/// then determines the default version based on the following criteria:
+///
+/// 1. The highest non-prerelease version that is not yanked.
+/// 2. The highest non-yanked version.
+/// 3. The highest version.
+///
+/// The default version is then written to the `default_versions` table.
+#[instrument(skip(conn))]
+pub async fn async_update_default_version(
+    crate_id: i32,
+    conn: &mut AsyncPgConnection,
+) -> QueryResult<()> {
+    use diesel_async::RunQueryDsl;
+
+    let default_version = async_calculate_default_version(crate_id, conn).await?;
+
+    debug!(
+        "Updating default version to {} (id: {})…",
+        default_version.num, default_version.id
+    );
+
+    diesel::insert_into(default_versions::table)
+        .values((
+            default_versions::crate_id.eq(crate_id),
+            default_versions::version_id.eq(default_version.id),
+        ))
+        .on_conflict(default_versions::crate_id)
+        .do_update()
+        .set(default_versions::version_id.eq(default_version.id))
+        .execute(conn)
+        .await?;
+
+    Ok(())
+}
+
 /// Updates the `default_versions` table entry for the specified crate.
 ///
 /// This function first loads all versions of the crate from the database,
@@ -79,6 +118,42 @@ pub fn update_default_version(crate_id: i32, conn: &mut impl Conn) -> QueryResul
     Ok(())
 }
 
+/// Verifies that the default version for the specified crate is up-to-date.
+#[instrument(skip(conn))]
+pub async fn async_verify_default_version(
+    crate_id: i32,
+    conn: &mut AsyncPgConnection,
+) -> QueryResult<()> {
+    use diesel_async::RunQueryDsl;
+
+    let calculated = async_calculate_default_version(crate_id, conn).await?;
+
+    let saved = default_versions::table
+        .select(default_versions::version_id)
+        .filter(default_versions::crate_id.eq(crate_id))
+        .first::<i32>(conn)
+        .await
+        .optional()?;
+
+    if let Some(saved) = saved {
+        if saved == calculated.id {
+            debug!("Default version for crate {crate_id} is up to date");
+        } else {
+            warn!(
+                "Default version for crate {crate_id} is outdated (expected: {saved}, actual: {})",
+                calculated.id,
+            );
+        }
+    } else {
+        warn!(
+            "Default version for crate {crate_id} is missing (expected: {})",
+            calculated.id
+        );
+    }
+
+    Ok(())
+}
+
 /// Verifies that the default version for the specified crate is up-to-date.
 #[instrument(skip(conn))]
 pub fn verify_default_version(crate_id: i32, conn: &mut impl Conn) -> QueryResult<()> {
@@ -111,6 +186,25 @@ pub fn verify_default_version(crate_id: i32, conn: &mut impl Conn) -> QueryResul
     Ok(())
 }
 
+async fn async_calculate_default_version(
+    crate_id: i32,
+    conn: &mut AsyncPgConnection,
+) -> QueryResult<Version> {
+    use diesel::result::Error::NotFound;
+    use diesel_async::RunQueryDsl;
+
+    debug!("Loading all versions for the crate…");
+    let versions = versions::table
+        .filter(versions::crate_id.eq(crate_id))
+        .select(Version::as_returning())
+        .load::<Version>(conn)
+        .await?;
+
+    debug!("Found {} versions", versions.len());
+
+    versions.into_iter().max().ok_or(NotFound)
+}
+
 fn calculate_default_version(crate_id: i32, conn: &mut impl Conn) -> QueryResult<Version> {
     use diesel::result::Error::NotFound;
     use diesel::RunQueryDsl;