fix: handle when fetch_one fails

burrito-project · Oct 7, 2024 · 1587c90 · 1587c90
1 parent b580ef2
commit 1587c90
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 10 deletions.
diff --git a/src/features/flags/handlers/get_flag.rs b/src/features/flags/handlers/get_flag.rs
@@ -1,8 +1,10 @@
+use rocket::http::Status;
+use rocket::response::status;
 use rocket::State;
 use serde_json::json;
 
 use crate::core::types::ApiResponse;
-use crate::core::AppState;
+use crate::core::{responses, AppState};
 use crate::features::flags;
 
 pub async fn get_flag_handler(flag: &str, state: &State<AppState>) -> ApiResponse {
@@ -11,9 +13,16 @@ pub async fn get_flag_handler(flag: &str, state: &State<AppState>) -> ApiRespons
         "SELECT * FROM flags WHERE name = $1 AND internal = false LIMIT 1;",
         flag,
     )
-    .fetch_one(&state.pool)
+    .fetch_optional(&state.pool)
     .await
     .unwrap();
 
-    Ok(json!(flag.value))
+    if flag.is_none() {
+        return Err(status::Custom(
+            Status::NotFound,
+            responses::error_response("Flag not found".to_string()),
+        ));
+    }
+
+    Ok(json!(flag.unwrap().value))
 }
diff --git a/src/features/notifications/handlers/delete_notification.rs b/src/features/notifications/handlers/delete_notification.rs
@@ -1,8 +1,10 @@
+use rocket::http::Status;
+use rocket::response::status;
 use rocket::State;
 use serde_json::json;
 
 use crate::core::types::ApiResponse;
-use crate::core::AppState;
+use crate::core::{responses, AppState};
 use crate::features::notifications::schemas;
 
 pub async fn delete_notification_handler(id: i32, state: &State<AppState>) -> ApiResponse {
@@ -11,9 +13,16 @@ pub async fn delete_notification_handler(id: i32, state: &State<AppState>) -> Ap
         "DELETE FROM notification_ads WHERE id = $1 RETURNING *;",
         id,
     )
-    .fetch_one(&state.pool)
+    .fetch_optional(&state.pool)
     .await
     .unwrap();
 
+    if deleted_notification.is_none() {
+        return Err(status::Custom(
+            Status::NotFound,
+            responses::error_response("Notification not found".to_string()),
+        ));
+    }
+
     Ok(json!(deleted_notification))
 }
diff --git a/src/features/updates/handlers/delete_app_version.rs b/src/features/updates/handlers/delete_app_version.rs
@@ -1,8 +1,10 @@
+use rocket::http::Status;
+use rocket::response::status;
 use rocket::State;
 use serde_json::json;
 
 use crate::core::types::ApiResponse;
-use crate::core::AppState;
+use crate::core::{responses, AppState};
 use crate::features::updates::schemas;
 
 pub async fn delete_app_version_handler(id: i32, state: &State<AppState>) -> ApiResponse {
@@ -11,9 +13,16 @@ pub async fn delete_app_version_handler(id: i32, state: &State<AppState>) -> Api
         "DELETE FROM app_versions WHERE id = $1 RETURNING *;",
         id,
     )
-    .fetch_one(&state.pool)
+    .fetch_optional(&state.pool)
     .await
     .unwrap();
 
-    Ok(json!(deleted_version))
+    if deleted_version.is_none() {
+        return Err(status::Custom(
+            Status::NotFound,
+            responses::error_response("Flag not found".to_string()),
+        ));
+    }
+
+    Ok(json!(deleted_version.unwrap()))
 }
diff --git a/src/features/updates/handlers/patch_app_version.rs b/src/features/updates/handlers/patch_app_version.rs
@@ -1,8 +1,10 @@
+use rocket::http::Status;
+use rocket::response::status;
 use rocket::State;
 use serde_json::json;
 
 use crate::core::types::ApiResponse;
-use crate::core::AppState;
+use crate::core::{responses, AppState};
 use crate::features::updates::schemas;
 
 pub async fn patch_app_version_handler(
@@ -33,7 +35,14 @@ pub async fn patch_app_version_handler(
         payload.release_notes,
     );
 
-    let updated_version = updated_version.fetch_one(&state.pool).await.unwrap();
+    let updated_version = updated_version.fetch_optional(&state.pool).await.unwrap();
+
+    if updated_version.is_none() {
+        return Err(status::Custom(
+            Status::NotFound,
+            responses::error_response("App version does not exist!"),
+        ));
+    }
 
     Ok(json!(updated_version))
 }