Add timestep option to dump_cell_edit_counts_over_time

src/commands/dump_cell_edit_counts_over_time.rs

@@ -1,15 +1,35 @@
 use crate::models::cell::{self, CellFileEditCount};
 use anyhow::Result;
-use chrono::{Duration, NaiveDateTime};
+use chrono::{Duration, NaiveDateTime, Months};
 use sqlx::postgres::PgPoolOptions;
-use std::{collections::HashMap, env};
+use std::{collections::HashMap, env, str::FromStr};
 use tokio::fs::File;
 use tokio::io::AsyncWriteExt;
 use tracing::{debug, info};
 
+#[derive(Debug)]
+pub enum TimeStep {
+    Day,
+    Week,
+    Month,
+}
+
+impl FromStr for TimeStep {
+    type Err = String;
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "day" => Ok(TimeStep::Day),
+            "week" => Ok(TimeStep::Week),
+            "month" => Ok(TimeStep::Month),
+            _ => Err(format!("invalid time step: {}", s)),
+        }
+    }
+}
+
 pub async fn dump_cell_edit_counts_over_time(
     start_date: NaiveDateTime,
     end_date: NaiveDateTime,
+    time_step: TimeStep,
     path: &str,
 ) -> Result<()> {
     let mut pool = PgPoolOptions::new()
@@ -28,7 +48,11 @@ pub async fn dump_cell_edit_counts_over_time(
                 .connect(&env::var("DATABASE_URL")?)
                 .await?;
         }
-        let next_date = current_date + Duration::weeks(1);
+        let next_date = match &time_step {
+            TimeStep::Day => current_date + Duration::days(1),
+            TimeStep::Week => current_date + Duration::weeks(1),
+            TimeStep::Month => current_date.checked_add_months(Months::new(1)).unwrap(),
+        };
         let mut cell_file_edit_counts = HashMap::new();
         let counts =
             cell::count_file_edits_in_time_range(&pool, "Skyrim.esm", 1, current_date, next_date)

src/commands/mod.rs

@@ -14,7 +14,7 @@ pub mod update;
 pub use download_tiles::download_tiles;
 pub use dump_cell_data::dump_cell_data;
 pub use dump_cell_edit_counts::dump_cell_edit_counts;
-pub use dump_cell_edit_counts_over_time::dump_cell_edit_counts_over_time;
+pub use dump_cell_edit_counts_over_time::{dump_cell_edit_counts_over_time, TimeStep};
 pub use dump_file_data::dump_file_data;
 pub use dump_games::dump_games;
 pub use dump_mod_cell_counts::dump_mod_cell_counts;

src/main.rs

@@ -1,6 +1,6 @@
 use anyhow::Result;
 use argh::FromArgs;
-use chrono::{NaiveDateTime, NaiveDate, Utc};
+use chrono::{NaiveDate, NaiveDateTime, Utc};
 use dotenv::dotenv;
 use sqlx::postgres::PgPoolOptions;
 use std::env;
@@ -16,7 +16,7 @@ use commands::{
     backfills::backfill_is_base_game, backfills::backfill_is_translation,
     backfills::deduplicate_interior_cells, download_tiles, dump_cell_data, dump_cell_edit_counts,
     dump_cell_edit_counts_over_time, dump_file_data, dump_games, dump_mod_cell_counts,
-    dump_mod_data, dump_mod_search_index, dump_plugin_data, update,
+    dump_mod_data, dump_mod_search_index, dump_plugin_data, update, TimeStep,
 };
 
 #[derive(FromArgs)]
@@ -42,10 +42,16 @@ struct Args {
     #[argh(option, short = 'e')]
     dump_edits: Option<String>,
 
-    /// file to output the cell mod edit counts over time as json
+    /// file to output the cell mod edit counts over time as json (time_step option required with 
+    /// this option)
     #[argh(option, short = 'E')]
     dump_edits_over_time: Option<String>,
 
+    /// the span of time to group cell edit counts into (day, week, or month) when dumping cell 
+    /// edits (only relevant for use with dump_edits_over_time option)
+    #[argh(option, short = 'T')]
+    time_step: Option<TimeStep>,
+
     /// folder to output all cell data as json files
     #[argh(option, short = 'c')]
     cell_data: Option<String>,
@@ -112,12 +118,20 @@ pub async fn main() -> Result<()> {
         return dump_cell_edit_counts(&pool, &path).await;
     }
     if let Some(path) = args.dump_edits_over_time {
-        return dump_cell_edit_counts_over_time(
-            NaiveDate::from_ymd_opt(2011, 11, 11).unwrap().and_hms_opt(0, 0, 0).unwrap(),
-            Utc::now().naive_utc(),
-            &path,
-        )
-        .await;
+        if let Some(time_step) = args.time_step {
+            return dump_cell_edit_counts_over_time(
+                NaiveDate::from_ymd_opt(2011, 11, 11)
+                    .unwrap()
+                    .and_hms_opt(0, 0, 0)
+                    .unwrap(),
+                Utc::now().naive_utc(),
+                time_step,
+                &path,
+            )
+            .await;
+        } else {
+            panic!("time_step option required with dump_edits_over_time option");
+        }
     }
     if let Some(dir) = args.cell_data {
         return dump_cell_data(&dir).await;