refactor: implement fetchall with generic type

This commit is contained in:
mokurin000
2025-08-02 19:39:39 +08:00
parent de330005b3
commit 9b53cb633c
9 changed files with 233 additions and 146 deletions

View File

@@ -0,0 +1,26 @@
use crate::title::methods::{APIExt, APIMethod};
#[crabtime::function]
fn api_implement(api_names: Vec<String>) {
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",
]);

View File

@@ -0,0 +1,21 @@
use crate::title::methods::HasUid;
#[crabtime::function]
fn uid_get_impl(api_names: Vec<String>) {
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"
]);

View File

@@ -53,31 +53,13 @@ pub trait APIExt {
type Response: for<'de> Deserialize<'de>; type Response: for<'de> Deserialize<'de>;
} }
#[crabtime::function] pub trait HasUid {
fn api_implement(api_names: Vec<String>) { fn get_uid(&self) -> u32;
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;
}
);
}
} }
api_implement!([ mod api_ext;
"Ping", mod has_uid;
"UserLoginApi", pub use api_ext::*;
"UserLogoutApi",
"GetUserDataApi",
"GetUserPreviewApi",
"GetUserRatingApi",
"GetUserMusicApi",
]);
#[cfg(test)] #[cfg(test)]
mod _test { mod _test {

View File

@@ -9,6 +9,12 @@ pub struct GetUserPreviewApi {
pub user_id: u32, pub user_id: u32,
} }
impl From<u32> for GetUserPreviewApi {
fn from(user_id: u32) -> Self {
Self { user_id }
}
}
#[derive(Debug, Clone, Serialize, Deserialize, Encode, Decode)] #[derive(Debug, Clone, Serialize, Deserialize, Encode, Decode)]
#[serde(rename_all = "camelCase")] #[serde(rename_all = "camelCase")]
pub struct GetUserPreviewApiResp { pub struct GetUserPreviewApiResp {

View File

@@ -6,7 +6,7 @@ authors = ["mokurin000"]
description = "CLI tool for SDGB protocol" description = "CLI tool for SDGB protocol"
[features] [features]
default = ["compio"] default = ["compio", "fetchall"]
compio = ["dep:compio", "sdgb-api/compio"] compio = ["dep:compio", "sdgb-api/compio"]
tokio = ["dep:tokio", "sdgb-api/tokio"] tokio = ["dep:tokio", "sdgb-api/tokio"]

View File

@@ -7,29 +7,36 @@ static DATABASE: LazyLock<redb::Database> = LazyLock::new(|| {
.create("players.redb") .create("players.redb")
.expect("failed to open database") .expect("failed to open database")
}); });
const DIFINITION: TableDefinition<'_, u32, Vec<u8>> = redb::TableDefinition::new("players");
pub const PLAYERS: TableDefinition<'_, u32, Vec<u8>> = redb::TableDefinition::new("players");
pub const PLAYER_B50: TableDefinition<'_, u32, Vec<u8>> = redb::TableDefinition::new("b50");
pub fn write_txn() -> Result<WriteTransaction, redb::Error> { pub fn write_txn() -> Result<WriteTransaction, redb::Error> {
Ok(DATABASE.begin_write()?) Ok(DATABASE.begin_write()?)
} }
pub fn open_table(write: &WriteTransaction) -> Result<Table<'_, u32, Vec<u8>>, redb::Error> {
Ok(write.open_table(DIFINITION)?)
}
pub fn read_txn() -> Result<ReadTransaction, redb::Error> { pub fn read_txn() -> Result<ReadTransaction, redb::Error> {
Ok(DATABASE.begin_read()?) Ok(DATABASE.begin_read()?)
} }
pub fn open_table_read( pub fn open_table<'a>(
write: &'a WriteTransaction,
definition: TableDefinition<'_, u32, Vec<u8>>,
) -> Result<Table<'a, u32, Vec<u8>>, redb::Error> {
Ok(write.open_table(definition)?)
}
pub fn open_table_ro(
read: &ReadTransaction, read: &ReadTransaction,
definition: TableDefinition<'_, u32, Vec<u8>>,
) -> Result<redb::ReadOnlyTable<u32, Vec<u8>>, redb::Error> { ) -> Result<redb::ReadOnlyTable<u32, Vec<u8>>, redb::Error> {
Ok(read.open_table(DIFINITION)?) Ok(read.open_table(definition)?)
} }
pub fn init_db() -> Result<(), redb::Error> { pub fn init_db() -> Result<(), redb::Error> {
let write_txn = DATABASE.begin_write()?; 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()?; write_txn.commit()?;
Ok(()) Ok(())
} }

View File

@@ -232,124 +232,16 @@ async fn main() -> Result<(), Box<dyn snafu::Error>> {
#[cfg(feature = "fetchall")] #[cfg(feature = "fetchall")]
Commands::ListAllUser { concurrency } => { Commands::ListAllUser { concurrency } => {
use futures_util::StreamExt; use sdgb_api::title::methods::GetUserPreviewApi;
use sdgb_api::bincode::borrow_decode_from_slice;
use std::io::{self, BufRead};
let mut user_ids = Vec::new(); use crate::{cache::PLAYERS, utils::helpers::cached_concurrent_fetch};
{ cached_concurrent_fetch::<GetUserPreviewApi>(&client, concurrency, PLAYERS).await?;
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<dyn snafu::Error>>::Ok(resp?)
})
.buffer_unordered(concurrency) // slower to avoid being banned
.filter_map(async |r| r.ok())
.collect::<Vec<_>>()
.await;
drop(collect);
let _ = write.commit();
} }
#[cfg(feature = "fetchall")] #[cfg(feature = "fetchall")]
Commands::ListAllUserDump { .. } => { Commands::ListAllUserDump {} => {
use std::{fs::OpenOptions, io::BufWriter}; use crate::{cache::PLAYERS, utils::helpers::dump_cache};
use redb::ReadableTable; dump_cache::<GetUserPreviewApiResp>("players.json", PLAYERS)?;
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::<bincode::config::LittleEndian>::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::<Vec<GetUserPreviewApiResp>>();
let writer = BufWriter::new(file);
serde_json::to_writer(writer, &user_ids)?;
info!("dumped {} user id", user_ids.len());
} }
Commands::Userdata { user_id } => { Commands::Userdata { user_id } => {

View File

@@ -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<D>(
output_path: impl AsRef<Path>,
definition: TableDefinition<'_, u32, Vec<u8>>,
) -> Result<(), Box<dyn snafu::Error>>
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::<bincode::config::LittleEndian>::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::<Vec<D>>();
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<A: APIExt>(
client: &AsyncClient,
concurrency: usize,
definition: TableDefinition<'_, u32, Vec<u8>>,
) -> Result<(), Box<dyn snafu::Error>>
where
A::Payload: From<u32>,
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::<A>(&client, <A as APIExt>::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<dyn snafu::Error>>::Ok(resp?)
})
.buffer_unordered(concurrency) // slower to avoid being banned
.filter_map(async |r| r.ok())
.collect::<Vec<_>>()
.await;
drop(collect);
let _ = write.commit();
Ok(())
}

View File

@@ -64,3 +64,6 @@ pub fn human_readable_display(
Ok(()) Ok(())
} }
#[cfg(feature = "fetchall")]
pub mod helpers;