From 868ec66b6826ae15aa46d7d7a97c0a7d30c18f27 Mon Sep 17 00:00:00 2001
From: Awiteb <a@4rs.nl>
Date: Wed, 17 Jul 2024 23:22:24 +0300
Subject: [PATCH] feat: Create unique index for `UserId` and `Target`

Signed-off-by: Awiteb <a@4rs.nl>
---
 crates/oxidetalis/src/database/user_status.rs | 83 ++++++++++++-------
 .../src/create_users_status.rs                | 12 +++
 2 files changed, 64 insertions(+), 31 deletions(-)

diff --git a/crates/oxidetalis/src/database/user_status.rs b/crates/oxidetalis/src/database/user_status.rs
index ab123f2..24d96ae 100644
--- a/crates/oxidetalis/src/database/user_status.rs
+++ b/crates/oxidetalis/src/database/user_status.rs
@@ -139,30 +139,38 @@ impl UsersStatusExt for DatabaseConnection {
         whitelister: &UserModel,
         target_public_key: &PublicKey,
     ) -> ServerResult<()> {
-        if self.is_whitelisted(whitelister, target_public_key).await? {
-            return Err(WsError::AlreadyOnTheWhitelist.into());
-        }
         if whitelister.public_key == target_public_key.to_string() {
             return Err(WsError::CannotAddSelfToWhitelist.into());
         }
 
-        if let Some(mut user) = get_user_status(self, whitelister, target_public_key)
-            .await?
-            .map(IntoActiveModel::into_active_model)
+        if let Some(mut user) = get_user_status(
+            self,
+            whitelister,
+            target_public_key,
+            AccessStatus::Blacklisted,
+        )
+        .await?
+        .map(IntoActiveModel::into_active_model)
         {
             user.status = Set(AccessStatus::Whitelisted);
             user.updated_at = Set(Utc::now());
             user.update(self).await?;
-        } else {
-            UsersStatusActiveModel {
-                user_id: Set(whitelister.id),
-                target: Set(target_public_key.to_string()),
-                status: Set(AccessStatus::Whitelisted),
-                updated_at: Set(Utc::now()),
-                ..Default::default()
+        } else if let Err(err) = (UsersStatusActiveModel {
+            user_id: Set(whitelister.id),
+            target: Set(target_public_key.to_string()),
+            status: Set(AccessStatus::Whitelisted),
+            updated_at: Set(Utc::now()),
+            ..Default::default()
+        }
+        .save(self)
+        .await)
+        {
+            match err.sql_err() {
+                Some(SqlErr::UniqueConstraintViolation(_)) => {
+                    return Err(WsError::AlreadyOnTheWhitelist.into());
+                }
+                _ => return Err(err.into()),
             }
-            .save(self)
-            .await?;
         }
 
         Ok(())
@@ -173,30 +181,38 @@ impl UsersStatusExt for DatabaseConnection {
         blacklister: &UserModel,
         target_public_key: &PublicKey,
     ) -> ServerResult<()> {
-        if self.is_blacklisted(blacklister, target_public_key).await? {
-            return Err(WsError::AlreadyOnTheBlacklist.into());
-        }
         if blacklister.public_key == target_public_key.to_string() {
             return Err(WsError::CannotAddSelfToBlacklist.into());
         }
 
-        if let Some(mut user) = get_user_status(self, blacklister, target_public_key)
-            .await?
-            .map(IntoActiveModel::into_active_model)
+        if let Some(mut user) = get_user_status(
+            self,
+            blacklister,
+            target_public_key,
+            AccessStatus::Whitelisted,
+        )
+        .await?
+        .map(IntoActiveModel::into_active_model)
         {
             user.status = Set(AccessStatus::Blacklisted);
             user.updated_at = Set(Utc::now());
             user.update(self).await?;
-        } else {
-            UsersStatusActiveModel {
-                user_id: Set(blacklister.id),
-                target: Set(target_public_key.to_string()),
-                status: Set(AccessStatus::Blacklisted),
-                updated_at: Set(Utc::now()),
-                ..Default::default()
+        } else if let Err(err) = (UsersStatusActiveModel {
+            user_id: Set(blacklister.id),
+            target: Set(target_public_key.to_string()),
+            status: Set(AccessStatus::Blacklisted),
+            updated_at: Set(Utc::now()),
+            ..Default::default()
+        }
+        .save(self)
+        .await)
+        {
+            match err.sql_err() {
+                Some(SqlErr::UniqueConstraintViolation(_)) => {
+                    return Err(WsError::AlreadyOnTheBlacklist.into());
+                }
+                _ => return Err(err.into()),
             }
-            .save(self)
-            .await?;
         }
 
         Ok(())
@@ -278,9 +294,14 @@ async fn get_user_status(
     conn: &DatabaseConnection,
     user: &UserModel,
     target_public_key: &PublicKey,
+    status: AccessStatus,
 ) -> ServerResult<Option<UsersStatusModel>> {
     user.find_related(UsersStatusEntity)
-        .filter(UsersStatusColumn::Target.eq(target_public_key.to_string()))
+        .filter(
+            UsersStatusColumn::Target
+                .eq(target_public_key.to_string())
+                .and(UsersStatusColumn::Status.eq(status)),
+        )
         .one(conn)
         .await
         .map_err(Into::into)
diff --git a/crates/oxidetalis_migrations/src/create_users_status.rs b/crates/oxidetalis_migrations/src/create_users_status.rs
index ad57d00..a0ad213 100644
--- a/crates/oxidetalis_migrations/src/create_users_status.rs
+++ b/crates/oxidetalis_migrations/src/create_users_status.rs
@@ -40,6 +40,18 @@ impl MigrationTrait for Migration {
                     .to_owned(),
             )
             .await?;
+        manager
+            .create_index(
+                Index::create()
+                    .if_not_exists()
+                    .name("sep_status")
+                    .table(UsersStatus::Table)
+                    .col(UsersStatus::UserId)
+                    .col(UsersStatus::Target)
+                    .unique()
+                    .to_owned(),
+            )
+            .await?;
 
         manager
             .create_table(