momento_functions/
response.rs

1use momento_functions_host::encoding::Encode;
2use momento_functions_host::http;
3use momento_functions_wit::function_web::exports::momento::functions::guest_function_web::Response;
4use std::error::Error;
5use std::fmt::{Debug, Display, Formatter};
6
7/// Values returned by a function implemented with the [crate::post!] macro must implement this trait.
8pub trait IntoWebResponse {
9    fn response(self) -> Response;
10}
11
12/// A WebError represents an error result produced by a function execution.
13/// Functionally, it is also just an HTTP response - however, this allows for writing
14/// functions with a return signature of `WebResult` if you are okay with all errors
15/// being converted to 500s and returned in the body.
16#[derive(Debug)]
17pub struct WebError {
18    source: Option<Box<dyn Error>>,
19    response: WebResponse,
20}
21
22impl WebError {
23    pub fn message(message: impl Into<String>) -> Self {
24        let message = message.into();
25        let response = WebResponse {
26            status: 500,
27            headers: vec![],
28            body: message.as_bytes().to_vec(),
29        };
30        Self {
31            source: None,
32            response,
33        }
34    }
35}
36
37impl<E: Error + 'static> From<E> for WebError {
38    fn from(e: E) -> Self {
39        let body = format!("An error occurred during function invocation: {e}");
40        Self {
41            source: Some(Box::new(e)),
42            response: WebResponse {
43                status: 500,
44                headers: vec![],
45                body: body.into_bytes(),
46            },
47        }
48    }
49}
50
51impl Display for WebError {
52    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
53        write!(f, "WebError(Source: {:?})", self.source)
54    }
55}
56
57/// A Result type for implementing functions. Allows you to use `?` within your function body
58/// to return a 500 with the error details.
59pub type WebResult<T> = Result<T, WebError>;
60
61impl<R> IntoWebResponse for Result<R, WebError>
62where
63    R: IntoWebResponse,
64{
65    fn response(self) -> Response {
66        match self {
67            Ok(r) => r.response(),
68            Err(e) => e.response.response(),
69        }
70    }
71}
72
73impl IntoWebResponse for http::Response {
74    fn response(self) -> Response {
75        Response {
76            status: self.status,
77            headers: self.headers.into_iter().map(Into::into).collect(),
78            body: self.body,
79        }
80    }
81}
82
83/// This represents a response from a web function.
84/// When constructed, it's a 200 response with no headers or body.
85/// You can set the status, headers, and body via [WebResponse::with_status], [WebResponse::with_headers],
86/// and [WebResponse::with_body] respectfully.
87#[derive(Debug)]
88pub struct WebResponse {
89    status: u16,
90    headers: Vec<(String, String)>,
91    body: Vec<u8>,
92}
93
94impl Default for WebResponse {
95    fn default() -> Self {
96        Self {
97            status: 200,
98            headers: vec![],
99            body: vec![],
100        }
101    }
102}
103
104impl WebResponse {
105    /// Creates a new default response.
106    pub fn new() -> Self {
107        Self::default()
108    }
109
110    /// Sets the response status.
111    pub fn with_status(mut self, status: u16) -> Self {
112        self.status = status;
113        self
114    }
115
116    /// Adds a header to the response.
117    pub fn header(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
118        self.headers.push((key.into(), value.into()));
119        self
120    }
121
122    /// Overrides the collection of headers for the response.
123    pub fn with_headers(mut self, headers: Vec<(String, String)>) -> Self {
124        self.headers = headers;
125        self
126    }
127
128    /// Sets the response body. If encoding the body fails, returns an error.
129    pub fn with_body<E: Encode>(mut self, body: E) -> Result<Self, E::Error> {
130        let body = body.try_serialize().map(Into::into)?;
131        self.body = body;
132        Ok(self)
133    }
134}
135
136impl IntoWebResponse for WebResponse {
137    fn response(self) -> Response {
138        Response {
139            status: self.status,
140            headers: self.headers.into_iter().map(Into::into).collect(),
141            body: self.body,
142        }
143    }
144}