cronback_lib/
ext.rs

1use tonic::metadata::MetadataMap;
2
3use crate::consts::{PROJECT_ID_HEADER, REQUEST_ID_HEADER};
4use crate::grpc_helpers::RequestContext;
5use crate::model::{ModelIdError, ValidShardedId};
6use crate::types::{ProjectId, RequestId};
7
8pub trait OptionExt {
9    type Item;
10    fn unwrap_ref(&self) -> &Self::Item;
11    fn unwrap_mut(&mut self) -> &mut Self::Item;
12}
13impl<T> OptionExt for Option<T> {
14    type Item = T;
15
16    fn unwrap_ref(&self) -> &T {
17        self.as_ref().unwrap()
18    }
19
20    fn unwrap_mut(&mut self) -> &mut T {
21        self.as_mut().unwrap()
22    }
23}
24
25pub trait GrpcMetadataMapExt {
26    fn set_project_id(&mut self, project_id: ValidShardedId<ProjectId>);
27    fn project_id(&self) -> Option<ProjectId>;
28
29    fn set_request_id(&mut self, request_id: RequestId);
30    fn request_id(&self) -> Option<RequestId>;
31}
32
33impl GrpcMetadataMapExt for MetadataMap {
34    fn set_project_id(&mut self, project_id: ValidShardedId<ProjectId>) {
35        self.insert(
36            PROJECT_ID_HEADER,
37            project_id
38                .to_string()
39                .parse()
40                .expect("ProjectId is not HTTP header-friendly!"),
41        );
42    }
43
44    fn set_request_id(&mut self, request_id: RequestId) {
45        self.insert(
46            REQUEST_ID_HEADER,
47            request_id
48                .to_string()
49                .parse()
50                .expect("RequestId is not HTTP header-friendly!"),
51        );
52    }
53
54    fn project_id(&self) -> Option<ProjectId> {
55        self.get(PROJECT_ID_HEADER)
56            .and_then(|project_id| project_id.to_str().ok())
57            .map(|project_id| ProjectId::from(project_id.to_string()))
58    }
59
60    fn request_id(&self) -> Option<RequestId> {
61        self.get(REQUEST_ID_HEADER)
62            .and_then(|request_id| request_id.to_str().ok())
63            .map(|request_id| RequestId::from(request_id.to_string()))
64    }
65}
66
67pub trait TonicRequestExt {
68    fn project_id(&self) -> Result<&ValidShardedId<ProjectId>, tonic::Status>;
69    fn request_id(&self) -> Result<&RequestId, tonic::Status>;
70    fn context(&self) -> Result<RequestContext, tonic::Status>;
71}
72
73impl<T> TonicRequestExt for tonic::Request<T> {
74    // ProjectId is added to extensions in CronbackRpcMiddleware
75    fn project_id(&self) -> Result<&ValidShardedId<ProjectId>, tonic::Status> {
76        let raw_id = self
77            .extensions()
78            .get::<Result<ValidShardedId<ProjectId>, ModelIdError>>()
79            .ok_or_else(|| {
80                tonic::Status::internal(format!(
81                    "{} was not set in GRPC headers!",
82                    PROJECT_ID_HEADER
83                ))
84            })?;
85        match raw_id {
86            | Ok(id) => Ok(id),
87            | Err(err) => {
88                Err(tonic::Status::internal(format!(
89                    "Invalid ProjectId was passed in headers: {}",
90                    err,
91                )))
92            }
93        }
94    }
95
96    fn request_id(&self) -> Result<&RequestId, tonic::Status> {
97        self.extensions().get::<RequestId>().ok_or_else(|| {
98            tonic::Status::internal(format!(
99                "{} was not set in GRPC headers!",
100                REQUEST_ID_HEADER,
101            ))
102        })
103    }
104
105    fn context(&self) -> Result<RequestContext, tonic::Status> {
106        Ok(RequestContext::new(
107            self.request_id()?.clone(),
108            self.project_id()?.clone(),
109        ))
110    }
111}
112
113// re-export Option extension for Prost lazy defaults
114pub use dto::traits::ProstOptionExt;
cronback_lib/ext.rs

cronback_lib/
ext.rs
