Add cli option to save cell data json to directory

.gitignore

@@ -2,3 +2,4 @@
 .env
 plugins.zip
 plugins
+cells

Cargo.lock

@@ -1731,6 +1731,7 @@ version = "1.0.64"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "799e97dc9fdae36a5c8b8f2cae9ce2ee9fdce2058c57a93e6099d919fd982f79"
 dependencies = [
+ "indexmap",
  "itoa",
  "ryu",
  "serde",
@@ -1928,6 +1929,7 @@ dependencies = [
  "once_cell",
  "proc-macro2",
  "quote",
+ "serde_json",
  "sha2",
  "sqlx-core",
  "sqlx-rt",

Cargo.toml

@@ -21,7 +21,7 @@ scraper = "0.12"
 seahash = "4.1"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
-sqlx = { version = "0.5", features = ["runtime-tokio-native-tls", "postgres", "migrate", "chrono"] }
+sqlx = { version = "0.5", features = ["runtime-tokio-native-tls", "postgres", "migrate", "chrono", "json"] }
 skyrim-cell-dump = "0.3.1"
 tempfile = "3.2"
 tokio = { version = "1.5.0", features = ["full"] }

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")
+}