Automatic prioritization of crates

src/bin/cratesfyi.rs

@@ -3,7 +3,7 @@ use std::path::{Path, PathBuf};
 
 use clap::{App, AppSettings, Arg, SubCommand};
 use cratesfyi::db::{add_path_into_database, connect_db};
-use cratesfyi::utils::add_crate_to_queue;
+use cratesfyi::utils::{add_crate_to_queue, remove_crate_priority, set_crate_priority};
 use cratesfyi::Server;
 use cratesfyi::{db, DocBuilder, DocBuilderOptions, RustwideBuilder};
 
@@ -139,7 +139,26 @@ pub fn main() {
                     .short("p")
                     .long("priority")
                     .help("Priority of build (default: 5) (new crate builds get priority 0)")
-                    .takes_value(true))))
+                    .takes_value(true)))
+            .subcommand(SubCommand::with_name("default-priority")
+                .about("Interactions with build queue priorities")
+                .setting(AppSettings::ArgRequiredElseHelp)
+                .subcommand(SubCommand::with_name("set")
+                    .about("Set all crates matching the given pattern to a priority level")
+                    .arg(Arg::with_name("PATTERN")
+                        .index(1)
+                        .required(true)
+                        .help("See https://www.postgresql.org/docs/current/functions-matching.html"))
+                    .arg(Arg::with_name("PRIORITY")
+                        .index(2)
+                        .required(true)
+                        .help("The priority to give crates matching PATTERN")))
+                .subcommand(SubCommand::with_name("remove")
+                    .about("Remove the prioritization of crates by the given pattern")
+                    .arg(Arg::with_name("PATTERN")
+                        .index(1)
+                        .required(true)
+                        .help("See https://www.postgresql.org/docs/current/functions-matching.html")))))
         .get_matches();
 
     if let Some(matches) = matches.subcommand_matches("build") {
@@ -298,6 +317,31 @@ pub fn main() {
                 priority,
             )
             .expect("Could not add crate to queue");
+        } else if let Some(matches) = matches.subcommand_matches("default-priority") {
+            if let Some(matches) = matches.subcommand_matches("set") {
+                let pattern = matches
+                    .value_of("PATTERN")
+                    .expect("You must give a pattern to match with");
+                let priority = clap::value_t!(matches.value_of("PRIORITY"), i32)
+                    .expect("You must give a priority for a crate");
+                let conn = connect_db().expect("Could not connect to the database");
+
+                set_crate_priority(&conn, pattern, priority)
+                    .expect("Could not set pattern's priority");
+            } else if let Some(matches) = matches.subcommand_matches("remove") {
+                let pattern = matches
+                    .value_of("PATTERN")
+                    .expect("You must give a pattern to remove");
+                let conn = connect_db().expect("Could not connect to the database");
+
+                if let Some(priority) = remove_crate_priority(&conn, pattern)
+                    .expect("Could not remove pattern's priority")
+                {
+                    println!("Removed pattern with priority {}", priority);
+                } else {
+                    println!("Pattern did not exist and so was not removed");
+                }
+            }
         }
     }
 }

src/db/migrate.rs

@@ -295,6 +295,20 @@ pub fn migrate(version: Option<Version>, conn: &Connection) -> CratesfyiResult<(
                 DROP FUNCTION normalize_crate_name;
             "
         ),
