Add cli option to save cell data json to directory

src/main.rs

@@ -6,12 +6,12 @@ use humansize::{file_size_opts, FileSize};
 use models::file::File;
 use models::game_mod::Mod;
 use reqwest::StatusCode;
-use serde::Serialize;
 use sqlx::postgres::PgPoolOptions;
 use std::collections::{HashMap, HashSet};
 use std::env;
 use std::io::Seek;
 use std::io::SeekFrom;
+use std::io::Write;
 use std::process::Command;
 use std::time::Duration;
 use tempfile::tempdir;
@@ -44,9 +44,13 @@ struct Args {
     /// the page number to start scraping for mods on nexus mods.
     page: usize,
 
-    /// output the cell mod edit counts as json
-    #[argh(switch, short = 'e')]
-    dump_edits: bool,
+    /// file to output the cell mod edit counts as json
+    #[argh(option, short = 'e')]
+    dump_edits: Option<String>,
+
+    /// folder to output all cell data as json files
+    #[argh(option, short = 'c')]
+    cell_data: Option<String>,
 }
 
 async fn extract_with_compress_tools(
@@ -209,7 +213,7 @@ pub async fn main() -> Result<()> {
 
     let args: Args = argh::from_env();
 
-    if args.dump_edits {
+    if let Some(dump_edits) = args.dump_edits {
         let mut cell_mod_edit_counts = HashMap::new();
         for x in -77..75 {
             for y in -50..44 {
@@ -218,7 +222,24 @@ pub async fn main() -> Result<()> {
                 }
             }
         }
-        println!("{}", serde_json::to_string(&cell_mod_edit_counts)?);
+        let mut file = std::fs::File::create(dump_edits)?;
+        write!(file, "{}", serde_json::to_string(&cell_mod_edit_counts)?)?;
+        return Ok(());
+    }
+
+    if let Some(cell_data_dir) = args.cell_data {
+        for x in -77..75 {
+            for y in -50..44 {
+                if let Ok(data) = cell::get_cell_data(&pool, "Skyrim.esm", 1, x, y).await {
+                    let path = format!("{}/{}", &cell_data_dir, x);
+                    let path = std::path::Path::new(&path);
+                    std::fs::create_dir_all(&path)?;
+                    let path = path.join(format!("{}.json", y));
+                    let mut file = std::fs::File::create(path)?;
+                    write!(file, "{}", serde_json::to_string(&data)?)?;
+                }
+            }
+        }
         return Ok(());
     }
 
@@ -231,7 +252,6 @@ pub async fn main() -> Result<()> {
         .connect_timeout(CONNECT_TIMEOUT)
         .build()?;
 
-
     while has_next_page {
         let page_span = info_span!("page", page);
         let _page_span = page_span.enter();

src/models/cell.rs

@@ -29,6 +29,18 @@ pub struct UnsavedCell<'a> {
     pub is_persistent: bool,
 }
 
+#[derive(Debug, Serialize, Deserialize, FromRow)]
+pub struct CellData {
+    pub form_id: i32,
+    pub x: Option<i32>,
+    pub y: Option<i32>,
+    pub is_persistent: bool,
+    pub plugins_count: Option<i64>,
+    pub files_count: Option<i64>,
+    pub mods_count: Option<i64>,
+    pub mods: Option<serde_json::Value>,
+}
+
 #[instrument(level = "debug", skip(pool))]
 pub async fn insert(
     pool: &sqlx::Pool<sqlx::Postgres>,
@@ -130,3 +142,40 @@ pub async fn count_mod_edits(
     .await
     .context("Failed to count mod edits on cell")
 }
+
+/// Returns cell properties plus a list of mods that edit the cell
+#[instrument(level = "debug", skip(pool))]
+pub async fn get_cell_data(
+    pool: &sqlx::Pool<sqlx::Postgres>,
+    master: &str,
+    world_id: i32,
+    x: i32,
+    y: i32,
+) -> Result<CellData> {
+    sqlx::query_as!(
+        CellData,
+        r#"SELECT
+                cells.x,
+                cells.y,
+                cells.is_persistent,
+                cells.form_id,
+                COUNT(DISTINCT plugins.id) as plugins_count,
+                COUNT(DISTINCT files.id) as files_count,
+                COUNT(DISTINCT mods.id) as mods_count,
+                json_agg(mods.*) as mods
+            FROM cells
+            JOIN plugin_cells on cells.id = cell_id
+            JOIN plugins ON plugins.id = plugin_id
+            JOIN files ON files.id = file_id
+            JOIN mods ON mods.id = mod_id
+            WHERE master = $1 AND world_id = $2 AND x = $3 and y = $4
+            GROUP BY cells.x, cells.y, cells.is_persistent, cells.form_id"#,
+        master,
+        world_id,
+        x,
+        y
+    )
+    .fetch_one(pool)
+    .await
+    .context("Failed get cell data")
+}