Automatically create and monitor partitions for version_downloads

This adds a new background job which is responsible for creating new
partitions for `version_downloads` 1 year in advance. The intent here
is to give an extremely large window to notice if for some reason this
job stops running. In the extremely unlikely event that a full year
passes, data will start going in the default partition and we will get
paged. Recovering from this will be a massive pain, but I've tried to
ensure it'll never happen.

Author: sgrif

diesel.toml

@@ -9,7 +9,6 @@ patch_file = "src/schema.patch"
 [print_schema.filter]
 except_tables = [
   "version_downloads_archive",
-  "version_downloads_default",
   "version_downloads_pre_2017",
   "version_downloads_2017",
   "version_downloads_2018_q1",

src/bin/enqueue-job.rs

@@ -11,6 +11,9 @@ fn main() -> Result<(), Error> {
     println!("Enqueueing background job: {}", job);
 
     match &*job {
+        "generate_version_downloads_partition" => {
+            Ok(tasks::generate_version_downloads_partition().enqueue(&conn)?)
+        }
         "update_downloads" => Ok(tasks::update_downloads().enqueue(&conn)?),
         "dump_db" => {
             let database_url = args.next().unwrap_or_else(|| env("READ_ONLY_REPLICA_URL"));

src/bin/monitor.rs

@@ -16,6 +16,7 @@ fn main() -> Result<(), Error> {
 
     check_stalled_background_jobs(&conn)?;
     check_spam_attack(&conn)?;
+    check_default_version_downloads_partition(&conn)?;
     Ok(())
 }
 
@@ -116,6 +117,34 @@ fn check_spam_attack(conn: &PgConnection) -> Result<(), Error> {
     Ok(())
 }
 
+fn check_default_version_downloads_partition(conn: &PgConnection) -> Result<(), Error> {
+    use cargo_registry::schema::version_downloads_default::dsl::*;
+
+    const EVENT_KEY: &str = "version_downloads_missing_partition";
+
+    println!("Checking for data in the default `version_downloads` partition");
+    let version_downloads_in_default_partition =
+        version_downloads_default.count().get_result::<i64>(conn)?;
+
+    let event = if version_downloads_in_default_partition > 0 {
+        on_call::Event::Trigger {
+            incident_key: Some(EVENT_KEY.into()),
+            description: format!(
+                "{} rows exist in the default `version_downloads` partition",
+                version_downloads_in_default_partition
+            ),
+        }
+    } else {
+        on_call::Event::Resolve {
+            incident_key: EVENT_KEY.into(),
+            description: Some("No records in default `version_downloads` partition".into()),
+        }
+    };
+
+    log_and_trigger_event(event)?;
+    Ok(())
+}
+
 fn log_and_trigger_event(event: on_call::Event) -> Result<(), Error> {
     match event {
         on_call::Event::Trigger {

src/schema.rs

@@ -866,6 +866,41 @@ table! {
     }
 }
 
+table! {
+    use diesel::sql_types::*;
+    use diesel_full_text_search::{TsVector as Tsvector};
+
+    /// Representation of the `version_downloads_default` table.
+    ///
+    /// (Automatically generated by Diesel.)
+    version_downloads_default (version_id, date) {
+        /// The `version_id` column of the `version_downloads_default` table.
+        ///
+        /// Its SQL type is `Int4`.
+        ///
+        /// (Automatically generated by Diesel.)
+        version_id -> Int4,
+        /// The `downloads` column of the `version_downloads_default` table.
+        ///
+        /// Its SQL type is `Int4`.
+        ///
+        /// (Automatically generated by Diesel.)
+        downloads -> Int4,
+        /// The `counted` column of the `version_downloads_default` table.
+        ///
+        /// Its SQL type is `Int4`.
+        ///
+        /// (Automatically generated by Diesel.)
+        counted -> Int4,
+        /// The `date` column of the `version_downloads_default` table.
+        ///
+        /// Its SQL type is `Date`.
+        ///
+        /// (Automatically generated by Diesel.)
+        date -> Date,
+    }
+}
+
 table! {
     use diesel::sql_types::*;
     use diesel_full_text_search::{TsVector as Tsvector};
@@ -1066,6 +1101,7 @@ allow_tables_to_appear_in_same_query!(
     users,
     version_authors,
     version_downloads,
+    version_downloads_default,
     version_owner_actions,
     versions,
     versions_published_by,

src/tasks.rs

@@ -1,5 +1,9 @@
 pub mod dump_db;
+mod generate_version_downloads_partition;
+#[cfg(test)]
+mod test_helpers;
 mod update_downloads;
 
 pub use dump_db::dump_db;
+pub use generate_version_downloads_partition::generate_version_downloads_partition;
 pub use update_downloads::update_downloads;

src/tasks/generate_version_downloads_partition.rs

@@ -0,0 +1,120 @@
+use crate::background_jobs::Environment;
+
+use chrono::{Datelike, Duration, NaiveDate, Utc};
+use diesel::prelude::*;
+use diesel::sql_query;
+use swirl::PerformError;
+
+table! {
+    information_schema.tables (table_schema, table_name) {
+        table_schema -> Text,
+        table_name -> Text,
+        table_type -> Text,
+    }
+}
+
+#[derive(PartialEq, Debug)]
+struct Quarter {
+    start: NaiveDate,
+    end: NaiveDate,
+    num: u32,
+}
+
+impl Quarter {
+    fn from_date(date: NaiveDate) -> Option<Self> {
+        let num = (date.month() + 2) / 3;
+        let start = date.with_day(1)?.with_month(num * 3 - 2)?;
+        let end = (start + Duration::days(93)).with_day(1)?;
+        Some(Self { num, start, end })
+    }
+}
+
+#[swirl::background_job]
+pub fn generate_version_downloads_partition(env: &Environment) -> Result<(), PerformError> {
+    let conn = env.connection()?;
+    generate_partition(&conn, Utc::today().naive_utc() + Duration::days(365))?;
+    Ok(())
+}
+
+fn generate_partition(conn: &PgConnection, today: NaiveDate) -> QueryResult<()> {
+    let quarter = Quarter::from_date(today).expect("could not determine start/end of quarter");
+    let table_name = format!("version_downloads_{}_q{}", today.year(), quarter.num);
+
+    if table_exists(conn, &table_name)? {
+        Ok(())
+    } else {
+        sql_query(format!(
+            "CREATE TABLE {} PARTITION OF version_downloads \
+             FOR VALUES FROM ('{}') TO ('{}')",
+            table_name, quarter.start, quarter.end
+        ))
+        .execute(conn)?;
+        Ok(())
+    }
+}
+
+fn table_exists(conn: &PgConnection, name: &str) -> QueryResult<bool> {
+    use self::tables::dsl::*;
+    use diesel::dsl::{exists, select};
+
+    select(exists(
+        tables
+            .filter(table_schema.eq("public"))
+            .filter(table_name.eq(name))
+            .filter(table_type.like("BASE TABLE")),
+    ))
+    .get_result(conn)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::super::test_helpers::*;
+    use super::*;
+
+    #[test]
+    fn generate_partition_creates_table_if_it_doesnt_exist() {
+        let conn = conn();
+
+        drop_table_if_exists(&conn, "version_downloads_2018_q1").unwrap();
+        drop_table_if_exists(&conn, "version_downloads_2019_q2").unwrap();
+
+        assert!(!table_exists(&conn, "version_downloads_2018_q1").unwrap());
+        assert!(!table_exists(&conn, "version_downloads_2019_q2").unwrap());
+
+        let q1_2018 = NaiveDate::from_ymd(2018, 2, 14);
+        let q2_2019 = NaiveDate::from_ymd(2019, 4, 1);
+        generate_partition(&conn, q1_2018).unwrap();
+
+        assert!(table_exists(&conn, "version_downloads_2018_q1").unwrap());
+        assert!(!table_exists(&conn, "version_downloads_2019_q2").unwrap());
+
+        generate_partition(&conn, q2_2019).unwrap();
+
+        assert!(table_exists(&conn, "version_downloads_2018_q1").unwrap());
+        assert!(table_exists(&conn, "version_downloads_2019_q2").unwrap());
+    }
+
+    #[test]
+    fn quarter_from_date() {
+        let q1_2018 = NaiveDate::from_ymd(2018, 2, 14);
+        let q4_2018 = NaiveDate::from_ymd(2018, 12, 1);
+        let q1 = Quarter {
+            start: NaiveDate::from_ymd(2018, 1, 1),
+            end: NaiveDate::from_ymd(2018, 4, 1),
+            num: 1,
+        };
+        let q4 = Quarter {
+            start: NaiveDate::from_ymd(2018, 10, 1),
+            end: NaiveDate::from_ymd(2019, 1, 1),
+            num: 4,
+        };
+
+        assert_eq!(Some(q1), Quarter::from_date(q1_2018));
+        assert_eq!(Some(q4), Quarter::from_date(q4_2018));
+    }
+
+    fn drop_table_if_exists(conn: &PgConnection, table_name: &str) -> QueryResult<()> {
+        sql_query(&format!("DROP TABLE IF EXISTS {}", table_name)).execute(conn)?;
+        Ok(())
+    }
+}

src/tasks/test_helpers.rs

@@ -0,0 +1,8 @@
+use crate::env;
+use diesel::prelude::*;
+
+pub(crate) fn conn() -> PgConnection {
+    let conn = PgConnection::establish(&env("TEST_DATABASE_URL")).unwrap();
+    conn.begin_test_transaction().unwrap();
+    conn
+}

src/tasks/update_downloads.rs

@@ -74,19 +74,11 @@ fn collect(conn: &PgConnection, rows: &[VersionDownload]) -> QueryResult<()> {
 
 #[cfg(test)]
 mod test {
+    use super::super::test_helpers::*;
     use super::*;
-    use crate::{
-        env,
-        models::{Crate, NewCrate, NewUser, NewVersion, User, Version},
-    };
+    use crate::models::{Crate, NewCrate, NewUser, NewVersion, User, Version};
     use std::collections::HashMap;
 
-    fn conn() -> PgConnection {
-        let conn = PgConnection::establish(&env("TEST_DATABASE_URL")).unwrap();
-        conn.begin_test_transaction().unwrap();
-        conn
-    }
-
     fn user(conn: &PgConnection) -> User {
         NewUser::new(2, "login", None, None, "access_token")
             .create_or_update(None, conn)