migrate limits & overrides to sqlx / async, introduce async test wrapper

src/bin/cratesfyi.rs

@@ -615,44 +615,47 @@ enum LimitsSubcommand {
 
 impl LimitsSubcommand {
     fn handle_args(self, ctx: BinContext) -> Result<()> {
-        let conn = &mut *ctx.conn()?;
-        match self {
-            Self::Get { crate_name } => {
-                let overrides = Overrides::for_crate(conn, &crate_name)?;
-                println!("sandbox limit overrides for {crate_name} = {overrides:?}");
-            }
+        ctx.runtime()?.block_on(async move {
+            let mut conn = ctx.pool()?.get_async().await?;
 
-            Self::List => {
-                for (crate_name, overrides) in Overrides::all(conn)? {
+            match self {
+                Self::Get { crate_name } => {
+                    let overrides = Overrides::for_crate(&mut conn, &crate_name).await?;
                     println!("sandbox limit overrides for {crate_name} = {overrides:?}");
                 }
-            }
 
-            Self::Set {
-                crate_name,
-                memory,
-                targets,
-                timeout,
-            } => {
-                let overrides = Overrides::for_crate(conn, &crate_name)?;
-                println!("previous sandbox limit overrides for {crate_name} = {overrides:?}");
-                let overrides = Overrides {
+                Self::List => {
+                    for (crate_name, overrides) in Overrides::all(&mut conn).await? {
+                        println!("sandbox limit overrides for {crate_name} = {overrides:?}");
+                    }
+                }
+
+                Self::Set {
+                    crate_name,
                     memory,
                     targets,
-                    timeout: timeout.map(Into::into),
-                };
-                Overrides::save(conn, &crate_name, overrides)?;
-                let overrides = Overrides::for_crate(conn, &crate_name)?;
-                println!("new sandbox limit overrides for {crate_name} = {overrides:?}");
-            }
+                    timeout,
+                } => {
+                    let overrides = Overrides::for_crate(&mut conn, &crate_name).await?;
+                    println!("previous sandbox limit overrides for {crate_name} = {overrides:?}");
+                    let overrides = Overrides {
+                        memory,
+                        targets,
+                        timeout: timeout.map(Into::into),
+                    };
+                    Overrides::save(&mut conn, &crate_name, overrides).await?;
+                    let overrides = Overrides::for_crate(&mut conn, &crate_name).await?;
+                    println!("new sandbox limit overrides for {crate_name} = {overrides:?}");
+                }
 
-            Self::Remove { crate_name } => {
-                let overrides = Overrides::for_crate(conn, &crate_name)?;
-                println!("previous overrides for {crate_name} = {overrides:?}");
-                Overrides::remove(conn, &crate_name)?;
+                Self::Remove { crate_name } => {
+                    let overrides = Overrides::for_crate(&mut conn, &crate_name).await?;
+                    println!("previous overrides for {crate_name} = {overrides:?}");
+                    Overrides::remove(&mut conn, &crate_name).await?;
+                }
             }
-        }
-        Ok(())
+            Ok(())
+        })
     }
 }
 

src/db/overrides.rs

@@ -1,5 +1,6 @@
 use crate::error::Result;
-use postgres::Client;
+use futures_util::stream::TryStreamExt;
+use sqlx::{postgres::PgRow, FromRow, Row};
 use std::time::Duration;
 
 #[derive(Default, Debug, Clone, Copy, Eq, PartialEq)]
@@ -9,75 +10,83 @@ pub struct Overrides {
     pub timeout: Option<Duration>,
 }
 
-impl Overrides {
-    pub fn all(conn: &mut Client) -> Result<Vec<(String, Self)>> {
-        Ok(conn
-            .query("SELECT * FROM sandbox_overrides", &[])?
-            .into_iter()
-            .map(|row| (row.get("crate_name"), Self::from_row(row)))
-            .collect())
-    }
-
-    pub fn for_crate(conn: &mut Client, krate: &str) -> Result<Option<Self>> {
-        Ok(conn
-            .query_opt(
-                "SELECT * FROM sandbox_overrides WHERE crate_name = $1",
-                &[&krate],
-            )?
-            .map(Self::from_row))
-    }
-
-    fn from_row(row: postgres::Row) -> Self {
-        Self {
+impl FromRow<'_, PgRow> for Overrides {
+    fn from_row(row: &PgRow) -> sqlx::Result<Self> {
+        Ok(Self {
             memory: row
-                .get::<_, Option<i64>>("max_memory_bytes")
+                .get::<Option<i64>, _>("max_memory_bytes")
                 .map(|i| i as usize),
-            targets: row.get::<_, Option<i32>>("max_targets").map(|i| i as usize),
+            targets: row.get::<Option<i32>, _>("max_targets").map(|i| i as usize),
             timeout: row
-                .get::<_, Option<i32>>("timeout_seconds")
+                .get::<Option<i32>, _>("timeout_seconds")
                 .map(|i| Duration::from_secs(i as u64)),
-        }
+        })
+    }
+}
+
+impl Overrides {
+    pub async fn all(conn: &mut sqlx::PgConnection) -> Result<Vec<(String, Self)>> {
+        Ok(sqlx::query("SELECT * FROM sandbox_overrides")
+            .fetch(conn)
+            .map_ok(|row| {
+                (
+                    row.get("crate_name"),
+                    Overrides::from_row(&row)
+                        .expect("this is fine because we never return Err(_) in from_row"),
+                )
+            })
+            .try_collect()
+            .await?)
     }
 
-    pub fn save(conn: &mut Client, krate: &str, overrides: Self) -> Result<()> {
+    pub async fn for_crate(conn: &mut sqlx::PgConnection, krate: &str) -> Result<Option<Self>> {
+        Ok(
+            sqlx::query_as("SELECT * FROM sandbox_overrides WHERE crate_name = $1")
+                .bind(krate)
+                .fetch_optional(conn)
+                .await?,
+        )
+    }
+
+    pub async fn save(conn: &mut sqlx::PgConnection, krate: &str, overrides: Self) -> Result<()> {
         if overrides.timeout.is_some() && overrides.targets.is_none() {
             tracing::warn!("setting `Overrides::timeout` implies a default `Overrides::targets = 1`, prefer setting this explicitly");
         }
 
-        if conn
-            .query_opt("SELECT id FROM crates WHERE crates.name = $1", &[&krate])?
+        if sqlx::query_scalar!("SELECT id FROM crates WHERE crates.name = $1", krate)
+            .fetch_optional(&mut *conn)
+            .await?
             .is_none()
         {
             tracing::warn!("setting overrides for unknown crate `{krate}`");
         }
 
-        conn.execute(
+        sqlx::query!(
             "
-                INSERT INTO sandbox_overrides (
-                    crate_name, max_memory_bytes, max_targets, timeout_seconds
-                )
-                VALUES ($1, $2, $3, $4)
-                ON CONFLICT (crate_name) DO UPDATE
-                    SET
-                        max_memory_bytes = $2,
-                        max_targets = $3,
-                        timeout_seconds = $4
-                ",
-            &[
-                &krate,
-                &overrides.memory.map(|i| i as i64),
-                &overrides.targets.map(|i| i as i32),
-                &overrides.timeout.map(|d| d.as_secs() as i32),
-            ],
-        )?;
+            INSERT INTO sandbox_overrides (
+                crate_name, max_memory_bytes, max_targets, timeout_seconds
+            )
+            VALUES ($1, $2, $3, $4)
+            ON CONFLICT (crate_name) DO UPDATE
+                SET
+                    max_memory_bytes = $2,
+                    max_targets = $3,
+                    timeout_seconds = $4
+            ",
+            krate,
+            overrides.memory.map(|i| i as i64),
+            overrides.targets.map(|i| i as i32),
+            overrides.timeout.map(|d| d.as_secs() as i32),
+        )
+        .execute(&mut *conn)
+        .await?;
         Ok(())
     }
 
-    pub fn remove(conn: &mut Client, krate: &str) -> Result<()> {
-        conn.execute(
-            "DELETE FROM sandbox_overrides WHERE crate_name = $1",
-            &[&krate],
-        )?;
+    pub async fn remove(conn: &mut sqlx::PgConnection, krate: &str) -> Result<()> {
+        sqlx::query!("DELETE FROM sandbox_overrides WHERE crate_name = $1", krate)
+            .execute(conn)
+            .await?;
         Ok(())
     }
 }
@@ -89,22 +98,23 @@ mod test {
 
     #[test]
     fn retrieve_overrides() {
-        wrapper(|env| {
-            let db = env.db();
+        async_wrapper(|env| async move {
+            let db = env.async_db().await;
+            let mut conn = db.async_conn().await;
 
             let krate = "hexponent";
 
             // no overrides
-            let actual = Overrides::for_crate(&mut db.conn(), krate)?;
+            let actual = Overrides::for_crate(&mut conn, krate).await?;
             assert_eq!(actual, None);
 
             // add partial overrides
             let expected = Overrides {
                 targets: Some(1),
                 ..Overrides::default()
             };
-            Overrides::save(&mut db.conn(), krate, expected)?;
-            let actual = Overrides::for_crate(&mut db.conn(), krate)?;
+            Overrides::save(&mut conn, krate, expected).await?;
+            let actual = Overrides::for_crate(&mut conn, krate).await?;
             assert_eq!(actual, Some(expected));
 
             // overwrite with full overrides
@@ -113,25 +123,25 @@ mod test {
                 targets: Some(1),
                 timeout: Some(Duration::from_secs(300)),
             };
-            Overrides::save(&mut db.conn(), krate, expected)?;
-            let actual = Overrides::for_crate(&mut db.conn(), krate)?;
+            Overrides::save(&mut conn, krate, expected).await?;
+            let actual = Overrides::for_crate(&mut conn, krate).await?;
             assert_eq!(actual, Some(expected));
 
             // overwrite with partial overrides
             let expected = Overrides {
                 memory: Some(1),
                 ..Overrides::default()
             };
-            Overrides::save(&mut db.conn(), krate, expected)?;
-            let actual = Overrides::for_crate(&mut db.conn(), krate)?;
+            Overrides::save(&mut conn, krate, expected).await?;
+            let actual = Overrides::for_crate(&mut conn, krate).await?;
             assert_eq!(actual, Some(expected));
 
             // remove overrides
-            Overrides::remove(&mut db.conn(), krate)?;
-            let actual = Overrides::for_crate(&mut db.conn(), krate)?;
+            Overrides::remove(&mut conn, krate).await?;
+            let actual = Overrides::for_crate(&mut conn, krate).await?;
             assert_eq!(actual, None);
 
             Ok(())
-        });
+        })
     }
 }