From a78fa4772b49a78a8e6f66e20eccd80b7f3f48ea Mon Sep 17 00:00:00 2001 From: Tyler Hallada Date: Mon, 17 Jan 2022 19:25:07 -0500 Subject: [PATCH] Add cli option to save cell data json to directory --- .gitignore | 3 ++- Cargo.lock | 2 ++ Cargo.toml | 2 +- src/main.rs | 34 +++++++++++++++++++++++++------- src/models/cell.rs | 49 ++++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 81 insertions(+), 9 deletions(-) diff --git a/.gitignore b/.gitignore index a192ad8..b5ed0f2 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ /target .env plugins.zip -plugins \ No newline at end of file +plugins +cells \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock index 4ad6b58..e157d49 100644 --- a/Cargo.lock +++ b/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", diff --git a/Cargo.toml b/Cargo.toml index 7701bb9..6eac460 100644 --- a/Cargo.toml +++ b/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"] } diff --git a/src/main.rs b/src/main.rs index f95c1ef..37c4512 100644 --- a/src/main.rs +++ b/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, + + /// folder to output all cell data as json files + #[argh(option, short = 'c')] + cell_data: Option, } 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(); diff --git a/src/models/cell.rs b/src/models/cell.rs index 468d6cb..b5730d0 100644 --- a/src/models/cell.rs +++ b/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, + pub y: Option, + pub is_persistent: bool, + pub plugins_count: Option, + pub files_count: Option, + pub mods_count: Option, + pub mods: Option, +} + #[instrument(level = "debug", skip(pool))] pub async fn insert( pool: &sqlx::Pool, @@ -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, + master: &str, + world_id: i32, + x: i32, + y: i32, +) -> Result { + 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") +}