87 lines
2.8 KiB
Rust
87 lines
2.8 KiB
Rust
use crate::models::cell::{self, CellFileEditCount};
|
|
use anyhow::Result;
|
|
use chrono::{Duration, NaiveDateTime, Months};
|
|
use sqlx::postgres::PgPoolOptions;
|
|
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()
|
|
.max_connections(5)
|
|
.connect(&env::var("DATABASE_URL")?)
|
|
.await?;
|
|
let mut i = 0;
|
|
let mut current_date = start_date;
|
|
while current_date <= end_date {
|
|
if i % 5 == 0 {
|
|
// There's a weird issue that slows down this query after 5 iterations. Recreating the
|
|
// connection pool seems to fix it. I don't know why.
|
|
info!("reconnecting to database");
|
|
pool = PgPoolOptions::new()
|
|
.max_connections(5)
|
|
.connect(&env::var("DATABASE_URL")?)
|
|
.await?;
|
|
}
|
|
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)
|
|
.await?;
|
|
for x in -77..75 {
|
|
for y in -50..44 {
|
|
let count: Option<&CellFileEditCount> = counts.iter().find(|c| c.x.unwrap() == x && c.y.unwrap() == y);
|
|
let count = count.map(|c| c.count).unwrap_or(Some(0)).unwrap();
|
|
debug!(x = x, y = y, count = count, "read cell edit count");
|
|
cell_file_edit_counts.insert(format!("{},{}", x, y), count);
|
|
}
|
|
}
|
|
|
|
let file_name = format!(
|
|
"{}/cell_edits_{}.json",
|
|
path,
|
|
current_date.format("%Y-%m-%d")
|
|
);
|
|
info!(
|
|
"writing {} cell edit counts to {}",
|
|
cell_file_edit_counts.values().sum::<i64>(),
|
|
file_name
|
|
);
|
|
let mut file = File::create(&file_name).await?;
|
|
file.write_all(serde_json::to_string(&cell_file_edit_counts)?.as_bytes()).await?;
|
|
|
|
current_date = next_date;
|
|
i += 1;
|
|
}
|
|
Ok(())
|
|
}
|