diff --git a/src/bin/cratesfyi.rs b/src/bin/cratesfyi.rs index fc7c995b9..f0720c179 100644 --- a/src/bin/cratesfyi.rs +++ b/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"); + } + } } } } diff --git a/src/db/migrate.rs b/src/db/migrate.rs index 8edb63958..e533f1abb 100644 --- a/src/db/migrate.rs +++ b/src/db/migrate.rs @@ -295,6 +295,20 @@ pub fn migrate(version: Option, 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 { diff --git a/src/docbuilder/queue.rs b/src/docbuilder/queue.rs index 39027cb6a..cf0cc3e49 100644 --- a/src/docbuilder/queue.rs +++ b/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; } diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 68a76f715..423eb9d54 100644 --- a/src/utils/mod.rs +++ b/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; diff --git a/src/utils/queue.rs b/src/utils/queue.rs index a1961c041..abe5472bf 100644 --- a/src/utils/queue.rs +++ b/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 { + // 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> { + 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(()) + }) + } +}