refactor: implement fetchall with generic type
This commit is contained in:
26
sdgb-api/src/title/methods/api_ext/mod.rs
Normal file
26
sdgb-api/src/title/methods/api_ext/mod.rs
Normal 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",
|
||||||
|
]);
|
||||||
21
sdgb-api/src/title/methods/has_uid/mod.rs
Normal file
21
sdgb-api/src/title/methods/has_uid/mod.rs
Normal 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"
|
||||||
|
]);
|
||||||
@@ -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 {
|
||||||
|
|||||||
@@ -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 {
|
||||||
|
|||||||
@@ -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"]
|
||||||
|
|||||||
23
sdgb-cli/src/cache/mod.rs
vendored
23
sdgb-cli/src/cache/mod.rs
vendored
@@ -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(())
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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 } => {
|
||||||
|
|||||||
150
sdgb-cli/src/utils/helpers/mod.rs
Normal file
150
sdgb-cli/src/utils/helpers/mod.rs
Normal 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(())
|
||||||
|
}
|
||||||
@@ -64,3 +64,6 @@ pub fn human_readable_display(
|
|||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(feature = "fetchall")]
|
||||||
|
pub mod helpers;
|
||||||
|
|||||||
Reference in New Issue
Block a user