Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add an encapsulated file stream in axum-extra to make it more conveni… #3047

Merged
merged 22 commits into from
Dec 4, 2024
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions axum-extra/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ version = "0.10.0-alpha.1"
default = ["tracing", "multipart"]

async-read-body = ["dep:tokio-util", "tokio-util?/io", "dep:tokio"]
fileStream = ["dep:tokio-util", "tokio-util?/io", "dep:tokio"]
YanHeDoki marked this conversation as resolved.
Show resolved Hide resolved
attachment = ["dep:tracing"]
error_response = ["dep:tracing", "tracing/std"]
cookie = ["dep:cookie"]
Expand Down
4 changes: 2 additions & 2 deletions axum-extra/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,8 @@
//! `tracing` | Log rejections from built-in extractors | Yes
//! `typed-routing` | Enables the [`TypedPath`](crate::routing::TypedPath) routing utilities | No
//! `typed-header` | Enables the [`TypedHeader`] extractor and response | No
//!
//! [`axum`]: https://crates.io/crates/axum
//! `fileStream` | Enables the [`fileStream`](crate::response::file_stream) response | No
jplatte marked this conversation as resolved.
Show resolved Hide resolved
//! [`axum`]: <https://crates.io/crates/axum>
YanHeDoki marked this conversation as resolved.
Show resolved Hide resolved
jplatte marked this conversation as resolved.
Show resolved Hide resolved

#![warn(
clippy::all,
Expand Down
212 changes: 212 additions & 0 deletions axum-extra/src/response/file_stream.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,212 @@
use axum::{
body,
response::{IntoResponse, Response},
BoxError,
};
use bytes::Bytes;
use futures_util::TryStream;
use http::{header, StatusCode};

/// Encapsulate the file stream.
/// The encapsulated file stream construct requires passing in a stream
/// # Examples
jplatte marked this conversation as resolved.
Show resolved Hide resolved
///
/// ```
/// use axum::{
/// http::StatusCode,
/// response::{Response, IntoResponse},
/// Router,
/// routing::get
/// };
/// use axum_extra::response::file_stream::FileStream;
/// use tokio::fs::File;
/// use tokio_util::io::ReaderStream ;
/// async fn file_stream() -> Result<Response, (StatusCode, String)> {
/// 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)
/// .file_name("test.txt");
//
/// Ok(file_stream_resp.into_response())
/// }
/// let app = Router::new().route("/FileStreamDownload", get(file_stream));
/// # let _: Router = app;
/// ```
#[derive(Debug)]
pub struct FileStream<S>
where
S: TryStream + Send + 'static,
S::Ok: Into<Bytes>,
S::Error: Into<BoxError>,
{
jplatte marked this conversation as resolved.
Show resolved Hide resolved
/// stream.
pub stream: S,
/// The file name of the file.
pub file_name: Option<String>,
/// The size of the file.
pub content_size: Option<u64>,
}

impl<S> FileStream<S>
where
S: TryStream + Send + 'static,
S::Ok: Into<Bytes>,
S::Error: Into<BoxError>,
{
/// Create a file stream.
pub fn new(stream: S) -> Self {
Self {
stream,
file_name: None,
content_size: None,
}
}

/// Set the file name of the file.
YanHeDoki marked this conversation as resolved.
Show resolved Hide resolved
pub fn file_name<T: Into<String>>(mut self, file_name: T) -> Self {
jplatte marked this conversation as resolved.
Show resolved Hide resolved
self.file_name = Some(file_name.into());
self
}

/// Set the size of the file.
pub fn content_size<T: Into<u64>>(mut self, len: T) -> Self {
jplatte marked this conversation as resolved.
Show resolved Hide resolved
self.content_size = Some(len.into());
jplatte marked this conversation as resolved.
Show resolved Hide resolved
self
}
}

