From 1f2d3146e6f1fe0598ec04b41027c9807e9d2f9c Mon Sep 17 00:00:00 2001 From: yanhe Date: Mon, 25 Nov 2024 22:21:53 +0800 Subject: [PATCH] Add a from_path function to make it easier to return a stream of files with file name and content size. --- axum-extra/src/response/file_stream.rs | 85 +++++++++++++++++++------- examples/stream-to-file/src/main.rs | 35 +++-------- 2 files changed, 72 insertions(+), 48 deletions(-) diff --git a/axum-extra/src/response/file_stream.rs b/axum-extra/src/response/file_stream.rs index 408d404adf..88d05a5116 100644 --- a/axum-extra/src/response/file_stream.rs +++ b/axum-extra/src/response/file_stream.rs @@ -1,3 +1,5 @@ +use std::{io, path::PathBuf}; + use axum::{ body, response::{IntoResponse, Response}, @@ -6,6 +8,11 @@ use axum::{ use bytes::Bytes; use futures_util::TryStream; use http::{header, StatusCode}; +use tokio::fs::File; +use tokio_util::io::ReaderStream; + +/// Alias for `tokio_util::io::ReaderStream`. +pub type AsyncReaderStream = ReaderStream; /// Encapsulate the file stream. /// The encapsulated file stream construct requires passing in a stream @@ -20,7 +27,7 @@ use http::{header, StatusCode}; /// }; /// use axum_extra::response::file_stream::FileStream; /// use tokio::fs::File; -/// use tokio_util::io::ReaderStream ; +/// use tokio_util::io::ReaderStream; /// async fn file_stream() -> Result { /// let stream=ReaderStream::new(File::open("test.txt").await.map_err(|e| (StatusCode::NOT_FOUND, format!("File not found: {e}")))?); /// let file_stream_resp = FileStream::new(stream) @@ -61,6 +68,54 @@ where } } + /// Create a file stream from a file path. + /// # Examples + /// ``` + /// use axum::{ + /// http::StatusCode, + /// response::{Response, IntoResponse}, + /// Router, + /// routing::get + /// }; + /// use axum_extra::response::file_stream::FileStream; + /// use std::path::PathBuf; + /// use tokio_util::io::ReaderStream; + /// use tokio::fs::File; + /// async fn file_stream() -> Result { + /// Ok(FileStream::>::from_path(PathBuf::from("test.txt")) + /// .await + /// .map_err(|e| (StatusCode::NOT_FOUND, format!("File not found: {e}")))? + /// .into_response()) + /// } + /// let app = Router::new().route("/FileStreamDownload", get(file_stream)); + /// # let _: Router = app; + /// ``` + pub async fn from_path(path: PathBuf) -> io::Result> { + // open file + let file = File::open(&path).await?; + let mut content_size = None; + let mut file_name = None; + + // get file metadata length + if let Ok(metadata) = file.metadata().await { + content_size = Some(metadata.len()); + } + + // get file name + if let Some(file_name_os) = path.file_name() { + if let Some(file_name_str) = file_name_os.to_str() { + file_name = Some(file_name_str.to_owned()); + } + } + + // return FileStream + Ok(FileStream { + stream: ReaderStream::new(file), + file_name, + content_size, + }) + } + /// Set the file name of the file. pub fn file_name>(mut self, file_name: T) -> Self { self.file_name = Some(file_name.into()); @@ -111,8 +166,7 @@ mod tests { use axum::{extract::Request, routing::get, Router}; use body::Body; use http_body_util::BodyExt; - use std::io::{Cursor, SeekFrom}; - use tokio::io::AsyncSeekExt; + use std::io::Cursor; use tokio_util::io::ReaderStream; use tower::ServiceExt; @@ -164,24 +218,13 @@ mod tests { } #[tokio::test] - async fn response_half_file() -> Result<(), Box> { + async fn response_from_path() -> Result<(), Box> { let app = Router::new().route( - "/half_file", + "/from_path", get(move || async move { - let mut file = tokio::fs::File::open("CHANGELOG.md").await.unwrap(); - - // get file size - let file_size = file.metadata().await.unwrap().len(); - - // seek to the middle of the file - let mid_position = file_size / 2; - file.seek(SeekFrom::Start(mid_position)).await.unwrap(); - - // response file stream - let stream = ReaderStream::new(file); - FileStream::new(stream) - .file_name("CHANGELOG.md") - .content_size(mid_position) + FileStream::::from_path("CHANGELOG.md".into()) + .await + .unwrap() .into_response() }), ); @@ -190,7 +233,7 @@ mod tests { let response = app .oneshot( Request::builder() - .uri("/half_file") + .uri("/from_path") .body(Body::empty()) .unwrap(), ) @@ -212,7 +255,7 @@ mod tests { let file = tokio::fs::File::open("CHANGELOG.md").await.unwrap(); // get file size - let content_length = file.metadata().await.unwrap().len() / 2; + let content_length = file.metadata().await.unwrap().len(); assert_eq!( response diff --git a/examples/stream-to-file/src/main.rs b/examples/stream-to-file/src/main.rs index 532e2d48ed..e443681a17 100644 --- a/examples/stream-to-file/src/main.rs +++ b/examples/stream-to-file/src/main.rs @@ -13,14 +13,14 @@ use axum::{ routing::{get, post}, BoxError, Router, }; -use axum_extra::response::file_stream::FileStream; +use axum_extra::response::file_stream::{AsyncReaderStream, FileStream}; use futures::{Stream, TryStreamExt}; use std::io; use tokio::{ fs::File, io::{AsyncReadExt, AsyncSeekExt, BufWriter}, }; -use tokio_util::io::{ReaderStream, StreamReader}; +use tokio_util::io::StreamReader; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; const UPLOADS_DIRECTORY: &str = "uploads"; @@ -133,31 +133,12 @@ async fn show_form2() -> Html<&'static str> { /// A simpler file download handler that uses the `FileStream` response. /// Returns the entire file as a stream. async fn simpler_file_download_handler() -> Response { - let Ok(file) = File::open("./CHANGELOG.md").await else { - return (StatusCode::INTERNAL_SERVER_ERROR, "Failed to open file").into_response(); - }; - - let Ok(file_metadata) = file.metadata().await else { - return ( - StatusCode::INTERNAL_SERVER_ERROR, - "Failed to get file metadata", - ) - .into_response(); - }; - - // Constructing a Stream with ReaderStream - let stream = ReaderStream::new(file); - - // Use FileStream to return and set some information. - // Will set application/octet-stream in the header. - let file_stream_resp = FileStream::new(stream) - .file_name("test.txt") - .content_size(file_metadata.len()); - - //It is also possible to set only the stream FileStream will be automatically set on the http header. - //let file_stream_resp = FileStream::new(stream); - - file_stream_resp.into_response() + //If you want to simply return a file as a stream + // you can use the from_path method directly, passing in the path of the file to construct a stream with a header and length. + FileStream::::from_path("./CHANGELOG.md".into()) + .await + .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "Failed to open file").into_response()) + .into_response() } /// If you want to control the returned files in more detail you can implement a Stream