feat: 实现全部的下载逻辑

This commit is contained in:
amtoaer
2024-03-30 14:48:43 +08:00
parent 1b0d657615
commit 9eadb70ec7
+78 -37
View File
@@ -7,33 +7,31 @@ use futures::stream::FuturesUnordered;
use futures::Future; use futures::Future;
use futures_util::{pin_mut, StreamExt}; use futures_util::{pin_mut, StreamExt};
use log::{error, info}; use log::{error, info};
use once_cell::sync::Lazy;
use sea_orm::entity::prelude::*; use sea_orm::entity::prelude::*;
use sea_orm::ActiveValue::Set; use sea_orm::ActiveValue::Set;
use sea_orm::TryIntoModel; use sea_orm::TryIntoModel;
use serde::Serialize; use serde_json::json;
use tinytemplate::TinyTemplate;
use tokio::fs; use tokio::fs;
use tokio::sync::Semaphore; use tokio::sync::Semaphore;
use super::status::Status; use super::status::PageStatus;
use super::utils::{unhandled_videos_pages, ModelWrapper, NFOMode, NFOSerializer}; use super::utils::{unhandled_videos_pages, ModelWrapper, NFOMode, NFOSerializer};
use crate::bilibili::{BestStream, BiliClient, FavoriteList, FilterOption, PageInfo, Video}; use crate::bilibili::{BestStream, BiliClient, FavoriteList, FilterOption, PageInfo, Video};
use crate::core::utils::{create_video_pages, create_videos, exist_labels, filter_videos, handle_favorite_info}; use crate::core::utils::{create_video_pages, create_videos, exist_labels, filter_videos, handle_favorite_info};
use crate::downloader::Downloader; use crate::downloader::Downloader;
use crate::Result; use crate::Result;
/// 用来拼接路径名称 static TEMPLATE: Lazy<handlebars::Handlebars> = Lazy::new(|| {
#[derive(Serialize)] let mut handlebars = handlebars::Handlebars::new();
struct Context<'a> { handlebars.register_template_string("video", "{{bvid}}").unwrap();
bvid: &'a str, handlebars.register_template_string("page", "{{bvid}}").unwrap();
name: &'a str, handlebars
pid: &'a str, });
}
pub async fn process_favorite(bili_client: &BiliClient, fid: &str, connection: &DatabaseConnection) -> Result<()> { pub async fn process_favorite(bili_client: &BiliClient, fid: &str, connection: &DatabaseConnection) -> Result<()> {
let favorite_model = refresh_favorite(bili_client, fid, connection).await?; let favorite_model = refresh_favorite(bili_client, fid, connection).await?;
download_favorite(bili_client, favorite_model, connection).await?; download_favorite(bili_client, favorite_model, connection).await
Ok(())
} }
pub async fn refresh_favorite( pub async fn refresh_favorite(
@@ -71,6 +69,7 @@ pub async fn refresh_favorite(
break; break;
} }
} }
info!("handle videos done");
Ok(favorite_model) Ok(favorite_model)
} }
@@ -79,6 +78,7 @@ pub async fn download_favorite(
favorite_model: favorite::Model, favorite_model: favorite::Model,
connection: &DatabaseConnection, connection: &DatabaseConnection,
) -> Result<()> { ) -> Result<()> {
info!("start to download!");
let unhandled_videos_pages = unhandled_videos_pages(&favorite_model, connection).await?; let unhandled_videos_pages = unhandled_videos_pages(&favorite_model, connection).await?;
let semaphore = Arc::new(Semaphore::new(3)); let semaphore = Arc::new(Semaphore::new(3));
let downloader = Downloader::default(); let downloader = Downloader::default();
@@ -98,6 +98,7 @@ pub async fn download_favorite(
error!("Error: {e}"); error!("Error: {e}");
} }
} }
info!("download done.");
Ok(()) Ok(())
} }
@@ -113,8 +114,6 @@ pub async fn download_video_pages(
if let Err(e) = permit { if let Err(e) = permit {
return Err(e.into()); return Err(e.into());
} }
let mut template = TinyTemplate::new();
let _ = template.add_template("video", "{bvid}");
let child_semaphore = Arc::new(Semaphore::new(5)); let child_semaphore = Arc::new(Semaphore::new(5));
let mut tasks = FuturesUnordered::new(); let mut tasks = FuturesUnordered::new();
for page_model in pages { for page_model in pages {
@@ -132,6 +131,14 @@ pub async fn download_video_pages(
error!("Error: {e}"); error!("Error: {e}");
} }
} }
let is_single_page = video_model.single_page.unwrap();
// 对于单页视频,page 的下载已经足够
// 对于多页视频,page 下载仅包含了分集内容,需要额外补上视频的 poster 的 tvshow.nfo
if !is_single_page {
let base_path = Path::new(&video_model.path);
generate_video_nfo(true, &video_model, base_path.join("tvshow.nfo")).await?;
download_video_poster(true, &video_model, downloader, base_path.join("poster.jpg")).await?;
}
Ok(()) Ok(())
} }
@@ -147,20 +154,18 @@ pub async fn download_page(
if let Err(e) = permit { if let Err(e) = permit {
return Err(e.into()); return Err(e.into());
} }
let mut status = Status::new(page_model.download_status); let mut status = PageStatus::new(page_model.download_status);
let seprate_status = status.should_run(); let seprate_status = status.should_run();
let is_single_page = video_model.single_page.unwrap(); let is_single_page = video_model.single_page.unwrap();
let base_path = Path::new(&video_model.path); let base_path = Path::new(&video_model.path);
let mut template = TinyTemplate::new();
// 这个文件名模板支持自定义 // 这个文件名模板支持自定义
let _ = template.add_template("video", "{bvid}"); let base_name = TEMPLATE.render(
let base_name = template.render( "page",
"video", &json!({
&Context { "bvid": &video_model.bvid,
bvid: &video_model.bvid, "name": &video_model.name,
name: &video_model.name, "pid": page_model.pid,
pid: &page_model.pid.to_string(), }),
},
)?; )?;
let (poster_path, video_path, nfo_path) = if is_single_page { let (poster_path, video_path, nfo_path) = if is_single_page {
( (
@@ -172,18 +177,18 @@ pub async fn download_page(
( (
base_path base_path
.join("Season 1") .join("Season 1")
.join(format!("{} - S01E{:2}-thumb.jpg", &base_name, page_model.pid)), .join(format!("{} - S01E{:0>2}-thumb.jpg", &base_name, page_model.pid)),
base_path base_path
.join("Season 1") .join("Season 1")
.join(format!("{} - S01E{:2}.mp4", &base_name, page_model.pid)), .join(format!("{} - S01E{:0>2}.mp4", &base_name, page_model.pid)),
base_path base_path
.join("Season 1") .join("Season 1")
.join(format!("{} - S01E{:2}.nfo", &base_name, page_model.pid)), .join(format!("{} - S01E{:0>2}.nfo", &base_name, page_model.pid)),
) )
}; };
let tasks: Vec<Pin<Box<dyn Future<Output = Result<()>>>>> = vec![ let tasks: Vec<Pin<Box<dyn Future<Output = Result<()>>>>> = vec![
// 暂时不支持下载字幕 // 暂时不支持下载字幕
Box::pin(download_poster( Box::pin(download_page_poster(
seprate_status[0], seprate_status[0],
video_model, video_model,
&page_model, &page_model,
@@ -198,7 +203,7 @@ pub async fn download_page(
downloader, downloader,
video_path, video_path,
)), )),
Box::pin(generate_nfo(seprate_status[2], video_model, &page_model, nfo_path)), Box::pin(generate_page_nfo(seprate_status[2], video_model, &page_model, nfo_path)),
]; ];
let results = futures::future::join_all(tasks).await; let results = futures::future::join_all(tasks).await;
status.update_status(&results); status.update_status(&results);
@@ -208,7 +213,7 @@ pub async fn download_page(
Ok(page_active_model.try_into_model().unwrap()) Ok(page_active_model.try_into_model().unwrap())
} }
pub async fn download_poster( pub async fn download_page_poster(
should_run: bool, should_run: bool,
video_model: &video::Model, video_model: &video::Model,
page_model: &page::Model, page_model: &page::Model,
@@ -223,8 +228,7 @@ pub async fn download_poster(
Some(url) => url.as_str(), Some(url) => url.as_str(),
None => video_model.cover.as_str(), None => video_model.cover.as_str(),
}; };
downloader.fetch(url, &poster_path).await?; downloader.fetch(url, &poster_path).await
Ok(())
} }
pub async fn download_video( pub async fn download_video(
@@ -272,7 +276,15 @@ pub async fn download_video(
Ok(()) Ok(())
} }
pub async fn generate_nfo( async fn generate_nfo(serializer: NFOSerializer<'_>, nfo_path: PathBuf) -> Result<()> {
if let Some(parent) = nfo_path.parent() {
fs::create_dir_all(parent).await?;
}
fs::write(nfo_path, serializer.generate_nfo().await?.as_bytes()).await?;
Ok(())
}
pub async fn generate_page_nfo(
should_run: bool, should_run: bool,
video_model: &video::Model, video_model: &video::Model,
page_model: &page::Model, page_model: &page::Model,
@@ -287,9 +299,38 @@ pub async fn generate_nfo(
} else { } else {
NFOSerializer(ModelWrapper::Page(page_model), NFOMode::EPOSODE) NFOSerializer(ModelWrapper::Page(page_model), NFOMode::EPOSODE)
}; };
if let Some(parent) = nfo_path.parent() { generate_nfo(nfo_serializer, nfo_path).await
fs::create_dir_all(parent).await?; }
pub async fn download_video_poster(
should_run: bool,
video_model: &video::Model,
downloader: &Downloader,
poster_path: PathBuf,
) -> Result<()> {
if !should_run {
return Ok(());
} }
fs::write(nfo_path, nfo_serializer.generate_nfo().await?.as_bytes()).await?; downloader.fetch(&video_model.cover, &poster_path).await
Ok(()) }
pub async fn generate_video_nfo(should_run: bool, video_model: &video::Model, nfo_path: PathBuf) -> Result<()> {
if !should_run {
return Ok(());
}
let nfo_serializer = NFOSerializer(ModelWrapper::Video(video_model), NFOMode::TVSHOW);
generate_nfo(nfo_serializer, nfo_path).await
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_template_usage() {
assert_eq!(
TEMPLATE.render("video", &json!({"bvid": "BV1b5411h7g7"})).unwrap(),
"BV1b5411h7g7"
);
}
} }