refactor: implement fetchall with generic type

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"]

sdgb-cli/src/cache/mod.rs

@@ -7,29 +7,36 @@ static DATABASE: LazyLock<redb::Database> = LazyLock::new(|| {
         .create("players.redb")
         .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> {
     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> {
     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,
+    definition: TableDefinition<'_, u32, Vec<u8>>,
 ) -> 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> {
     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(())
 }

sdgb-cli/src/main.rs

@@ -232,124 +232,16 @@ async fn main() -> Result<(), Box<dyn snafu::Error>> {
 
         #[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<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();
+            use crate::{cache::PLAYERS, utils::helpers::cached_concurrent_fetch};
+            cached_concurrent_fetch::<GetUserPreviewApi>(&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::<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());
+            dump_cache::<GetUserPreviewApiResp>("players.json", PLAYERS)?;
         }
 
         Commands::Userdata { user_id } => {

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<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(())
+}

sdgb-cli/src/utils/mod.rs

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