diff --git a/sdgb-api/src/title/methods/api_ext/mod.rs b/sdgb-api/src/title/methods/api_ext/mod.rs new file mode 100644 index 0000000..95afbe5 --- /dev/null +++ b/sdgb-api/src/title/methods/api_ext/mod.rs @@ -0,0 +1,26 @@ +use crate::title::methods::{APIExt, APIMethod}; + +#[crabtime::function] +fn api_implement(api_names: Vec) { + for api_name in api_names { + crabtime::output!( + pub struct {{api_name}}; + + impl APIExt for {{api_name}} { + const METHOD: APIMethod = APIMethod::{{api_name}}; + type Payload = crate::title::model::{{api_name}}; + type Response = crate::title::model::{{api_name}}Resp; + } + ); + } +} + +api_implement!([ + "Ping", + "UserLoginApi", + "UserLogoutApi", + "GetUserDataApi", + "GetUserPreviewApi", + "GetUserRatingApi", + "GetUserMusicApi", +]); diff --git a/sdgb-api/src/title/methods/has_uid/mod.rs b/sdgb-api/src/title/methods/has_uid/mod.rs new file mode 100644 index 0000000..811c1f3 --- /dev/null +++ b/sdgb-api/src/title/methods/has_uid/mod.rs @@ -0,0 +1,21 @@ +use crate::title::methods::HasUid; + +#[crabtime::function] +fn uid_get_impl(api_names: Vec) { + for api_name in api_names { + crabtime::output!( + impl HasUid for crate::title::model::{{api_name}}Resp { + fn get_uid(&self) -> u32 { + self.user_id + } + } + ); + } +} + +uid_get_impl!([ + "GetUserDataApi", + "GetUserMusicApi", + "GetUserPreviewApi", + "GetUserRatingApi" +]); diff --git a/sdgb-api/src/title/methods/mod.rs b/sdgb-api/src/title/methods/mod.rs index e6a973a..f6d33b9 100644 --- a/sdgb-api/src/title/methods/mod.rs +++ b/sdgb-api/src/title/methods/mod.rs @@ -53,31 +53,13 @@ pub trait APIExt { type Response: for<'de> Deserialize<'de>; } -#[crabtime::function] -fn api_implement(api_names: Vec) { - for api_name in api_names { - let api_name = api_name; - crabtime::output!( - pub struct {{api_name}}; - - impl APIExt for {{api_name}} { - const METHOD: APIMethod = APIMethod::{{api_name}}; - type Payload = crate::title::model::{{api_name}}; - type Response = crate::title::model::{{api_name}}Resp; - } - ); - } +pub trait HasUid { + fn get_uid(&self) -> u32; } -api_implement!([ - "Ping", - "UserLoginApi", - "UserLogoutApi", - "GetUserDataApi", - "GetUserPreviewApi", - "GetUserRatingApi", - "GetUserMusicApi", -]); +mod api_ext; +mod has_uid; +pub use api_ext::*; #[cfg(test)] mod _test { diff --git a/sdgb-api/src/title/model/get_user_preview_api/mod.rs b/sdgb-api/src/title/model/get_user_preview_api/mod.rs index f2de7e0..48e69cc 100644 --- a/sdgb-api/src/title/model/get_user_preview_api/mod.rs +++ b/sdgb-api/src/title/model/get_user_preview_api/mod.rs @@ -9,6 +9,12 @@ pub struct GetUserPreviewApi { pub user_id: u32, } +impl From for GetUserPreviewApi { + fn from(user_id: u32) -> Self { + Self { user_id } + } +} + #[derive(Debug, Clone, Serialize, Deserialize, Encode, Decode)] #[serde(rename_all = "camelCase")] pub struct GetUserPreviewApiResp { diff --git a/sdgb-cli/Cargo.toml b/sdgb-cli/Cargo.toml index d25c6f5..662717d 100644 --- a/sdgb-cli/Cargo.toml +++ b/sdgb-cli/Cargo.toml @@ -6,7 +6,7 @@ authors = ["mokurin000"] description = "CLI tool for SDGB protocol" [features] -default = ["compio"] +default = ["compio", "fetchall"] compio = ["dep:compio", "sdgb-api/compio"] tokio = ["dep:tokio", "sdgb-api/tokio"] diff --git a/sdgb-cli/src/cache/mod.rs b/sdgb-cli/src/cache/mod.rs index b5bdfaa..8813e2a 100644 --- a/sdgb-cli/src/cache/mod.rs +++ b/sdgb-cli/src/cache/mod.rs @@ -7,29 +7,36 @@ static DATABASE: LazyLock = LazyLock::new(|| { .create("players.redb") .expect("failed to open database") }); -const DIFINITION: TableDefinition<'_, u32, Vec> = redb::TableDefinition::new("players"); + +pub const PLAYERS: TableDefinition<'_, u32, Vec> = redb::TableDefinition::new("players"); +pub const PLAYER_B50: TableDefinition<'_, u32, Vec> = redb::TableDefinition::new("b50"); pub fn write_txn() -> Result { Ok(DATABASE.begin_write()?) } -pub fn open_table(write: &WriteTransaction) -> Result>, redb::Error> { - Ok(write.open_table(DIFINITION)?) -} - pub fn read_txn() -> Result { Ok(DATABASE.begin_read()?) } -pub fn open_table_read( +pub fn open_table<'a>( + write: &'a WriteTransaction, + definition: TableDefinition<'_, u32, Vec>, +) -> Result>, redb::Error> { + Ok(write.open_table(definition)?) +} + +pub fn open_table_ro( read: &ReadTransaction, + definition: TableDefinition<'_, u32, Vec>, ) -> Result>, redb::Error> { - Ok(read.open_table(DIFINITION)?) + Ok(read.open_table(definition)?) } pub fn init_db() -> Result<(), redb::Error> { let write_txn = DATABASE.begin_write()?; - write_txn.open_table(DIFINITION)?; + write_txn.open_table(PLAYERS)?; + write_txn.open_table(PLAYER_B50)?; write_txn.commit()?; Ok(()) } diff --git a/sdgb-cli/src/main.rs b/sdgb-cli/src/main.rs index 9ab17e9..aba1b8f 100644 --- a/sdgb-cli/src/main.rs +++ b/sdgb-cli/src/main.rs @@ -232,124 +232,16 @@ async fn main() -> Result<(), Box> { #[cfg(feature = "fetchall")] Commands::ListAllUser { concurrency } => { - use futures_util::StreamExt; - use sdgb_api::bincode::borrow_decode_from_slice; - use std::io::{self, BufRead}; + use sdgb_api::title::methods::GetUserPreviewApi; - let mut user_ids = Vec::new(); - { - let mut stdin = io::stdin().lock(); - let mut buf = String::new(); - - while stdin.read_line(&mut buf).is_ok_and(|size| size != 0) { - if buf.is_empty() { - continue; - } - - let user_id: u32 = buf.trim().parse()?; - buf.clear(); - user_ids.push(user_id); - } - } - - let _ = cache::init_db(); - let read = cache::read_txn()?; - let write = cache::write_txn()?; - let config = sdgb_api::bincode::config::Configuration::< - sdgb_api::bincode::config::LittleEndian, - >::default() - .with_no_limit(); - - info!("number of user_id: {}", user_ids.len()); - - let collect = futures_util::stream::iter(user_ids) - .map(async |user_id| { - { - let cache_table = cache::open_table_read(&read)?; - let data = cache_table.get(user_id)?; - if let Some(data) = data { - let decoded: (GetUserPreviewApiResp, _) = - borrow_decode_from_slice(&data.value(), config)?; - - return Ok(decoded.0); - } - } - - if EARLY_QUIT.load(Ordering::Relaxed) { - return Err("early skip due to ctrl-c")?; - } - - let resp = Sdgb1_50::request::<_, GetUserPreviewApiResp>( - &client, - APIMethod::GetUserPreviewApi, - user_id, - GetUserPreviewApi { user_id }, - ) - .await; - - match &resp { - Ok(resp) => { - use sdgb_api::bincode::encode_to_vec; - - info!("found: {user_id}"); - - if let Ok(mut table) = cache::open_table(&write) - && let Ok(encoded) = encode_to_vec(resp, config) - { - _ = table.insert(resp.user_id, encoded); - } - } - Err(sdgb_api::ApiError::JSON { .. }) => {} - Err(e) => { - error!("preview failed: {e}"); - } - } - - Result::<_, Box>::Ok(resp?) - }) - .buffer_unordered(concurrency) // slower to avoid being banned - .filter_map(async |r| r.ok()) - .collect::>() - .await; - drop(collect); - - let _ = write.commit(); + use crate::{cache::PLAYERS, utils::helpers::cached_concurrent_fetch}; + cached_concurrent_fetch::(&client, concurrency, PLAYERS).await?; } #[cfg(feature = "fetchall")] - Commands::ListAllUserDump { .. } => { - use std::{fs::OpenOptions, io::BufWriter}; + Commands::ListAllUserDump {} => { + use crate::{cache::PLAYERS, utils::helpers::dump_cache}; - use redb::ReadableTable; - use sdgb_api::bincode::{self, borrow_decode_from_slice}; - - use crate::cache::{open_table_read, read_txn}; - - let file = OpenOptions::new() - .create(true) - .truncate(true) - .write(true) - .open("players.json")?; - - #[cfg(file_lock_ready)] - file.try_lock()?; - - let txn = read_txn()?; - let table = open_table_read(&txn)?; - - let config = bincode::config::Configuration::::default() - .with_no_limit(); - - let user_ids = table - .iter()? - .flatten() - .map(|d| borrow_decode_from_slice(&d.1.value(), config)) - .flatten() - .map(|(value, _)| value) - .collect::>(); - - let writer = BufWriter::new(file); - serde_json::to_writer(writer, &user_ids)?; - info!("dumped {} user id", user_ids.len()); + dump_cache::("players.json", PLAYERS)?; } Commands::Userdata { user_id } => { diff --git a/sdgb-cli/src/utils/helpers/mod.rs b/sdgb-cli/src/utils/helpers/mod.rs new file mode 100644 index 0000000..ceda18c --- /dev/null +++ b/sdgb-cli/src/utils/helpers/mod.rs @@ -0,0 +1,150 @@ +use std::{fs::OpenOptions, io::BufWriter}; +use std::{ + io::{self, BufRead}, + path::Path, + sync::atomic::Ordering, +}; + +use futures_util::StreamExt; +use nyquest_preset::nyquest::AsyncClient; + +use redb::ReadableTable; +use redb::TableDefinition; +use serde::Serialize; +use spdlog::{error, info}; + +use sdgb_api::bincode; +use sdgb_api::title::MaiVersionExt as _; +use sdgb_api::title::{ + Sdgb1_50, + methods::{APIExt, HasUid}, +}; + +use bincode::{BorrowDecode, Encode, borrow_decode_from_slice}; + +use crate::{EARLY_QUIT, cache}; + +pub fn dump_cache( + output_path: impl AsRef, + definition: TableDefinition<'_, u32, Vec>, +) -> Result<(), Box> +where + D: for<'d> BorrowDecode<'d, ()> + Serialize, +{ + let file = OpenOptions::new() + .create(true) + .truncate(true) + .write(true) + .open(output_path)?; + + #[cfg(file_lock_ready)] + file.try_lock()?; + + let txn = cache::read_txn()?; + let table = cache::open_table_ro(&txn, definition)?; + + let config = + bincode::config::Configuration::::default().with_no_limit(); + + let user_ids = table + .iter()? + .flatten() + .map(|d| borrow_decode_from_slice(&d.1.value(), config)) + .flatten() + .map(|(value, _)| value) + .collect::>(); + + let writer = BufWriter::new(file); + serde_json::to_writer(writer, &user_ids)?; + info!("dumped {} user id", user_ids.len()); + + Ok(()) +} + +pub async fn cached_concurrent_fetch( + client: &AsyncClient, + concurrency: usize, + definition: TableDefinition<'_, u32, Vec>, +) -> Result<(), Box> +where + A::Payload: From, + A::Response: Encode + for<'a> BorrowDecode<'a, ()> + HasUid, +{ + let mut user_ids = Vec::new(); + { + let mut stdin = io::stdin().lock(); + let mut buf = String::new(); + + while stdin.read_line(&mut buf).is_ok_and(|size| size != 0) { + if buf.is_empty() { + continue; + } + + let user_id: u32 = buf.trim().parse()?; + buf.clear(); + user_ids.push(user_id); + } + } + + let _ = cache::init_db(); + let read = cache::read_txn()?; + let write = cache::write_txn()?; + let config = sdgb_api::bincode::config::Configuration::< + sdgb_api::bincode::config::LittleEndian, + >::default() + .with_no_limit(); + + info!("number of user_id: {}", user_ids.len()); + + let collect = futures_util::stream::iter(user_ids) + .map(async |user_id| { + { + let cache_table = cache::open_table_ro(&read, definition)?; + let data = cache_table.get(user_id)?; + if let Some(data) = data { + let decoded: (A::Response, _) = + borrow_decode_from_slice(&data.value(), config)?; + + return Ok(decoded.0); + } + } + + if EARLY_QUIT.load(Ordering::Relaxed) { + return Err("early skip due to ctrl-c")?; + } + + let resp = + Sdgb1_50::request_ext::(&client, ::Payload::from(user_id), user_id) + .await; + + match &resp { + Ok(resp) => { + use sdgb_api::bincode::encode_to_vec; + + use crate::cache::PLAYERS; + + info!("found: {user_id}"); + + if let Ok(mut table) = cache::open_table(&write, PLAYERS) + && let Ok(encoded) = encode_to_vec(resp, config) + { + _ = table.insert(resp.get_uid(), encoded); + } + } + Err(sdgb_api::ApiError::JSON { .. }) => {} + Err(e) => { + error!("preview failed: {e}"); + } + } + + Result::<_, Box>::Ok(resp?) + }) + .buffer_unordered(concurrency) // slower to avoid being banned + .filter_map(async |r| r.ok()) + .collect::>() + .await; + drop(collect); + + let _ = write.commit(); + Ok(()) +} diff --git a/sdgb-cli/src/utils/mod.rs b/sdgb-cli/src/utils/mod.rs index f132aa4..67d5fd2 100644 --- a/sdgb-cli/src/utils/mod.rs +++ b/sdgb-cli/src/utils/mod.rs @@ -64,3 +64,6 @@ pub fn human_readable_display( Ok(()) } + +#[cfg(feature = "fetchall")] +pub mod helpers;