Add error kind (#590)

jamsocket · Jan 31, 2024 · 7041219 · 7041219
1 parent 389e40c
commit 7041219
Show file tree

Hide file tree

Showing 5 changed files with 73 additions and 21 deletions.
diff --git a/plane/src/controller/connect.rs b/plane/src/controller/connect.rs
@@ -1,4 +1,4 @@
-use super::error::err_to_response;
+use super::error::{err_to_response, ApiErrorKind};
 use super::Controller;
 use crate::database::connect::ConnectError;
 use crate::types::{ConnectRequest, ConnectResponse};
@@ -11,46 +11,55 @@ fn connect_error_to_response(connect_error: &ConnectError) -> Response {
             connect_error,
             StatusCode::INTERNAL_SERVER_ERROR,
             "Failed to acquire lock.",
+            ApiErrorKind::FailedToAcquireKey,
         ),
         ConnectError::KeyUnheldNoSpawnConfig => err_to_response(
             connect_error,
             StatusCode::CONFLICT,
             "Lock is unheld but no spawn config was provided.",
+            ApiErrorKind::KeyUnheldNoSpawnConfig,
         ),
         ConnectError::KeyHeldUnhealthy => err_to_response(
             connect_error,
             StatusCode::INTERNAL_SERVER_ERROR,
             "Lock is held but unhealthy.",
+            ApiErrorKind::KeyHeldUnhealthy,
         ),
         ConnectError::KeyHeld { .. } => err_to_response(
             connect_error,
             StatusCode::CONFLICT,
             "Lock is held but tag does not match.",
+            ApiErrorKind::KeyHeld,
         ),
         ConnectError::NoDroneAvailable => err_to_response(
             connect_error,
             StatusCode::INTERNAL_SERVER_ERROR,
             "No active drone available.",
+            ApiErrorKind::NoDroneAvailable,
         ),
         ConnectError::FailedToRemoveKey => err_to_response(
             connect_error,
             StatusCode::CONFLICT,
             "Failed to remove lock.",
+            ApiErrorKind::FailedToRemoveKey,
         ),
-        ConnectError::Sql(_) | ConnectError::Serialization(_) => err_to_response(
+        ConnectError::DatabaseError(_) | ConnectError::Serialization(_) => err_to_response(
             connect_error,
             StatusCode::INTERNAL_SERVER_ERROR,
             "Internal error.",
+            ApiErrorKind::Other,
         ),
         ConnectError::NoClusterProvided => err_to_response(
             connect_error,
             StatusCode::BAD_REQUEST,
             "No cluster provided, and no default cluster for this controller.",
+            ApiErrorKind::NoClusterProvided,
         ),
         ConnectError::Other(_) => err_to_response(
             connect_error,
             StatusCode::INTERNAL_SERVER_ERROR,
             "Internal error.",
+            ApiErrorKind::Other,
         ),
     }
 }

