Update reqwest-response example (#2852)

This commit is contained in:
redforks 2024-07-30 14:55:15 +08:00 committed by GitHub
parent 304cbee7c9
commit 8dc371e9a2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -4,18 +4,16 @@
//! cargo run -p example-reqwest-response //! cargo run -p example-reqwest-response
//! ``` //! ```
use std::{convert::Infallible, time::Duration};
use axum::http::{HeaderMap, StatusCode};
use axum::{ use axum::{
body::{Body, Bytes}, body::{Body, Bytes},
extract::State, extract::State,
http::{HeaderName, HeaderValue}, http::StatusCode,
response::{IntoResponse, Response}, response::{IntoResponse, Response},
routing::get, routing::get,
Router, Router,
}; };
use reqwest::Client; use reqwest::Client;
use std::{convert::Infallible, time::Duration};
use tokio_stream::StreamExt; use tokio_stream::StreamExt;
use tower_http::trace::TraceLayer; use tower_http::trace::TraceLayer;
use tracing::Span; use tracing::Span;
@ -34,7 +32,7 @@ async fn main() {
let client = Client::new(); let client = Client::new();
let app = Router::new() let app = Router::new()
.route("/", get(proxy_via_reqwest)) .route("/", get(stream_reqwest_response))
.route("/stream", get(stream_some_data)) .route("/stream", get(stream_some_data))
// Add some logging so we can see the streams going through // Add some logging so we can see the streams going through
.layer(TraceLayer::new_for_http().on_body_chunk( .layer(TraceLayer::new_for_http().on_body_chunk(
@ -51,7 +49,7 @@ async fn main() {
axum::serve(listener, app).await.unwrap(); axum::serve(listener, app).await.unwrap();
} }
async fn proxy_via_reqwest(State(client): State<Client>) -> Response { async fn stream_reqwest_response(State(client): State<Client>) -> Response {
let reqwest_response = match client.get("http://127.0.0.1:3000/stream").send().await { let reqwest_response = match client.get("http://127.0.0.1:3000/stream").send().await {
Ok(res) => res, Ok(res) => res,
Err(err) => { Err(err) => {
@ -60,18 +58,8 @@ async fn proxy_via_reqwest(State(client): State<Client>) -> Response {
} }
}; };
let mut response_builder = Response::builder().status(reqwest_response.status().as_u16()); let mut response_builder = Response::builder().status(reqwest_response.status());
*response_builder.headers_mut().unwrap() = reqwest_response.headers().clone();
// Here the mapping of headers is required due to reqwest and axum differ on the http crate versions
let mut headers = HeaderMap::with_capacity(reqwest_response.headers().len());
headers.extend(reqwest_response.headers().into_iter().map(|(name, value)| {
let name = HeaderName::from_bytes(name.as_ref()).unwrap();
let value = HeaderValue::from_bytes(value.as_ref()).unwrap();
(name, value)
}));
*response_builder.headers_mut().unwrap() = headers;
response_builder response_builder
.body(Body::from_stream(reqwest_response.bytes_stream())) .body(Body::from_stream(reqwest_response.bytes_stream()))
// This unwrap is fine because the body is empty here // This unwrap is fine because the body is empty here