Introduce lambda_http::IntoResponse (awslabs#31)

softprops · davidbarsky · commit 771e5983111d · 2018-12-04T11:23:01.000-05:00
diff --git a/lambda-http/examples/basic.rs b/lambda-http/examples/basic.rs
@@ -3,7 +3,7 @@ extern crate lambda_runtime as runtime;
 extern crate log;
 extern crate simple_logger;
 
-use http::{lambda, Body, Request, RequestExt, Response};
+use http::{lambda, Body, IntoResponse, Request, RequestExt, Response};
 use runtime::{error::HandlerError, Context};
 
 use log::error;
@@ -16,14 +16,14 @@ fn main() -> Result<(), Box<dyn Error>> {
     Ok(())
 }
 
-fn my_handler(e: Request, c: Context) -> Result<Response<Body>, HandlerError> {
+fn my_handler(e: Request, c: Context) -> Result<impl IntoResponse, HandlerError> {
     Ok(match e.query_string_parameters().get("first_name") {
-        Some(first_name) => Response::new(format!("Hello, {}!", first_name).into()),
+        Some(first_name) => format!("Hello, {}!", first_name).into_response(),
         _ => {
             error!("Empty first name in request {}", c.aws_request_id);
             Response::builder()
                 .status(400)
-                .body::<Body>("Empty first name".into())
+                .body("Empty first name".into())
                 .expect("failed to render response")
         }
     })
diff --git a/lambda-http/src/lib.rs b/lambda-http/src/lib.rs
@@ -9,7 +9,7 @@
 //! extern crate lambda_runtime as lambda;
 //!
 //! use lambda::{Context, HandlerError};
-//! use lambda_http::{Body, Request, Response, RequestExt};
+//! use lambda_http::{Request, IntoResponse, RequestExt};
 //!
 //! fn main() {
 //!   lambda!(handler)
@@ -18,16 +18,14 @@
 //! fn handler(
 //!   request: Request,
 //!   ctx: Context
-//! ) -> Result<Response<Body>, HandlerError> {
+//! ) -> Result<impl IntoResponse, HandlerError> {
 //!   Ok(
-//!     Response::new(
 //!       format!(
 //!         "hello {}",
 //!         request.query_string_parameters()
 //!           .get("name")
 //!           .unwrap_or_else(|| "stranger")
-//!       ).into()
-//!     )
+//!       )
 //!   )
 //! }
 //! ```
@@ -60,22 +58,23 @@ pub use body::Body;
 pub use ext::RequestExt;
 use request::GatewayRequest;
 use response::GatewayResponse;
+pub use response::IntoResponse;
 pub use strmap::StrMap;
 
 /// Type alias for `http::Request`s with a fixed `lambda_http::Body` body
 pub type Request = HttpRequest<Body>;
 
 /// Functions acting as API Gateway handlers must conform to this type.
-pub trait Handler {
+pub trait Handler<R> {
     /// Run the handler.
-    fn run(&mut self, event: Request, ctx: Context) -> Result<Response<Body>, HandlerError>;
+    fn run(&mut self, event: Request, ctx: Context) -> Result<R, HandlerError>;
 }
 
-impl<F> Handler for F
+impl<F, R> Handler<R> for F
 where
-    F: FnMut(Request, Context) -> Result<Response<Body>, HandlerError>,
+    F: FnMut(Request, Context) -> Result<R, HandlerError>,
 {
-    fn run(&mut self, event: Request, ctx: Context) -> Result<Response<Body>, HandlerError> {
+    fn run(&mut self, event: Request, ctx: Context) -> Result<R, HandlerError> {
         (*self)(event, ctx)
     }
 }
@@ -88,11 +87,17 @@ where
 ///
 /// # Panics
 /// The function panics if the Lambda environment variables are not set.
-pub fn start(f: impl Handler, runtime: Option<TokioRuntime>) {
+pub fn start<R>(f: impl Handler<R>, runtime: Option<TokioRuntime>)
+where
+    R: IntoResponse,
+{
     // handler requires a mutable ref
     let mut func = f;
     lambda::start(
-        |req: GatewayRequest, ctx: Context| func.run(req.into(), ctx).map(GatewayResponse::from),
+        |req: GatewayRequest, ctx: Context| {
+            func.run(req.into(), ctx)
+                .map(|resp| GatewayResponse::from(resp.into_response()))
+        },
         runtime,
     )
 }
diff --git a/lambda-http/src/response.rs b/lambda-http/src/response.rs
@@ -5,7 +5,7 @@ use std::ops::Not;
 
 use http::{
     header::{HeaderMap, HeaderValue},
-    Response as HttpResponse,
+    Response,
 };
 use serde::{
     ser::{Error as SerError, SerializeMap},
@@ -50,11 +50,11 @@ where
     map.end()
 }
 
-impl<T> From<HttpResponse<T>> for GatewayResponse
+impl<T> From<Response<T>> for GatewayResponse
 where
     T: Into<Body>,
 {
-    fn from(value: HttpResponse<T>) -> Self {
+    fn from(value: Response<T>) -> Self {
         let (parts, bod) = value.into_parts();
         let (is_base64_encoded, body) = match bod.into() {
             Body::Empty => (false, None),
@@ -70,11 +70,89 @@ where
     }
 }
 
+/// A conversion of self into a `Response`
+///
+/// Implementations for `Response<B> where B: Into<Body>`,
+/// `B where B: Into<Body>` and `serde_json::Value` are provided
+/// by default
+///
+/// # example
+///
+/// ```rust
+/// use lambda_http::{Body, IntoResponse, Response};
+///
+/// assert_eq!(
+///   "hello".into_response().body(),
+///   Response::new(Body::from("hello")).body()
+/// );
+/// ```
+pub trait IntoResponse {
+    /// Return a translation of `self` into a `Response<Body>`
+    fn into_response(self) -> Response<Body>;
+}
+
+impl<B> IntoResponse for Response<B>
+where
+    B: Into<Body>,
+{
+    fn into_response(self) -> Response<Body> {
+        let (parts, body) = self.into_parts();
+        Response::from_parts(parts, body.into())
+    }
+}
+
+impl<B> IntoResponse for B
+where
+    B: Into<Body>,
+{
+    fn into_response(self) -> Response<Body> {
+        Response::new(self.into())
+    }
+}
+
+impl IntoResponse for serde_json::Value {
+    fn into_response(self) -> Response<Body> {
+        Response::builder()
+            .header(http::header::CONTENT_TYPE, "application/json")
+            .body(
+                serde_json::to_string(&self)
+                    .expect("unable to serialize serde_json::Value")
+                    .into(),
+            )
+            .expect("unable to build http::Response")
+    }
+}
+
 #[cfg(test)]
 mod tests {
 
-    use super::GatewayResponse;
-    use serde_json;
+    use super::{Body, GatewayResponse, IntoResponse};
+    use serde_json::{self, json};
+
+    #[test]
+    fn json_into_response() {
+        let response = json!({ "hello": "lambda"}).into_response();
+        match response.body() {
+            Body::Text(json) => assert_eq!(json, r#"{"hello":"lambda"}"#),
+            _ => panic!("invalid body"),
+        }
+        assert_eq!(
+            response
+                .headers()
+                .get(http::header::CONTENT_TYPE)
+                .map(|h| h.to_str().expect("invalid header")),
+            Some("application/json")
+        )
+    }
+
+    #[test]
+    fn text_into_response() {
+        let response = "text".into_response();
+        match response.body() {
+            Body::Text(text) => assert_eq!(text, "text"),
+            _ => panic!("invalid body"),
+        }
+    }
 
     #[test]
     fn default_response() {