+        migration!(
+            context,
+            // version
+            11,
+            // description
+            "Allow crates to be given a different default priority",
+            // upgrade query
+            "CREATE TABLE crate_priorities (
+                pattern VARCHAR NOT NULL UNIQUE,
+                priority INTEGER NOT NULL
+            );",
+            // downgrade query
+            "DROP TABLE crate_priorities;",
+        ),
     ];
 
     for migration in migrations {

src/docbuilder/queue.rs

@@ -3,7 +3,7 @@
 use super::{DocBuilder, RustwideBuilder};
 use crate::db::connect_db;
 use crate::error::Result;
-use crate::utils::add_crate_to_queue;
+use crate::utils::{add_crate_to_queue, get_crate_priority};
 use crates_index_diff::{ChangeKind, Index};
 
 impl DocBuilder {
@@ -19,7 +19,9 @@ impl DocBuilder {
         changes.reverse();
 
         for krate in changes.iter().filter(|k| k.kind != ChangeKind::Yanked) {
-            add_crate_to_queue(&conn, &krate.name, &krate.version, 0).ok();
+            let priority = get_crate_priority(&conn, &krate.name)?;
+            add_crate_to_queue(&conn, &krate.name, &krate.version, priority).ok();
+
             debug!("{}-{} added into build queue", krate.name, krate.version);
             add_count += 1;
         }

src/utils/mod.rs

@@ -5,7 +5,9 @@ pub(crate) use self::copy::copy_doc_dir;
 pub use self::daemon::start_daemon;
 pub use self::github_updater::github_updater;
 pub use self::html::extract_head_and_body;
-pub use self::queue::add_crate_to_queue;
+pub use self::queue::{
+    add_crate_to_queue, get_crate_priority, remove_crate_priority, set_crate_priority,
+};
 pub use self::release_activity_updater::update_release_activity;
 pub(crate) use self::rustc_version::parse_rustc_version;
 

src/utils/queue.rs

@@ -3,6 +3,50 @@
 use crate::error::Result;
 use postgres::Connection;
 
+const DEFAULT_PRIORITY: i32 = 0;
+
+/// Get the build queue priority for a crate
+pub fn get_crate_priority(conn: &Connection, name: &str) -> Result<i32> {
+    // Search the `priority` table for a priority where the crate name matches the stored pattern
+    let query = conn.query(
+        "SELECT priority FROM crate_priorities WHERE $1 LIKE pattern LIMIT 1",
+        &[&name],
+    )?;
+
+    // If no match is found, return the default priority
+    if let Some(row) = query.iter().next() {
+        Ok(row.get(0))
+    } else {
+        Ok(DEFAULT_PRIORITY)
+    }
+}
+
+/// Set all crates that match [`pattern`] to have a certain priority
+///
+/// Note: `pattern` is used in a `LIKE` statement, so it must follow the postgres like syntax
+///
+/// [`pattern`]: https://www.postgresql.org/docs/8.3/functions-matching.html
+pub fn set_crate_priority(conn: &Connection, pattern: &str, priority: i32) -> Result<()> {
+    conn.query(
+        "INSERT INTO crate_priorities (pattern, priority) VALUES ($1, $2)",
+        &[&pattern, &priority],
+    )?;
+
+    Ok(())
+}
+
+/// Remove a pattern from the priority table, returning the priority that it was associated with or `None`
+/// if nothing was removed
+pub fn remove_crate_priority(conn: &Connection, pattern: &str) -> Result<Option<i32>> {
+    let query = conn.query(
+        "DELETE FROM crate_priorities WHERE pattern = $1 RETURNING priority",
+        &[&pattern],
+    )?;
+
+    Ok(query.iter().next().map(|row| row.get(0)))
+}
+
+/// Adds a crate to the build queue to be built by rustdoc. `priority` should be gotten from `get_crate_priority`
 pub fn add_crate_to_queue(
     conn: &Connection,
     name: &str,
@@ -13,5 +57,144 @@ pub fn add_crate_to_queue(
         "INSERT INTO queue (name, version, priority) VALUES ($1, $2, $3)",
         &[&name, &version, &priority],
     )?;
+
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::test::wrapper;
+
+    #[test]
+    fn set_priority() {
+        wrapper(|env| {
+            let db = env.db();
+
+            set_crate_priority(&db.conn(), "cratesfyi-%", -100)?;
+            assert_eq!(get_crate_priority(&db.conn(), "cratesfyi-database")?, -100);
+            assert_eq!(get_crate_priority(&db.conn(), "cratesfyi-")?, -100);
+            assert_eq!(get_crate_priority(&db.conn(), "cratesfyi-s3")?, -100);
+            assert_eq!(get_crate_priority(&db.conn(), "cratesfyi-webserver")?, -100);
+            assert_eq!(
+                get_crate_priority(&db.conn(), "cratesfyi")?,
+                DEFAULT_PRIORITY
+            );
+
+            set_crate_priority(&db.conn(), "_c_", 100)?;
+            assert_eq!(get_crate_priority(&db.conn(), "rcc")?, 100);
+            assert_eq!(get_crate_priority(&db.conn(), "rc")?, DEFAULT_PRIORITY);
+
+            set_crate_priority(&db.conn(), "hexponent", 10)?;
+            assert_eq!(get_crate_priority(&db.conn(), "hexponent")?, 10);
+            assert_eq!(
+                get_crate_priority(&db.conn(), "hexponents")?,
+                DEFAULT_PRIORITY
+            );
+            assert_eq!(
+                get_crate_priority(&db.conn(), "floathexponent")?,
+                DEFAULT_PRIORITY
+            );
+
+            Ok(())
+        })
+    }
+
+    #[test]
+    fn remove_priority() {
+        wrapper(|env| {
+            let db = env.db();
+
+            set_crate_priority(&db.conn(), "cratesfyi-%", -100)?;
+            assert_eq!(get_crate_priority(&db.conn(), "cratesfyi-")?, -100);
+
+            assert_eq!(
+                remove_crate_priority(&db.conn(), "cratesfyi-%")?,
+                Some(-100)
+            );
+            assert_eq!(
+                get_crate_priority(&db.conn(), "cratesfyi-")?,
+                DEFAULT_PRIORITY
+            );
+
+            Ok(())
+        })
+    }
+
+    #[test]
+    fn get_priority() {
+        wrapper(|env| {
+            let db = env.db();
+
+            set_crate_priority(&db.conn(), "cratesfyi-%", -100)?;
+
+            assert_eq!(get_crate_priority(&db.conn(), "cratesfyi-database")?, -100);
+            assert_eq!(get_crate_priority(&db.conn(), "cratesfyi-")?, -100);
+            assert_eq!(get_crate_priority(&db.conn(), "cratesfyi-s3")?, -100);
+            assert_eq!(get_crate_priority(&db.conn(), "cratesfyi-webserver")?, -100);
+            assert_eq!(
+                get_crate_priority(&db.conn(), "unrelated")?,
+                DEFAULT_PRIORITY
+            );
+
+            Ok(())
+        })
+    }
+
+    #[test]
+    fn get_default_priority() {
+        wrapper(|env| {
+            let db = env.db();
+
+            assert_eq!(
+                get_crate_priority(&db.conn(), "cratesfyi")?,
+                DEFAULT_PRIORITY
+            );
+            assert_eq!(get_crate_priority(&db.conn(), "rcc")?, DEFAULT_PRIORITY);
+            assert_eq!(get_crate_priority(&db.conn(), "lasso")?, DEFAULT_PRIORITY);
+            assert_eq!(
+                get_crate_priority(&db.conn(), "hexponent")?,
+                DEFAULT_PRIORITY
+            );
+            assert_eq!(
+                get_crate_priority(&db.conn(), "rust4lyfe")?,
+                DEFAULT_PRIORITY
+            );
+
+            Ok(())
+        })
+    }
+
+    #[test]
+    fn add_to_queue() {
+        wrapper(|env| {
+            let db = env.db();
+
+            let test_crates = [
+                ("rcc", "0.1.0", 2),
+                ("lasso", "0.1.0", -1),
+                ("hexponent", "0.1.0", 0),
+                ("destroy-all-humans", "0.0.0-alpha", -100000),
+                ("totally-not-destroying-humans", "0.0.1", 0),
+            ];
+
+            for (name, version, priority) in test_crates.iter() {
+                add_crate_to_queue(&db.conn(), name, version, *priority)?;
+
+                let query = db.conn().query(
+                    "SELECT name, version, priority FROM queue WHERE name = $1",
+                    &[&name],
+                )?;
+
+                assert!(query.len() == 1);
+                let row = query.iter().next().unwrap();
+
+                assert_eq!(&row.get::<_, String>(0), name);
+                assert_eq!(&row.get::<_, String>(1), version);
+                assert_eq!(row.get::<_, i32>(2), *priority);
+            }
+
+            Ok(())
+        })
+    }
+}