perf: Improve error handling for count method

zaunchat · Jul 15, 2022 · 73b37e0 · 73b37e0
1 parent ee88e4b
commit 73b37e0
Show file tree

Hide file tree

Showing 6 changed files with 13 additions and 14 deletions.
diff --git a/src/middlewares/member.rs b/src/middlewares/member.rs
@@ -22,7 +22,7 @@ pub async fn handle<B: std::marker::Send>(
     let Path(ID { server_id, .. }) = req.extract::<Path<ID>>().await.unwrap();
     let user = req.extensions().get::<User>().unwrap();
 
-    let exists = Member::count(&format!("id = {} AND server_id = {}", user.id, server_id)).await;
+    let exists = Member::count(&format!("id = {} AND server_id = {}", user.id, server_id)).await?;
 
     if exists == 0 {
         return Err(Error::UnknownServer);

diff --git a/src/routes/invites/join.rs b/src/routes/invites/join.rs
@@ -23,9 +23,9 @@ pub async fn join(Extension(user): Extension<User>, Path(code): Path<String>) ->
                 return Err(Error::MissingAccess);
             }
 
-            let count = Member::count(&format!("server_id = {}", server_id)).await;
+            let count = Member::count(&format!("server_id = {}", server_id)).await?;
 
-            if count > *MAX_SERVER_MEMBERS {
+            if count >= *MAX_SERVER_MEMBERS {
                 return Err(Error::MaximumChannels);
             }
 

diff --git a/src/routes/servers/channels/create.rs b/src/routes/servers/channels/create.rs
@@ -22,9 +22,9 @@ pub async fn create(
         .await?
         .has(Permissions::MANAGE_CHANNELS)?;
 
-    let count = Channel::count(&format!("server_id = {}", server_id)).await;
+    let count = Channel::count(&format!("server_id = {}", server_id)).await?;
 
-    if count > *MAX_SERVER_CHANNELS {
+    if count >= *MAX_SERVER_CHANNELS {
         return Err(Error::MaximumChannels);
     }
 

diff --git a/src/routes/servers/create.rs b/src/routes/servers/create.rs
@@ -16,9 +16,9 @@ pub async fn create(
     Extension(user): Extension<User>,
     ValidatedJson(data): ValidatedJson<CreateServerOptions>,
 ) -> Result<Json<Server>> {
-    let count = Member::count(&format!("id = {}", user.id)).await;
+    let count = Member::count(&format!("id = {}", user.id)).await?;
 
-    if count > *MAX_SERVERS {
+    if count >= *MAX_SERVERS {
         return Err(Error::MaximumServers);
     }
 

diff --git a/src/routes/servers/roles/create.rs b/src/routes/servers/roles/create.rs
@@ -24,9 +24,9 @@ pub async fn create(
         .await?
         .has(Permissions::MANAGE_ROLES)?;
 
-    let count = Role::count(&format!("server_id = {}", server_id)).await;
+    let count = Role::count(&format!("server_id = {}", server_id)).await?;
 
-    if count > *MAX_SERVER_ROLES {
+    if count >= *MAX_SERVER_ROLES {
         return Err(Error::MaximumRoles);
     }
 

diff --git a/src/structures/base.rs b/src/structures/base.rs
@@ -4,16 +4,15 @@ use sqlx::{Encode, Postgres, Type};
 
 #[async_trait]
 pub trait Base: Model<Postgres> {
-    async fn count(filter: &str) -> u64 {
-        sqlx::query(&format!(
+    async fn count(filter: &str) -> Result<u64, sqlx::Error> {
+        Ok(sqlx::query(&format!(
             "SELECT COUNT(*) FROM {} WHERE {}",
             Self::table_name(),
             filter
         ))
         .execute(pool())
-        .await
-        .unwrap()
-        .rows_affected()
+        .await?
+        .rows_affected())
     }
 
     async fn save(self) -> Result<Self, ormlite::Error> {