feat(types): Add gRPC Richer Error Model support (QuotaFailure) (#1204)

flemosr · web-flow · commit 03b4735bb4ba · 2023-01-04T14:24:09.000-05:00
diff --git a/tonic-types/src/lib.rs b/tonic-types/src/lib.rs
@@ -46,7 +46,8 @@ pub use pb::Status;
 mod richer_error;
 
 pub use richer_error::{
-    BadRequest, DebugInfo, ErrorDetail, ErrorDetails, FieldViolation, RetryInfo, StatusExt,
+    BadRequest, DebugInfo, ErrorDetail, ErrorDetails, FieldViolation, QuotaFailure, QuotaViolation,
+    RetryInfo, StatusExt,
 };
 
 mod sealed {
diff --git a/tonic-types/src/richer_error/error_details/mod.rs b/tonic-types/src/richer_error/error_details/mod.rs
@@ -1,6 +1,8 @@
 use std::time;
 
-use super::std_messages::{BadRequest, DebugInfo, FieldViolation, RetryInfo};
+use super::std_messages::{
+    BadRequest, DebugInfo, FieldViolation, QuotaFailure, QuotaViolation, RetryInfo,
+};
 
 pub(crate) mod vec;
 
@@ -17,6 +19,9 @@ pub struct ErrorDetails {
     /// This field stores [`DebugInfo`] data, if any.
     pub(crate) debug_info: Option<DebugInfo>,
 
+    /// This field stores [`QuotaFailure`] data, if any.
+    pub(crate) quota_failure: Option<QuotaFailure>,
+
     /// This field stores [`BadRequest`] data, if any.
     pub(crate) bad_request: Option<BadRequest>,
 }
@@ -35,6 +40,7 @@ impl ErrorDetails {
         ErrorDetails {
             retry_info: None,
             debug_info: None,
+            quota_failure: None,
             bad_request: None,
         }
     }
@@ -76,6 +82,46 @@ impl ErrorDetails {
         }
     }
 
+    /// Generates an [`ErrorDetails`] struct with [`QuotaFailure`] details and
+    /// remaining fields set to `None`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use tonic_types::{ErrorDetails, QuotaViolation};
+    ///
+    /// let err_details = ErrorDetails::with_quota_failure(vec![
+    ///     QuotaViolation::new("subject 1", "description 1"),
+    ///     QuotaViolation::new("subject 2", "description 2"),
+    /// ]);
+    /// ```
+    pub fn with_quota_failure(violations: Vec<QuotaViolation>) -> Self {
+        ErrorDetails {
+            quota_failure: Some(QuotaFailure::new(violations)),
+            ..ErrorDetails::new()
+        }
+    }
+
+    /// Generates an [`ErrorDetails`] struct with [`QuotaFailure`] details (one
+    /// [`QuotaViolation`] set) and remaining fields set to `None`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use tonic_types::{ErrorDetails};
+    ///
+    /// let err_details = ErrorDetails::with_quota_failure_violation("subject", "description");
+    /// ```
+    pub fn with_quota_failure_violation(
+        subject: impl Into<String>,
+        description: impl Into<String>,
+    ) -> Self {
+        ErrorDetails {
+            quota_failure: Some(QuotaFailure::with_violation(subject, description)),
+            ..ErrorDetails::new()
+        }
+    }
+
     /// Generates an [`ErrorDetails`] struct with [`BadRequest`] details and
     /// remaining fields set to `None`.
     ///
@@ -129,6 +175,11 @@ impl ErrorDetails {
         self.debug_info.clone()
     }
 
+    /// Get [`QuotaFailure`] details, if any
+    pub fn quota_failure(&self) -> Option<QuotaFailure> {
+        self.quota_failure.clone()
+    }
+
     /// Get [`BadRequest`] details, if any
     pub fn bad_request(&self) -> Option<BadRequest> {
         self.bad_request.clone()
@@ -175,6 +226,78 @@ impl ErrorDetails {
         self
     }
 
+    /// Set [`QuotaFailure`] details. Can be chained with other `.set_` and
+    /// `.add_` [`ErrorDetails`] methods.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use tonic_types::{ErrorDetails, QuotaViolation};
+    ///
+    /// let mut err_details = ErrorDetails::new();
+    ///
+    /// err_details.set_quota_failure(vec![
+    ///     QuotaViolation::new("subject 1", "description 1"),
+    ///     QuotaViolation::new("subject 2", "description 2"),
+    /// ]);
+    /// ```
+    pub fn set_quota_failure(&mut self, violations: Vec<QuotaViolation>) -> &mut Self {
+        self.quota_failure = Some(QuotaFailure::new(violations));
+        self
+    }
+
+    /// Adds a [`QuotaViolation`] to [`QuotaFailure`] details. Sets
+    /// [`QuotaFailure`] details if it is not set yet. Can be chained with
+    /// other `.set_` and `.add_` [`ErrorDetails`] methods.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use tonic_types::{ErrorDetails};
+    ///
+    /// let mut err_details = ErrorDetails::new();
+    ///
+    /// err_details.add_quota_failure_violation("subject", "description");
+    /// ```
+    pub fn add_quota_failure_violation(
+        &mut self,
+        subject: impl Into<String>,
+        description: impl Into<String>,
+    ) -> &mut Self {
+        match &mut self.quota_failure {
+            Some(quota_failure) => {
+                quota_failure.add_violation(subject, description);
+            }
+            None => {
+                self.quota_failure = Some(QuotaFailure::with_violation(subject, description));
+            }
+        };
+        self
+    }
+
+    /// Returns `true` if [`QuotaFailure`] is set and its `violations` vector
+    /// is not empty, otherwise returns `false`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use tonic_types::{ErrorDetails};
+    ///
+    /// let mut err_details = ErrorDetails::with_quota_failure(vec![]);
+    ///
+    /// assert_eq!(err_details.has_quota_failure_violations(), false);
+    ///
+    /// err_details.add_quota_failure_violation("subject", "description");
+    ///
+    /// assert_eq!(err_details.has_quota_failure_violations(), true);
+    /// ```
+    pub fn has_quota_failure_violations(&self) -> bool {
+        if let Some(quota_failure) = &self.quota_failure {
+            return !quota_failure.violations.is_empty();
+        }
+        false
+    }
+
     /// Set [`BadRequest`] details. Can be chained with other `.set_` and
     /// `.add_` [`ErrorDetails`] methods.
     ///
diff --git a/tonic-types/src/richer_error/error_details/vec.rs b/tonic-types/src/richer_error/error_details/vec.rs
@@ -1,4 +1,4 @@
-use super::super::std_messages::{BadRequest, DebugInfo, RetryInfo};
+use super::super::std_messages::{BadRequest, DebugInfo, QuotaFailure, RetryInfo};
 
 /// Wraps the structs corresponding to the standard error messages, allowing
 /// the implementation and handling of vectors containing any of them.
@@ -11,6 +11,9 @@ pub enum ErrorDetail {
     /// Wraps the [`DebugInfo`] struct.
     DebugInfo(DebugInfo),
 
+    /// Wraps the [`QuotaFailure`] struct.
+    QuotaFailure(QuotaFailure),
+
     /// Wraps the [`BadRequest`] struct.
     BadRequest(BadRequest),
 }
@@ -27,6 +30,12 @@ impl From<DebugInfo> for ErrorDetail {
     }
 }
 
+impl From<QuotaFailure> for ErrorDetail {
+    fn from(err_detail: QuotaFailure) -> Self {
+        ErrorDetail::QuotaFailure(err_detail)
+    }
+}
+
 impl From<BadRequest> for ErrorDetail {
     fn from(err_detail: BadRequest) -> Self {
         ErrorDetail::BadRequest(err_detail)
diff --git a/tonic-types/src/richer_error/mod.rs b/tonic-types/src/richer_error/mod.rs
@@ -8,7 +8,9 @@ mod std_messages;
 use super::pb;
 
 pub use error_details::{vec::ErrorDetail, ErrorDetails};
-pub use std_messages::{BadRequest, DebugInfo, FieldViolation, RetryInfo};
+pub use std_messages::{
+    BadRequest, DebugInfo, FieldViolation, QuotaFailure, QuotaViolation, RetryInfo,
+};
 
 trait IntoAny {
     fn into_any(self) -> Any;
@@ -288,6 +290,28 @@ pub trait StatusExt: crate::sealed::Sealed {
     /// ```
     fn get_details_debug_info(&self) -> Option<DebugInfo>;
 
+    /// Get first [`QuotaFailure`] details found on `tonic::Status`, if any.
+    /// If some `prost::DecodeError` occurs, returns `None`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use tonic::{Status, Response};
+    /// use tonic_types::{StatusExt};
+    ///
+    /// fn handle_request_result<T>(req_result: Result<Response<T>, Status>) {
+    ///     match req_result {
+    ///         Ok(_) => {},
+    ///         Err(status) => {
+    ///             if let Some(quota_failure) = status.get_details_quota_failure() {
+    ///                 // Handle quota_failure details
+    ///             }
+    ///         }
+    ///     };
+    /// }
+    /// ```
+    fn get_details_quota_failure(&self) -> Option<QuotaFailure>;
+
     /// Get first [`BadRequest`] details found on `tonic::Status`, if any. If
     /// some `prost::DecodeError` occurs, returns `None`.
     ///
@@ -332,6 +356,10 @@ impl StatusExt for tonic::Status {
             conv_details.push(debug_info.into_any());
         }
 
+        if let Some(quota_failure) = details.quota_failure {
+            conv_details.push(quota_failure.into_any());
+        }
+
         if let Some(bad_request) = details.bad_request {
             conv_details.push(bad_request.into_any());
         }
@@ -363,6 +391,9 @@ impl StatusExt for tonic::Status {
                 ErrorDetail::DebugInfo(debug_info) => {
                     conv_details.push(debug_info.into_any());
                 }
+                ErrorDetail::QuotaFailure(quota_failure) => {
+                    conv_details.push(quota_failure.into_any());
+                }
                 ErrorDetail::BadRequest(bad_req) => {
                     conv_details.push(bad_req.into_any());
                 }
@@ -400,6 +431,9 @@ impl StatusExt for tonic::Status {
                 DebugInfo::TYPE_URL => {
                     details.debug_info = Some(DebugInfo::from_any(any)?);
                 }
+                QuotaFailure::TYPE_URL => {
+                    details.quota_failure = Some(QuotaFailure::from_any(any)?);
+                }
                 BadRequest::TYPE_URL => {
                     details.bad_request = Some(BadRequest::from_any(any)?);
                 }
@@ -427,6 +461,9 @@ impl StatusExt for tonic::Status {
                 DebugInfo::TYPE_URL => {
                     details.push(DebugInfo::from_any(any)?.into());
                 }
+                QuotaFailure::TYPE_URL => {
+                    details.push(QuotaFailure::from_any(any)?.into());
+                }
                 BadRequest::TYPE_URL => {
                     details.push(BadRequest::from_any(any)?.into());
                 }
@@ -469,6 +506,20 @@ impl StatusExt for tonic::Status {
         None
     }
 
+    fn get_details_quota_failure(&self) -> Option<QuotaFailure> {
+        let status = pb::Status::decode(self.details()).ok()?;
+
+        for any in status.details.into_iter() {
+            if any.type_url.as_str() == QuotaFailure::TYPE_URL {
+                if let Ok(detail) = QuotaFailure::from_any(any) {
+                    return Some(detail);
+                }
+            }
+        }
+
+        None
+    }
+
     fn get_details_bad_request(&self) -> Option<BadRequest> {
         let status = pb::Status::decode(self.details()).ok()?;
 
@@ -489,7 +540,7 @@ mod tests {
     use std::time::Duration;
     use tonic::{Code, Status};
 
-    use super::{BadRequest, DebugInfo, ErrorDetails, RetryInfo, StatusExt};
+    use super::{BadRequest, DebugInfo, ErrorDetails, QuotaFailure, RetryInfo, StatusExt};
 
     #[test]
     fn gen_status_with_details() {
@@ -501,6 +552,7 @@ mod tests {
                 vec!["trace3".into(), "trace2".into(), "trace1".into()],
                 "details",
             )
+            .add_quota_failure_violation("clientip:<ip address>", "description")
             .add_bad_request_violation("field", "description");
 
         let fmt_details = format!("{:?}", err_details);
@@ -512,6 +564,7 @@ mod tests {
                 "details",
             )
             .into(),
+            QuotaFailure::with_violation("clientip:<ip address>", "description").into(),
             BadRequest::with_violation("field", "description").into(),
         ];
 
diff --git a/tonic-types/src/richer_error/std_messages/mod.rs b/tonic-types/src/richer_error/std_messages/mod.rs
@@ -6,6 +6,10 @@ mod debug_info;
 
 pub use debug_info::DebugInfo;
 
+mod quota_failure;
+
+pub use quota_failure::{QuotaFailure, QuotaViolation};
+
 mod bad_request;
 
 pub use bad_request::{BadRequest, FieldViolation};
diff --git a/tonic-types/src/richer_error/std_messages/quota_failure.rs b/tonic-types/src/richer_error/std_messages/quota_failure.rs