impl<S> IntoResponse for FileStream<S>
where
S: TryStream + Send + 'static,
S::Ok: Into<Bytes>,
S::Error: Into<BoxError>,
{
fn into_response(self) -> Response {
let mut resp = Response::builder().header(header::CONTENT_TYPE, "application/octet-stream");

if let Some(file_name) = self.file_name {
resp = resp.header(
header::CONTENT_DISPOSITION,
format!("attachment; filename=\"{}\"", file_name),
);
};
YanHeDoki marked this conversation as resolved.
Show resolved Hide resolved

if let Some(content_size) = self.content_size {
resp = resp.header(header::CONTENT_LENGTH, content_size);
};
YanHeDoki marked this conversation as resolved.
Show resolved Hide resolved

resp.body(body::Body::from_stream(self.stream))
.unwrap_or_else(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("build FileStream responsec error:{}", e),
YanHeDoki marked this conversation as resolved.
Show resolved Hide resolved
)
.into_response()
})
}
}

#[cfg(test)]
mod tests {
use super::*;
use axum::{extract::Request, routing::get, Router};
use body::Body;
use http_body_util::BodyExt;
use tokio::io::AsyncSeekExt;
use std::io::{Cursor, SeekFrom};
use tokio_util::io::ReaderStream;
use tower::ServiceExt;

#[tokio::test]
async fn response_file_stream() -> Result<(), Box<dyn std::error::Error>> {
let app = Router::new().route(
"/file",
get(|| async {
// Simulating a file stream
let file_content = b"Hello, this is the simulated file content!".to_vec();
let size = file_content.len() as u64;
let reader = Cursor::new(file_content);

// response file stream
let stream = ReaderStream::new(reader);
let resp = FileStream::new(stream)
.file_name("test")
.content_size(size)
.into_response();
resp
}),
);

// Simulating a GET request
let response = app
.oneshot(Request::builder().uri("/file").body(Body::empty())?)
.await?;

// Validate Response Status Code
assert_eq!(response.status(), StatusCode::OK);

// Validate Response Headers
assert_eq!(
response.headers().get("content-type").unwrap(),
"application/octet-stream"
);
assert_eq!(
response.headers().get("content-disposition").unwrap(),
"attachment; filename=\"test\""
);
assert_eq!(response.headers().get("content-length").unwrap(), "42");

// Validate Response Body
let body: &[u8] = &response.into_body().collect().await?.to_bytes();
assert_eq!(
std::str::from_utf8(body)?,
"Hello, this is the simulated file content!"
);
Ok(())
}

#[tokio::test]
async fn response_half_file() -> Result<(), Box<dyn std::error::Error>> {
let app = Router::new().route(
"/half_file",
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);
let resp = FileStream::new(stream)
.file_name("CHANGELOG.md")
.content_size(mid_position)
.into_response();
resp
}),
);

// Simulating a GET request
let response = app
.oneshot(Request::builder().uri("/half_file").body(Body::empty()).unwrap())
.await
.unwrap();

// Validate Response Status Code
assert_eq!(response.status(), StatusCode::OK);

// Validate Response Headers
assert_eq!(
response.headers().get("content-type").unwrap(),
"application/octet-stream"
);
assert_eq!(
response.headers().get("content-disposition").unwrap(),
"attachment; filename=\"CHANGELOG.md\""
);
assert_eq!(response.headers().get("content-length").unwrap(), "8098");
YanHeDoki marked this conversation as resolved.
Show resolved Hide resolved
Ok(())
}
}
4 changes: 4 additions & 0 deletions axum-extra/src/response/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,10 @@ pub mod multiple;
#[cfg(feature = "error_response")]
mod error_response;

#[cfg(feature = "fileStream")]
/// Module for handling file streams.
pub mod file_stream;

#[cfg(feature = "error_response")]
pub use error_response::InternalServerError;

Expand Down
2 changes: 2 additions & 0 deletions examples/stream-to-file/Cargo.toml
YanHeDoki marked this conversation as resolved.
Show resolved Hide resolved
Original file line number Diff line number Diff line change
Expand Up @@ -11,3 +11,5 @@ tokio = { version = "1.0", features = ["full"] }
tokio-util = { version = "0.7", features = ["io"] }
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
async-stream = "0.3"
axum-extra = { path = "../../axum-extra", features = ["fileStream"] }
YanHeDoki marked this conversation as resolved.
Show resolved Hide resolved
Loading
Loading