diff --git a/plane/src/controller/drone.rs b/plane/src/controller/drone.rs
@@ -1,4 +1,4 @@
-use super::Controller;
+use super::{error::ApiErrorKind, Controller};
 use crate::{
     controller::error::IntoApiError,
     database::{
@@ -214,10 +214,11 @@ pub async fn handle_drone_socket(
     connect_info: ConnectInfo<SocketAddr>,
     ws: WebSocketUpgrade,
 ) -> Result<impl IntoResponse, Response> {
-    let cluster: ClusterName = cluster
-        .parse()
-        .ok()
-        .or_status(StatusCode::BAD_REQUEST, "Invalid cluster name")?;
+    let cluster: ClusterName = cluster.parse().ok().or_status(
+        StatusCode::BAD_REQUEST,
+        "Invalid cluster name",
+        ApiErrorKind::InvalidClusterName,
+    )?;
     let ip = connect_info.0.ip();
     Ok(ws.on_upgrade(move |socket| drone_socket(cluster, socket, controller, ip)))
 }
diff --git a/plane/src/controller/error.rs b/plane/src/controller/error.rs
@@ -10,9 +10,25 @@ use std::{
     fmt::{Debug, Display},
 };
 
+#[derive(Serialize, Deserialize, Debug, Clone, Copy)]
+pub enum ApiErrorKind {
+    FailedToAcquireKey,
+    KeyUnheldNoSpawnConfig,
+    KeyHeldUnhealthy,
+    KeyHeld,
+    NoDroneAvailable,
+    FailedToRemoveKey,
+    DatabaseError,
+    NoClusterProvided,
+    NotFound,
+    InvalidClusterName,
+    Other,
+}
+
 #[derive(thiserror::Error, Debug, Serialize, Deserialize)]
 pub struct ApiError {
     pub id: String,
+    pub kind: ApiErrorKind,
     pub message: String,
 }
 
@@ -22,7 +38,12 @@ impl Display for ApiError {
     }
 }
 
-pub fn err_to_response<E: Debug>(error: E, status: StatusCode, user_message: &str) -> Response {
+pub fn err_to_response<E: Debug>(
+    error: E,
+    status: StatusCode,
+    user_message: &str,
+    code: ApiErrorKind,
+) -> Response {
     let err_id = random_string();
 
     if status.is_server_error() {
@@ -42,41 +63,61 @@ pub fn err_to_response<E: Debug>(error: E, status: StatusCode, user_message: &st
     let result = ApiError {
         id: err_id.clone(),
         message: user_message.to_string(),
+        kind: code,
     };
 
     (status, Json(result)).into_response()
 }
 
 pub trait IntoApiError<T>: Sized {
     fn or_not_found(self, user_message: &str) -> Result<T, Response> {
-        self.or_status(StatusCode::NOT_FOUND, user_message)
+        self.or_status(StatusCode::NOT_FOUND, user_message, ApiErrorKind::NotFound)
     }
 
     fn or_internal_error(self, user_message: &str) -> Result<T, Response> {
-        self.or_status(StatusCode::INTERNAL_SERVER_ERROR, user_message)
+        self.or_status(
+            StatusCode::INTERNAL_SERVER_ERROR,
+            user_message,
+            ApiErrorKind::Other,
+        )
     }
 
-    fn or_status(self, status: StatusCode, user_message: &str) -> Result<T, Response>;
+    fn or_status(
+        self,
+        status: StatusCode,
+        user_message: &str,
+        code: ApiErrorKind,
+    ) -> Result<T, Response>;
 }
 
 impl<T, E: Error> IntoApiError<T> for Result<T, E> {
-    fn or_status(self, status: StatusCode, user_message: &str) -> Result<T, Response> {
+    fn or_status(
+        self,
+        status: StatusCode,
+        user_message: &str,
+        code: ApiErrorKind,
+    ) -> Result<T, Response> {
         match self {
             Ok(v) => Ok(v),
             Err(error) => {
-                let err = err_to_response(&error, status, user_message);
+                let err = err_to_response(&error, status, user_message, code);
                 Err(err)
             }
         }
     }
 }
 
 impl<T> IntoApiError<T> for Option<T> {
-    fn or_status(self, status: StatusCode, user_message: &str) -> Result<T, Response> {
+    fn or_status(
+        self,
+        status: StatusCode,
+        user_message: &str,
+        code: ApiErrorKind,
+    ) -> Result<T, Response> {
         match self {
             Some(v) => Ok(v),
             None => {
-                let err = err_to_response("Missing value.", status, user_message);
+                let err = err_to_response("Missing value.", status, user_message, code);
                 Err(err)
             }
         }

diff --git a/plane/src/controller/proxy.rs b/plane/src/controller/proxy.rs
@@ -1,4 +1,4 @@
-use super::Controller;
+use super::{error::ApiErrorKind, Controller};
 use crate::{
     controller::error::IntoApiError,
     protocol::{
@@ -132,10 +132,11 @@ pub async fn handle_proxy_socket(
     connect_info: ConnectInfo<SocketAddr>,
     ws: WebSocketUpgrade,
 ) -> Result<impl IntoResponse, Response> {
-    let cluster: ClusterName = cluster
-        .parse()
-        .ok()
-        .or_status(StatusCode::BAD_REQUEST, "Invalid cluster name")?;
+    let cluster: ClusterName = cluster.parse().ok().or_status(
+        StatusCode::BAD_REQUEST,
+        "Invalid cluster name",
+        ApiErrorKind::InvalidClusterName,
+    )?;
     let ip = connect_info.ip();
     Ok(ws.on_upgrade(move |socket| proxy_socket(cluster, socket, controller, ip)))
 }
diff --git a/plane/src/database/connect.rs b/plane/src/database/connect.rs
@@ -50,7 +50,7 @@ pub enum ConnectError {
     FailedToAcquireKey,
 
     #[error("SQL error: {0}")]
-    Sql(#[from] sqlx::Error),
+    DatabaseError(#[from] sqlx::Error),
 
     #[error("Serialization error: {0}")]
     Serialization(#[from] serde_json::Error),