Make lambda::Context a first-class part of the Handler API (#233)

Author: softprops

lambda-attributes/src/lib.rs

@@ -38,42 +38,48 @@ pub fn lambda(attr: TokenStream, item: TokenStream) -> TokenStream {
     }
 
     let result = match inputs.len() {
-        1 => {
-            let event = match inputs.first().unwrap() {
+        2 => {
+            let event = match inputs.first().expect("expected event argument") {
                 FnArg::Typed(arg) => arg,
                 _ => {
                     let tokens = quote_spanned! { inputs.span() =>
-                        compile_error!("fn main must take a fully formed argument");
+                        compile_error!("fn main's first argument must be fully formed");
                     };
                     return TokenStream::from(tokens);
                 }
             };
-            let arg_name = &event.pat;
-            let arg_type = &event.ty;
+            let event_name = &event.pat;
+            let event_type = &event.ty;
+            let context = match inputs.iter().nth(1).expect("expected context argument") {
+                FnArg::Typed(arg) => arg,
+                _ => {
+                    let tokens = quote_spanned! { inputs.span() =>
+                        compile_error!("fn main's second argument must be fully formed");
+                    };
+                    return TokenStream::from(tokens);
+                }
+            };
+            let context_name = &context.pat;
+            let context_type = &context.ty;
 
             if is_http(&args) {
                 quote_spanned! { input.span() =>
-                    use lambda_http::lambda::LambdaCtx;
 
                     #(#attrs)*
                     #asyncness fn main() {
-                        async fn actual(#arg_name: #arg_type) #ret {
-                            #body
-                        }
+                        async fn actual(#event_name: #event_type, #context_name: #context_type) #ret #body
+
                         let f = lambda_http::handler(actual);
                         lambda_http::lambda::run(f).await.unwrap();
                     }
                 }
             } else {
                 quote_spanned! { input.span() =>
 
-                    use lambda::LambdaCtx;
-
                     #(#attrs)*
                     #asyncness fn main() {
-                        async fn actual(#arg_name: #arg_type) #ret {
-                            #body
-                        }
+                        async fn actual(#event_name: #event_type, #context_name: #context_type) #ret #body
+
                         let f = lambda::handler_fn(actual);
                         lambda::run(f).await.unwrap();
                     }
@@ -82,7 +88,7 @@ pub fn lambda(attr: TokenStream, item: TokenStream) -> TokenStream {
         }
         _ => {
             let tokens = quote_spanned! { inputs.span() =>
-                compile_error!("The #[lambda] macro can accept only a single argument.");
+                compile_error!("The #[lambda] macro can expects two arguments: a triggered event and lambda context.");
             };
             return TokenStream::from(tokens);
         }

lambda-http/examples/hello-http-without-macros.rs

@@ -1,4 +1,8 @@
-use lambda_http::{handler, lambda, IntoResponse, Request, RequestExt, Response};
+use lambda_http::{
+    handler,
+    lambda::{self, Context},
+    IntoResponse, Request, RequestExt, Response,
+};
 
 type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
 
@@ -8,7 +12,7 @@ async fn main() -> Result<(), Error> {
     Ok(())
 }
 
-async fn func(event: Request) -> Result<impl IntoResponse, Error> {
+async fn func(event: Request, _: Context) -> Result<impl IntoResponse, Error> {
     Ok(match event.query_string_parameters().get("first_name") {
         Some(first_name) => format!("Hello, {}!", first_name).into_response(),
         _ => Response::builder()

lambda-http/examples/hello-http.rs

@@ -1,9 +1,12 @@
-use lambda_http::{lambda, IntoResponse, Request};
+use lambda_http::{
+    lambda::{lambda, Context},
+    IntoResponse, Request,
+};
 
 type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
 
 #[lambda(http)]
 #[tokio::main]
-async fn main(_: Request) -> Result<impl IntoResponse, Error> {
+async fn main(_: Request, _: Context) -> Result<impl IntoResponse, Error> {
     Ok("👋 world")
 }

lambda-http/src/ext.rs

@@ -67,7 +67,7 @@ impl Error for PayloadError {
 /// as well as `{"x":1, "y":2}` respectively.
 ///
 /// ```rust,no_run
-/// use lambda_http::{handler, lambda, Body, IntoResponse, Request, Response, RequestExt};
+/// use lambda_http::{handler, lambda::{self, Context}, Body, IntoResponse, Request, Response, RequestExt};
 /// use serde_derive::Deserialize;
 ///
 /// type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
@@ -87,7 +87,8 @@ impl Error for PayloadError {
 /// }
 ///
 /// async fn add(
-///   request: Request
+///   request: Request,
+///   _: Context
 /// ) -> Result<Response<Body>, Error> {
 ///   let args: Args = request.payload()
 ///     .unwrap_or_else(|_parse_err| None)

lambda-http/src/lib.rs

@@ -22,13 +22,13 @@
 //! The full body of your `main` function will be executed on **every** invocation of your lambda task.
 //!
 //! ```rust,no_run
-//! use lambda_http::{lambda, Request, IntoResponse};
+//! use lambda_http::{lambda::{lambda, Context}, Request, IntoResponse};
 //!
 //! type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
 //!
 //! #[lambda(http)]
 //! #[tokio::main]
-//! async fn main(_: Request) -> Result<impl IntoResponse, Error> {
+//! async fn main(_: Request, _: Context) -> Result<impl IntoResponse, Error> {
 //!     Ok("👋 world!")
 //! }
 //! ```
@@ -48,7 +48,7 @@
 //! async fn main() -> Result<(), Error> {
 //!     // initialize dependencies once here for the lifetime of your
 //!     // lambda task
-//!     lambda::run(handler(|request| async { Ok("👋 world!") })).await?;
+//!     lambda::run(handler(|request, context| async { Ok("👋 world!") })).await?;
 //!     Ok(())
 //! }
 //!
@@ -60,7 +60,7 @@
 //! with the [`RequestExt`](trait.RequestExt.html) trait.
 //!
 //! ```rust,no_run
-//! use lambda_http::{handler, lambda, IntoResponse, Request, RequestExt};
+//! use lambda_http::{handler, lambda::{self, Context}, IntoResponse, Request, RequestExt};
 //!
 //! type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
 //!
@@ -72,6 +72,7 @@
 //!
 //! async fn hello(
 //!     request: Request,
+//!     _: Context
 //! ) -> Result<impl IntoResponse, Error> {
 //!     Ok(format!(
 //!         "hello {}",
@@ -90,7 +91,7 @@ extern crate maplit;
 
 pub use http::{self, Response};
 use lambda::Handler as LambdaHandler;
-pub use lambda::{self};
+pub use lambda::{self, Context};
 pub use lambda_attributes::lambda;
 
 mod body;
@@ -103,7 +104,7 @@ use crate::{request::LambdaRequest, response::LambdaResponse};
 use std::{
     future::Future,
     pin::Pin,
-    task::{Context, Poll},
+    task::{Context as TaskContext, Poll},
 };
 
 /// Error type that lambdas may result in
@@ -123,7 +124,7 @@ pub trait Handler: Sized {
     /// The type of Future this Handler will return
     type Fut: Future<Output = Result<Self::Response, Self::Error>> + 'static;
     /// Function used to execute handler behavior
-    fn call(&mut self, event: Request) -> Self::Fut;
+    fn call(&mut self, event: Request, context: Context) -> Self::Fut;
 }
 
 /// Adapts a [`Handler`](trait.Handler.html) to the `lambda::run` interface
@@ -134,15 +135,15 @@ pub fn handler<H: Handler>(handler: H) -> Adapter<H> {
 /// An implementation of `Handler` for a given closure return a `Future` representing the computed response
 impl<F, R, Fut> Handler for F
 where
-    F: FnMut(Request) -> Fut,
+    F: FnMut(Request, Context) -> Fut,
     R: IntoResponse,
     Fut: Future<Output = Result<R, Error>> + Send + 'static,
 {
     type Response = R;
     type Error = Error;
     type Fut = Fut;
-    fn call(&mut self, event: Request) -> Self::Fut {
-        (*self)(event)
+    fn call(&mut self, event: Request, context: Context) -> Self::Fut {
+        (*self)(event, context)
     }
 }
 
@@ -157,7 +158,7 @@ where
     R: IntoResponse,
 {
     type Output = Result<LambdaResponse, E>;
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
+    fn poll(mut self: Pin<&mut Self>, cx: &mut TaskContext) -> Poll<Self::Output> {
         match self.fut.as_mut().poll(cx) {
             Poll::Ready(result) => {
                 Poll::Ready(result.map(|resp| LambdaResponse::from_response(self.is_alb, resp.into_response())))
@@ -182,17 +183,17 @@ impl<H: Handler> Handler for Adapter<H> {
     type Response = H::Response;
     type Error = H::Error;
     type Fut = H::Fut;
-    fn call(&mut self, event: Request) -> Self::Fut {
-        self.handler.call(event)
+    fn call(&mut self, event: Request, context: Context) -> Self::Fut {
+        self.handler.call(event, context)
     }
 }
 
 impl<H: Handler> LambdaHandler<LambdaRequest<'_>, LambdaResponse> for Adapter<H> {
     type Error = H::Error;
     type Fut = TransformResponse<H::Response, Self::Error>;
-    fn call(&mut self, event: LambdaRequest<'_>) -> Self::Fut {
+    fn call(&mut self, event: LambdaRequest<'_>, context: Context) -> Self::Fut {
         let is_alb = event.is_alb();
-        let fut = Box::pin(self.handler.call(event.into()));
+        let fut = Box::pin(self.handler.call(event.into(), context));
         TransformResponse { is_alb, fut }
     }
 }

lambda/examples/hello-with-ctx.rs

@@ -1,4 +1,4 @@
-use lambda::handler_fn;
+use lambda::{handler_fn, Context};
 use serde_json::Value;
 
 type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
@@ -10,6 +10,6 @@ async fn main() -> Result<(), Error> {
     Ok(())
 }
 
-async fn func(event: Value) -> Result<Value, Error> {
+async fn func(event: Value, _: Context) -> Result<Value, Error> {
     Ok(event)
 }

lambda/examples/hello-without-macro.rs

@@ -1,10 +1,10 @@
-use lambda::lambda;
+use lambda::{lambda, Context};
 use serde_json::Value;
 
 type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
 
 #[lambda]
 #[tokio::main]
-async fn main(event: Value) -> Result<Value, Error> {
+async fn main(event: Value, _: Context) -> Result<Value, Error> {
     Ok(event)
 }

lambda/examples/hello.rs

@@ -17,26 +17,23 @@
 //!    to the the `lambda::run` function, which launches and runs the Lambda runtime.
 //!
 //! An asynchronous function annotated with the `#[lambda]` attribute must
-//! accept an argument of type `A` which implements [`serde::Deserialize`] and
+//! accept an argument of type `A` which implements [`serde::Deserialize`], a [`lambda::Context`] and
 //! return a `Result<B, E>`, where `B` implements [`serde::Serializable`]. `E` is
 //! any type that implements `Into<Box<dyn std::error::Error + Send + Sync + 'static>>`.
 //!
-//! Optionally, the `#[lambda]` annotated function can accept an argument
-//! of [`lambda::LambdaCtx`].
-//!
 //! ```no_run
-//! use lambda::{lambda};
+//! use lambda::{lambda, Context};
 //! use serde_json::Value;
 //!
 //! type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
 //!
 //! #[lambda]
 //! #[tokio::main]
-//! async fn main(event: Value) -> Result<Value, Error> {
+//! async fn main(event: Value, _: Context) -> Result<Value, Error> {
 //!     Ok(event)
 //! }
 //! ```
-pub use crate::types::LambdaCtx;
+pub use crate::types::Context;
 use client::Client;
 use futures::stream::{Stream, StreamExt};
 use genawaiter::{sync::gen, yield_};
@@ -93,7 +90,7 @@ impl Config {
 }
 
 tokio::task_local! {
-    pub static INVOCATION_CTX: types::LambdaCtx;
+    pub static INVOCATION_CTX: types::Context;
 }
 
 /// A trait describing an asynchronous function `A` to `B.
@@ -102,8 +99,8 @@ pub trait Handler<A, B> {
     type Error;
     /// The future response value of this handler.
     type Fut: Future<Output = Result<B, Self::Error>>;
-    /// Process the incoming event and return the response asynchronously.
-    fn call(&mut self, event: A) -> Self::Fut;
+    /// Process the incoming event and `Context` then return the response asynchronously.
+    fn call(&mut self, event: A, context: Context) -> Self::Fut;
 }
 
 /// Returns a new `HandlerFn` with the given closure.
@@ -119,15 +116,14 @@ pub struct HandlerFn<F> {
 
 impl<F, A, B, Error, Fut> Handler<A, B> for HandlerFn<F>
 where
-    F: Fn(A) -> Fut,
+    F: Fn(A, Context) -> Fut,
     Fut: Future<Output = Result<B, Error>> + Send,
     Error: Into<Error> + fmt::Debug,
 {
     type Error = Error;
     type Fut = Fut;
-    fn call(&mut self, req: A) -> Self::Fut {
-        // we pass along the context here
-        (self.f)(req)
+    fn call(&mut self, req: A, ctx: Context) -> Self::Fut {
+        (self.f)(req, ctx)
     }
 }
 
@@ -136,7 +132,7 @@ where
 ///
 /// # Example
 /// ```no_run
-/// use lambda::{handler_fn};
+/// use lambda::{handler_fn, Context};
 /// use serde_json::Value;
 ///
 /// type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
@@ -148,7 +144,7 @@ where
 ///     Ok(())
 /// }
 ///
-/// async fn func(event: Value) -> Result<Value, Error> {
+/// async fn func(event: Value, _: Context) -> Result<Value, Error> {
 ///     Ok(event)
 /// }
 /// ```
@@ -211,12 +207,12 @@ where
         let event = event?;
         let (parts, body) = event.into_parts();
 
-        let ctx: LambdaCtx = LambdaCtx::try_from(parts.headers)?;
+        let ctx: Context = Context::try_from(parts.headers)?;
         let body = hyper::body::to_bytes(body).await?;
         let body = serde_json::from_slice(&body)?;
 
         let request_id = &ctx.request_id.clone();
-        let f = INVOCATION_CTX.scope(ctx, { handler.call(body) });
+        let f = INVOCATION_CTX.scope(ctx.clone(), handler.call(body, ctx));
 
         let req = match f.await {
             Ok(res) => EventCompletionRequest { request_id, body: res }.into_req()?,