poem_lambda/
lib.rs

1//! Poem for AWS Lambda.
2
3#![doc(html_favicon_url = "https://raw.githubusercontent.com/poem-web/poem/master/favicon.ico")]
4#![doc(html_logo_url = "https://raw.githubusercontent.com/poem-web/poem/master/logo.png")]
5#![forbid(unsafe_code)]
6#![deny(unreachable_pub)]
7#![cfg_attr(docsrs, feature(doc_cfg))]
8#![warn(missing_docs)]
9
10use std::{io::ErrorKind, ops::Deref, sync::Arc};
11
12pub use lambda_http::lambda_runtime::Error;
13use lambda_http::{
14    lambda_runtime, service_fn, Body as LambdaBody, Request as LambdaRequest, RequestExt,
15};
16use poem::{Body, Endpoint, EndpointExt, FromRequest, IntoEndpoint, Request, RequestBody, Result};
17
18/// The Lambda function execution context.
19///
20/// It implements [`poem::FromRequest`], so it can be used as an extractor.
21///
22/// # Example
23///
24/// ```
25/// use poem::{handler, Request};
26/// use poem_lambda::Context;
27///
28/// #[handler]
29/// fn index(req: &Request, ctx: &Context) {
30///     println!("request_id: {}", ctx.request_id);
31/// }
32/// ```
33#[derive(Debug, Clone)]
34pub struct Context(pub lambda_runtime::Context);
35
36impl Deref for Context {
37    type Target = lambda_runtime::Context;
38
39    fn deref(&self) -> &Self::Target {
40        &self.0
41    }
42}
43
44/// Starts the AWS Lambda runtime.
45///
46/// # Example
47///
48/// ```no_run
49/// use poem::handler;
50/// use poem_lambda::Error;
51///
52/// #[handler]
53/// fn index() -> &'static str {
54///     "hello"
55/// }
56///
57/// #[tokio::main]
58/// async fn main() -> Result<(), Error> {
59///     poem_lambda::run(index).await
60/// }
61/// ```
62pub async fn run(ep: impl IntoEndpoint) -> Result<(), Error> {
63    let ep = Arc::new(ep.map_to_response().into_endpoint());
64    lambda_http::run(service_fn(move |req: LambdaRequest| {
65        let ctx = req.lambda_context();
66        let ep = ep.clone();
67        async move {
68            let mut req: Request = from_lambda_request(req);
69            req.extensions_mut().insert(Context(ctx));
70
71            let resp = ep.get_response(req).await;
72            let (parts, body) = resp.into_parts();
73            let data = body
74                .into_vec()
75                .await
76                .map_err(|_| std::io::Error::new(ErrorKind::Other, "invalid request"))?;
77            let mut lambda_resp = poem::http::Response::new(if data.is_empty() {
78                LambdaBody::Empty
79            } else {
80                match String::from_utf8(data) {
81                    Ok(data) => LambdaBody::Text(data),
82                    Err(err) => LambdaBody::Binary(err.into_bytes()),
83                }
84            });
85            *lambda_resp.status_mut() = parts.status;
86            *lambda_resp.version_mut() = parts.version;
87            *lambda_resp.headers_mut() = parts.headers;
88            *lambda_resp.extensions_mut() = parts.extensions;
89
90            Ok::<_, Error>(lambda_resp)
91        }
92    }))
93    .await
94}
95
96fn from_lambda_request(req: LambdaRequest) -> Request {
97    let (parts, lambda_body) = req.into_parts();
98    let body = match lambda_body {
99        LambdaBody::Empty => Body::empty(),
100        LambdaBody::Text(data) => Body::from_string(data),
101        LambdaBody::Binary(data) => Body::from_vec(data),
102    };
103    let mut req = Request::builder()
104        .method(parts.method)
105        .uri(parts.uri)
106        .version(parts.version)
107        .body(body);
108    *req.headers_mut() = parts.headers;
109    *req.extensions_mut() = parts.extensions;
110    req
111}
112
113impl<'a> FromRequest<'a> for &'a Context {
114    async fn from_request(req: &'a Request, _body: &mut RequestBody) -> Result<Self> {
115        let ctx = match req.extensions().get::<Context>() {
116            Some(ctx) => ctx,
117            None => panic!("Lambda runtime is required."),
118        };
119        Ok(ctx)
120    }
121}
poem_lambda/lib.rs

poem_lambda/
lib.rs
