diff --git a/sr_download/src/db.rs b/sr_download/src/db_part.rs similarity index 95% rename from sr_download/src/db.rs rename to sr_download/src/db_part.rs index 4e0d826..db6c8e1 100644 --- a/sr_download/src/db.rs +++ b/sr_download/src/db_part.rs @@ -1,7 +1,7 @@ use blake3::Hasher; use sea_orm::{ - ActiveModelTrait, ConnectOptions, Database, DatabaseConnection, EntityTrait, - IntoActiveModel, ModelTrait, QueryOrder, QuerySelect, TransactionTrait, + ActiveModelTrait, ConnectOptions, Database, DatabaseConnection, EntityTrait, IntoActiveModel, + ModelTrait, QueryOrder, QuerySelect, TransactionTrait, }; use tracing::{event, Level}; @@ -59,19 +59,21 @@ pub enum CoverStrategy { /// 如果失败, 会返回 Err /// 如果成功, 会返回 Ok(true) /// 如果数据已经存在, 会根据策略返回 -pub async fn save_data_to_db( +pub async fn save_data_to_db( save_id: SaveId, - save_type: SaveType, + save_type: T, data: D, cover_strategy: Option, db: &DatabaseConnection, ) -> anyhow::Result where D: Into, + T: Into, { // 干活之前, 先检查一下数据是否已经存在 // 如果已经存在, 那就根据策略来处理 let cover_strategy = cover_strategy.unwrap_or_default(); + let save_type: SaveType = save_type.into(); let exitst_data: Option = { model::main_data::Entity::find_by_id(save_id as i32) .select_only() diff --git a/sr_download/src/main.rs b/sr_download/src/main.rs index 1206235..e955d98 100644 --- a/sr_download/src/main.rs +++ b/sr_download/src/main.rs @@ -1,32 +1,32 @@ +use futures::future::select_all; use std::{ops::Range, path::Path}; use tracing::{event, Level}; -use futures::{future::select_all, task}; mod config; -mod db; +mod db_part; mod model; mod net; pub type SaveId = u32; pub const TEXT_DATA_MAX_LEN: usize = 1024; -use net::DownloadFile; +use model::sea_orm_active_enums::SaveType; async fn big_worker(db: sea_orm::DatabaseConnection, work_range: Range) { let client = net::Downloader::default(); for work_id in work_range { - match client.try_download_as_any(work_id).await { + match match client.try_download_as_any(work_id).await { Some(file) => { - match file { - DownloadFile::Save(data) => { - - }, - DownloadFile::Ship(data) => { - - }, - } + let save_type = (&file).into(); + db_part::save_data_to_db(work_id, save_type, file.take_data(), None, &db) } - }; + None => db_part::save_data_to_db(work_id, SaveType::None, "".to_string(), None, &db), + } + .await + { + Ok(_) => event!(Level::INFO, "Save data {} success", work_id), + Err(e) => event!(Level::WARN, "Save data {} failed: {:?}", work_id, e), + } } } @@ -43,14 +43,14 @@ async fn main() -> anyhow::Result<()> { } }; - let db_connect = db::connect(&conf).await?; - let start_id = db::find_max_id(&db_connect).await; + let db_connect = db_part::connect(&conf).await?; + let start_id = db_part::find_max_id(&db_connect).await; event!(Level::INFO, "Starting download from save_id: {}", start_id); // 1321469 end let end_id: SaveId = 1321469; - + let mut current_id = start_id; let batch_size = 100; @@ -59,19 +59,30 @@ async fn main() -> anyhow::Result<()> { let max_works = 10; for _ in 0..10 { let end = current_id + batch_size; - works.push(tokio::spawn(big_worker(db_connect.clone(), current_id..end))); + works.push(tokio::spawn(big_worker( + db_connect.clone(), + current_id..end, + ))); current_id = end; } while current_id < end_id || !works.is_empty() { while current_id < end_id && works.len() < max_works { let end = current_id + batch_size; - works.push(tokio::spawn(big_worker(db_connect.clone(), current_id..end))); + works.push(tokio::spawn(big_worker( + db_connect.clone(), + current_id..end, + ))); current_id = end; } if !works.is_empty() { let (result, index, remain) = select_all(works).await; - event!(Level::INFO, "worker {} finish with result: {:?}", index, result); + event!( + Level::INFO, + "worker {} finish with result: {:?}", + index, + result + ); works = remain; } } diff --git a/sr_download/src/net.rs b/sr_download/src/net.rs index bf53843..4c7b5ac 100644 --- a/sr_download/src/net.rs +++ b/sr_download/src/net.rs @@ -1,7 +1,7 @@ use reqwest::Client; use std::time::Duration; -use crate::SaveId; +use crate::{model::sea_orm_active_enums::SaveType, SaveId}; pub struct Downloader { pub client: Client, @@ -23,20 +23,32 @@ impl DownloadFile { _ => None, } } - pub fn as_save(&self) -> Option<&str> { match self { DownloadFile::Save(s) => Some(s), _ => None, } } - pub fn is_ship(&self) -> bool { - self.as_ship().is_some() + matches!(self, DownloadFile::Ship(_)) } - pub fn is_save(&self) -> bool { - self.as_save().is_some() + matches!(self, DownloadFile::Save(_)) + } + pub fn take_data(self) -> String { + match self { + DownloadFile::Ship(s) => s, + DownloadFile::Save(s) => s, + } + } +} + +impl From<&DownloadFile> for SaveType { + fn from(file: &DownloadFile) -> Self { + match file { + DownloadFile::Ship(_) => SaveType::Ship, + DownloadFile::Save(_) => SaveType::Save, + } } }