From 0ed02a9a46fe76e17fa7e1441db893c87674079f Mon Sep 17 00:00:00 2001
From: jachii <70832612+dobecad@users.noreply.github.com>
Date: Mon, 3 Jul 2023 03:35:59 -0400
Subject: [PATCH] Add `axum::extract::Query::try_from_uri` (#2058)

Co-authored-by: David Pedersen <david.pdrsn@gmail.com>
---
 axum/CHANGELOG.md         |  2 ++
 axum/src/extract/query.rs | 64 ++++++++++++++++++++++++++++++++++++---
 2 files changed, 62 insertions(+), 4 deletions(-)

diff --git a/axum/CHANGELOG.md b/axum/CHANGELOG.md
index 9d5388d1..b0b9c68f 100644
--- a/axum/CHANGELOG.md
+++ b/axum/CHANGELOG.md
@@ -53,6 +53,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - **breaking:** Only inherit fallbacks for routers nested with `Router::nest`.
   Routers nested with `Router::nest_service` will no longer inherit fallbacks ([#1956])
 - **fixed:** Don't remove the `Sec-WebSocket-Key` header in `WebSocketUpgrade` ([#1972])
+- **added:** Add `axum::extract::Query::try_from_uri` ([#2058])
 - **added:** Implement `IntoResponse` for `Box<str>` and `Box<[u8]>` ([#2035])
 
 [#1664]: https://github.com/tokio-rs/axum/pull/1664
@@ -64,6 +65,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 [#1868]: https://github.com/tokio-rs/axum/pull/1868
 [#1956]: https://github.com/tokio-rs/axum/pull/1956
 [#1972]: https://github.com/tokio-rs/axum/pull/1972
+[#2058]: https://github.com/tokio-rs/axum/pull/2058
 
 # 0.6.17 (25. April, 2023)
 
diff --git a/axum/src/extract/query.rs b/axum/src/extract/query.rs
index 81b28b61..37a40771 100644
--- a/axum/src/extract/query.rs
+++ b/axum/src/extract/query.rs
@@ -1,6 +1,6 @@
 use super::{rejection::*, FromRequestParts};
 use async_trait::async_trait;
-use http::request::Parts;
+use http::{request::Parts, Uri};
 use serde::de::DeserializeOwned;
 
 /// Extractor that deserializes query strings into some type.
@@ -55,10 +55,38 @@ where
     type Rejection = QueryRejection;
 
     async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
-        let query = parts.uri.query().unwrap_or_default();
-        let value =
+        Self::try_from_uri(&parts.uri)
+    }
+}
+
+impl<T> Query<T>
+where
+    T: DeserializeOwned,
+{
+    /// Attempts to construct a [`Query`] from a reference to a [`Uri`].
+    ///
+    /// # Example
+    /// ```
+    /// use axum::extract::Query;
+    /// use http::Uri;
+    /// use serde::Deserialize;
+    ///
+    /// #[derive(Deserialize)]
+    /// struct ExampleParams {
+    ///     foo: String,
+    ///     bar: u32,
+    /// }
+    ///
+    /// let uri: Uri = "http://example.com/path?foo=hello&bar=42".parse().unwrap();
+    /// let result: Query<ExampleParams> = Query::try_from_uri(&uri).unwrap();
+    /// assert_eq!(result.foo, String::from("hello"));
+    /// assert_eq!(result.bar, 42);
+    /// ```
+    pub fn try_from_uri(value: &Uri) -> Result<Self, QueryRejection> {
+        let query = value.query().unwrap_or_default();
+        let params =
             serde_urlencoded::from_str(query).map_err(FailedToDeserializeQueryString::from_err)?;
-        Ok(Query(value))
+        Ok(Query(params))
     }
 }
 
@@ -137,4 +165,32 @@ mod tests {
         let res = client.get("/?n=hi").send().await;
         assert_eq!(res.status(), StatusCode::BAD_REQUEST);
     }
+
+    #[test]
+    fn test_try_from_uri() {
+        #[derive(Deserialize)]
+        struct TestQueryParams {
+            foo: String,
+            bar: u32,
+        }
+        let uri: Uri = "http://example.com/path?foo=hello&bar=42".parse().unwrap();
+        let result: Query<TestQueryParams> = Query::try_from_uri(&uri).unwrap();
+        assert_eq!(result.foo, String::from("hello"));
+        assert_eq!(result.bar, 42);
+    }
+
+    #[test]
+    fn test_try_from_uri_with_invalid_query() {
+        #[derive(Deserialize)]
+        struct TestQueryParams {
+            _foo: String,
+            _bar: u32,
+        }
+        let uri: Uri = "http://example.com/path?foo=hello&bar=invalid"
+            .parse()
+            .unwrap();
+        let result: Result<Query<TestQueryParams>, _> = Query::try_from_uri(&uri);
+
+        assert!(result.is_err());
+    }
 }