aws_smithy_http_server_python/middleware/
request.rs1use std::mem;
9use std::sync::Arc;
10
11use aws_smithy_http_server::body::Body;
12use http::{request::Parts, Request};
13use pyo3::{
14 exceptions::{PyRuntimeError, PyValueError},
15 prelude::*,
16};
17use tokio::sync::Mutex;
18
19use super::{PyHeaderMap, PyMiddlewareError};
20
21#[pyclass(name = "Request")]
23#[derive(Debug)]
24pub struct PyRequest {
25 parts: Option<Parts>,
26 headers: PyHeaderMap,
27 body: Arc<Mutex<Option<Body>>>,
28}
29
30impl PyRequest {
31 pub fn new(request: Request<Body>) -> Self {
33 let (mut parts, body) = request.into_parts();
34 let headers = mem::take(&mut parts.headers);
35 Self {
36 parts: Some(parts),
37 headers: PyHeaderMap::new(headers),
38 body: Arc::new(Mutex::new(Some(body))),
39 }
40 }
41
42 pub fn take_inner(&mut self) -> Option<Request<Body>> {
46 let headers = self.headers.take_inner()?;
47 let mut parts = self.parts.take()?;
48 parts.headers = headers;
49 let body = {
50 let body = mem::take(&mut self.body);
51 let body = Arc::try_unwrap(body).ok()?;
52 body.into_inner().take()?
53 };
54 Some(Request::from_parts(parts, body))
55 }
56}
57
58#[pymethods]
59impl PyRequest {
60 #[getter]
64 fn method(&self) -> PyResult<String> {
65 self.parts
66 .as_ref()
67 .map(|parts| parts.method.to_string())
68 .ok_or_else(|| PyMiddlewareError::RequestGone.into())
69 }
70
71 #[getter]
75 fn uri(&self) -> PyResult<String> {
76 self.parts
77 .as_ref()
78 .map(|parts| parts.uri.to_string())
79 .ok_or_else(|| PyMiddlewareError::RequestGone.into())
80 }
81
82 #[setter]
86 fn set_uri(&mut self, uri_str: String) -> PyResult<()> {
87 self.parts.as_mut().map_or_else(
88 || Err(PyMiddlewareError::RequestGone.into()),
89 |parts| {
90 parts.uri = uri_str.parse().map_err(|e: http::uri::InvalidUri| {
91 PyValueError::new_err(format!(
92 "URI `{}` cannot be parsed. Error: {}",
93 uri_str, e
94 ))
95 })?;
96 Ok(())
97 },
98 )
99 }
100
101 #[getter]
105 fn version(&self) -> PyResult<String> {
106 self.parts
107 .as_ref()
108 .map(|parts| format!("{:?}", parts.version))
109 .ok_or_else(|| PyMiddlewareError::RequestGone.into())
110 }
111
112 #[getter]
116 fn headers(&self) -> PyHeaderMap {
117 self.headers.clone()
118 }
119
120 #[getter]
125 fn body<'p>(&self, py: Python<'p>) -> PyResult<&'p PyAny> {
126 let body = self.body.clone();
127 pyo3_asyncio::tokio::future_into_py(py, async move {
128 let body = {
129 let mut body_guard = body.lock().await;
130 let body = body_guard.take().ok_or(PyMiddlewareError::RequestGone)?;
131 let body = hyper::body::to_bytes(body)
132 .await
133 .map_err(|err| PyRuntimeError::new_err(err.to_string()))?;
134 let buf = body.clone();
135 body_guard.replace(Body::from(body));
136 buf
137 };
138 Ok(body.to_vec())
140 })
141 }
142
143 #[setter]
145 fn set_body(&mut self, buf: &[u8]) {
146 self.body = Arc::new(Mutex::new(Some(Body::from(buf.to_owned()))));
147 }
148}