feat: 重构优化部分 API，支持重置全体失败的任务 (#351)

crates/bili_sync/src/api/handler.rs

@@ -1,27 +1,30 @@
+use std::collections::HashSet;
 use std::sync::Arc;
 
-use anyhow::{Result, anyhow};
+use anyhow::Result;
 use axum::extract::{Extension, Path, Query};
 use bili_sync_entity::*;
-use bili_sync_migration::{Expr, OnConflict};
+use bili_sync_migration::Expr;
 use sea_orm::{
-    ColumnTrait, DatabaseConnection, EntityTrait, IntoActiveModel, PaginatorTrait, QueryFilter, QueryOrder,
-    QuerySelect, Set, TransactionTrait, Unchanged,
+    ColumnTrait, DatabaseConnection, EntityTrait, PaginatorTrait, QueryFilter, QueryOrder, QuerySelect,
+    TransactionTrait,
 };
 use utoipa::OpenApi;
 
 use crate::api::auth::OpenAPIAuth;
 use crate::api::error::InnerApiError;
+use crate::api::helper::{update_page_download_status, update_video_download_status};
 use crate::api::request::VideosRequest;
 use crate::api::response::{
-    PageInfo, ResetVideoResponse, VideoInfo, VideoResponse, VideoSource, VideoSourcesResponse, VideosResponse,
+    PageInfo, ResetAllVideosResponse, ResetVideoResponse, VideoInfo, VideoResponse, VideoSource, VideoSourcesResponse,
+    VideosResponse,
 };
 use crate::api::wrapper::{ApiError, ApiResponse};
 use crate::utils::status::{PageStatus, VideoStatus};
 
 #[derive(OpenApi)]
 #[openapi(
-    paths(get_video_sources, get_videos, get_video, reset_video),
+    paths(get_video_sources, get_videos, get_video, reset_video, reset_all_videos),
     modifiers(&OpenAPIAuth),
     security(
         ("Token" = []),
@@ -40,33 +43,35 @@ pub struct ApiDoc;
 pub async fn get_video_sources(
     Extension(db): Extension<Arc<DatabaseConnection>>,
 ) -> Result<ApiResponse<VideoSourcesResponse>, ApiError> {
-    Ok(ApiResponse::ok(VideoSourcesResponse {
-        collection: collection::Entity::find()
+    let (collection, favorite, submission, watch_later) = tokio::try_join!(
+        collection::Entity::find()
             .select_only()
             .columns([collection::Column::Id, collection::Column::Name])
             .into_model::<VideoSource>()
-            .all(db.as_ref())
-            .await?,
-        favorite: favorite::Entity::find()
+            .all(db.as_ref()),
+        favorite::Entity::find()
             .select_only()
             .columns([favorite::Column::Id, favorite::Column::Name])
             .into_model::<VideoSource>()
-            .all(db.as_ref())
-            .await?,
-        submission: submission::Entity::find()
+            .all(db.as_ref()),
+        submission::Entity::find()
             .select_only()
             .column(submission::Column::Id)
             .column_as(submission::Column::UpperName, "name")
             .into_model::<VideoSource>()
-            .all(db.as_ref())
-            .await?,
-        watch_later: watch_later::Entity::find()
+            .all(db.as_ref()),
+        watch_later::Entity::find()
             .select_only()
             .column(watch_later::Column::Id)
             .column_as(Expr::value("稍后再看"), "name")
             .into_model::<VideoSource>()
             .all(db.as_ref())
-            .await?,
+    )?;
+    Ok(ApiResponse::ok(VideoSourcesResponse {
+        collection,
+        favorite,
+        submission,
+        watch_later,
     }))
 }
 
@@ -103,25 +108,15 @@ pub async fn get_videos(
     let (page, page_size) = if let (Some(page), Some(page_size)) = (params.page, params.page_size) {
         (page, page_size)
     } else {
-        (1, 10)
+        (0, 10)
     };
     Ok(ApiResponse::ok(VideosResponse {
         videos: query
             .order_by_desc(video::Column::Id)
-            .select_only()
-            .columns([
-                video::Column::Id,
-                video::Column::Name,
-                video::Column::UpperName,
-                video::Column::DownloadStatus,
-            ])
-            .into_tuple::<(i32, String, String, u32)>()
+            .into_partial_model::<VideoInfo>()
             .paginate(db.as_ref(), page_size)
             .fetch_page(page)
-            .await?
-            .into_iter()
-            .map(VideoInfo::from)
-            .collect(),
+            .await?,
         total_count,
     }))
 }
@@ -138,40 +133,22 @@ pub async fn get_video(
     Path(id): Path<i32>,
     Extension(db): Extension<Arc<DatabaseConnection>>,
 ) -> Result<ApiResponse<VideoResponse>, ApiError> {
-    let video_info = video::Entity::find_by_id(id)
-        .select_only()
-        .columns([
-            video::Column::Id,
-            video::Column::Name,
-            video::Column::UpperName,
-            video::Column::DownloadStatus,
-        ])
-        .into_tuple::<(i32, String, String, u32)>()
-        .one(db.as_ref())
-        .await?
-        .map(VideoInfo::from);
+    let (video_info, pages_info) = tokio::try_join!(
+        video::Entity::find_by_id(id)
+            .into_partial_model::<VideoInfo>()
+            .one(db.as_ref()),
+        page::Entity::find()
+            .filter(page::Column::VideoId.eq(id))
+            .order_by_asc(page::Column::Cid)
+            .into_partial_model::<PageInfo>()
+            .all(db.as_ref())
+    )?;
     let Some(video_info) = video_info else {
         return Err(InnerApiError::NotFound(id).into());
     };
-    let pages = page::Entity::find()
-        .filter(page::Column::VideoId.eq(id))
-        .order_by_asc(page::Column::Pid)
-        .select_only()
-        .columns([
-            page::Column::Id,
-            page::Column::Pid,
-            page::Column::Name,
-            page::Column::DownloadStatus,
-        ])
-        .into_tuple::<(i32, i32, String, u32)>()
-        .all(db.as_ref())
-        .await?
-        .into_iter()
-        .map(PageInfo::from)
-        .collect();
     Ok(ApiResponse::ok(VideoResponse {
         video: video_info,
-        pages,
+        pages: pages_info,
     }))
 }
 
@@ -180,73 +157,129 @@ pub async fn get_video(
     post,
     path = "/api/videos/{id}/reset",
     responses(
-        (status = 200, body = ApiResponse<ResetVideoResponse> ),
+        (status = 200, body = ApiResponse<ResetVideoResponse>),
     )
 )]
 pub async fn reset_video(
     Path(id): Path<i32>,
     Extension(db): Extension<Arc<DatabaseConnection>>,
 ) -> Result<ApiResponse<ResetVideoResponse>, ApiError> {
-    let txn = db.begin().await?;
-    let video_status: Option<u32> = video::Entity::find_by_id(id)
-        .select_only()
-        .column(video::Column::DownloadStatus)
-        .into_tuple()
-        .one(&txn)
-        .await?;
-    let Some(video_status) = video_status else {
-        return Err(anyhow!(InnerApiError::NotFound(id)).into());
+    let (video_info, pages_info) = tokio::try_join!(
+        video::Entity::find_by_id(id)
+            .into_partial_model::<VideoInfo>()
+            .one(db.as_ref()),
+        page::Entity::find()
+            .filter(page::Column::VideoId.eq(id))
+            .order_by_asc(page::Column::Cid)
+            .into_partial_model::<PageInfo>()
+            .all(db.as_ref())
+    )?;
+    let Some(mut video_info) = video_info else {
+        return Err(InnerApiError::NotFound(id).into());
     };
-    let resetted_pages_model: Vec<_> = page::Entity::find()
-        .filter(page::Column::VideoId.eq(id))
-        .all(&txn)
-        .await?
+    let resetted_pages_info = pages_info
         .into_iter()
-        .filter_map(|mut model| {
-            let mut page_status = PageStatus::from(model.download_status);
+        .filter_map(|mut page_info| {
+            let mut page_status = PageStatus::from(page_info.download_status);
             if page_status.reset_failed() {
-                model.download_status = page_status.into();
-                Some(model)
+                page_info.download_status = page_status.into();
+                Some(page_info)
             } else {
                 None
             }
         })
-        .collect();
-    let mut video_status = VideoStatus::from(video_status);
-    let mut should_update_video = video_status.reset_failed();
-    if !resetted_pages_model.is_empty() {
-        // 视频状态标志的第 5 位表示是否有分 P 下载失败，如果有需要重置的分页，需要同时重置视频的该状态
-        video_status.set(4, 0);
-        should_update_video = true;
+        .collect::<Vec<_>>();
+    let mut video_status = VideoStatus::from(video_info.download_status);
+    let mut video_resetted = video_status.reset_failed();
+    if !resetted_pages_info.is_empty() {
+        video_status.set(4, 0); //  将“分P下载”重置为 0
+        video_resetted = true;
     }
-    if should_update_video {
-        video::Entity::update(video::ActiveModel {
-            id: Unchanged(id),
-            download_status: Set(video_status.into()),
-            ..Default::default()
-        })
-        .exec(&txn)
-        .await?;
+    let resetted_videos_info = if video_resetted {
+        video_info.download_status = video_status.into();
+        vec![video_info.clone()]
+    } else {
+        vec![]
+    };
+    let resetted = !resetted_videos_info.is_empty() || !resetted_pages_info.is_empty();
+    if resetted {
+        let txn = db.begin().await?;
+        if !resetted_videos_info.is_empty() {
+            // 只可能有 1 个元素，所以不用 batch
+            update_video_download_status(&txn, &resetted_videos_info, None).await?;
+        }
+        if !resetted_pages_info.is_empty() {
+            update_page_download_status(&txn, &resetted_pages_info, Some(500)).await?;
+        }
+        txn.commit().await?;
     }
-    let resetted_pages_id: Vec<_> = resetted_pages_model.iter().map(|model| model.id).collect();
-    let resetted_pages_model: Vec<page::ActiveModel> = resetted_pages_model
-        .into_iter()
-        .map(|model| model.into_active_model())
-        .collect();
-    for page_trunk in resetted_pages_model.chunks(50) {
-        page::Entity::insert_many(page_trunk.to_vec())
-            .on_conflict(
-                OnConflict::column(page::Column::Id)
-                    .update_column(page::Column::DownloadStatus)
-                    .to_owned(),
-            )
-            .exec(&txn)
-            .await?;
-    }
-    txn.commit().await?;
     Ok(ApiResponse::ok(ResetVideoResponse {
-        resetted: should_update_video,
-        video: id,
-        pages: resetted_pages_id,
+        resetted,
+        video: video_info,
+        pages: resetted_pages_info,
+    }))
+}
+
+/// 重置所有视频和页面的失败状态为未下载状态，这样在下次下载任务中会触发重试
+#[utoipa::path(
+    post,
+    path = "/api/videos/reset-all",
+    responses(
+        (status = 200, body = ApiResponse<ResetAllVideosResponse>),
+    )
+)]
+pub async fn reset_all_videos(
+    Extension(db): Extension<Arc<DatabaseConnection>>,
+) -> Result<ApiResponse<ResetAllVideosResponse>, ApiError> {
+    // 先查询所有视频和页面数据
+    let (all_videos, all_pages) = tokio::try_join!(
+        video::Entity::find().into_partial_model::<VideoInfo>().all(db.as_ref()),
+        page::Entity::find().into_partial_model::<PageInfo>().all(db.as_ref())
+    )?;
+    let resetted_pages_info = all_pages
+        .into_iter()
+        .filter_map(|mut page_info| {
+            let mut page_status = PageStatus::from(page_info.download_status);
+            if page_status.reset_failed() {
+                page_info.download_status = page_status.into();
+                Some(page_info)
+            } else {
+                None
+            }
+        })
+        .collect::<Vec<_>>();
+    let video_ids_with_resetted_pages: HashSet<i32> = resetted_pages_info.iter().map(|page| page.video_id).collect();
+    let resetted_videos_info = all_videos
+        .into_iter()
+        .filter_map(|mut video_info| {
+            let mut video_status = VideoStatus::from(video_info.download_status);
+            let mut video_resetted = video_status.reset_failed();
+            if video_ids_with_resetted_pages.contains(&video_info.id) {
+                video_status.set(4, 0); // 将"分P下载"重置为 0
+                video_resetted = true;
+            }
+            if video_resetted {
+                video_info.download_status = video_status.into();
+                Some(video_info)
+            } else {
+                None
+            }
+        })
+        .collect::<Vec<_>>();
+    let resetted = !(resetted_videos_info.is_empty() && resetted_pages_info.is_empty());
+    if resetted {
+        let txn = db.begin().await?;
+        if !resetted_videos_info.is_empty() {
+            update_video_download_status(&txn, &resetted_videos_info, Some(500)).await?;
+        }
+        if !resetted_pages_info.is_empty() {
+            update_page_download_status(&txn, &resetted_pages_info, Some(500)).await?;
+        }
+        txn.commit().await?;
+    }
+    Ok(ApiResponse::ok(ResetAllVideosResponse {
+        resetted,
+        resetted_videos_count: resetted_videos_info.len(),
+        resetted_pages_count: resetted_pages_info.len(),
     }))
 }