From cc9420bb99939a24a758628987b3e42b1eff8afa Mon Sep 17 00:00:00 2001
From: Nikita Tikhonov <nekitdevofficial@gmail.com>
Date: Tue, 20 Feb 2024 04:35:47 +0300
Subject: [PATCH] Handle `QueryAssertionError` in following users/playlists.

Fixes #8.
---
 melody/kit/endpoints/v1/self.py | 24 +++++++++++++++++++++---
 melody/kit/errors/core.py       |  2 ++
 melody/kit/errors/users.py      | 24 ++++++++++++++++++++++++
 schema/default.esdl             |  7 +++++++
 4 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/melody/kit/endpoints/v1/self.py b/melody/kit/endpoints/v1/self.py
index 569f160..4fba91f 100644
--- a/melody/kit/endpoints/v1/self.py
+++ b/melody/kit/endpoints/v1/self.py
@@ -2,6 +2,7 @@
 from uuid import UUID
 
 from aiofiles import open as async_open
+from edgedb import QueryAssertionError
 from fastapi import Body, Depends
 from fastapi.responses import FileResponse
 from typing_extensions import Annotated
@@ -13,7 +14,12 @@
 from melody.kit.dependencies.images import ImageDependency
 from melody.kit.dependencies.request_urls import RequestURLDependency
 from melody.kit.enums import EntityType, Platform, PrivacyType, Tag
-from melody.kit.errors.users import UserImageNotFound, UserNotFound
+from melody.kit.errors.users import (
+    UserFollowSelfForbidden,
+    UserFollowSelfPlaylistsForbidden,
+    UserImageNotFound,
+    UserNotFound,
+)
 from melody.kit.models.pagination import Pagination
 from melody.kit.models.user import (
     UserAlbums,
@@ -438,7 +444,13 @@ async def get_self_following(
 async def add_self_following(
     context: FollowingWriteTokenDependency, ids: UUIDListDependency
 ) -> None:
-    await database.add_user_following(user_id=context.user_id, ids=ids)
+    self_id = context.user_id
+
+    try:
+        await database.add_user_following(user_id=self_id, ids=ids)
+
+    except QueryAssertionError:
+        raise UserFollowSelfForbidden(self_id) from None
 
 
 @v1.delete(
@@ -489,7 +501,13 @@ async def get_self_followed_playlists(
 async def add_self_followed_playlists(
     context: LibraryWriteTokenDependency, ids: UUIDListDependency
 ) -> None:
-    await database.add_user_followed_playlists(user_id=context.user_id, ids=ids)
+    self_id = context.user_id
+
+    try:
+        await database.add_user_followed_playlists(user_id=self_id, ids=ids)
+
+    except QueryAssertionError:
+        raise UserFollowSelfPlaylistsForbidden(self_id) from None
 
 
 @v1.delete(
diff --git a/melody/kit/errors/core.py b/melody/kit/errors/core.py
index 34e952a..8ecccc9 100644
--- a/melody/kit/errors/core.py
+++ b/melody/kit/errors/core.py
@@ -49,6 +49,8 @@ class ErrorCode(Enum):
     USER_NOT_FOUND = 13501
     USER_INACCESSIBLE = 13502
     USER_IMAGE_NOT_FOUND = 13503
+    USER_FOLLOW_SELF_FORBIDDEN = 13504
+    USER_FOLLOW_SELF_PLAYLISTS_FORBIDDEN = 13505
 
     CLIENT_ERROR = 13600
     CLIENT_NOT_FOUND = 13601
diff --git a/melody/kit/errors/users.py b/melody/kit/errors/users.py
index 8bc3be8..d25bde6 100644
--- a/melody/kit/errors/users.py
+++ b/melody/kit/errors/users.py
@@ -11,6 +11,8 @@
     "UserNotFound",
     "UserInaccessible",
     "UserImageNotFound",
+    "UserFollowSelfForbidden",
+    "UserFollowSelfPlaylistsForbidden",
 )
 
 
@@ -63,3 +65,25 @@ def __init__(self, user_id: UUID) -> None:
 class UserImageNotFound(UserError):
     def __init__(self, user_id: UUID) -> None:
         super().__init__(user_id, user_image_not_found(user_id))
+
+
+USER_CAN_NOT_FOLLOW_SELF = "user `{}` can not follow themselves"
+user_can_not_follow_self = USER_CAN_NOT_FOLLOW_SELF.format
+
+
+@default_code(ErrorCode.USER_FOLLOW_SELF_FORBIDDEN)
+@default_status_code(status.HTTP_403_FORBIDDEN)
+class UserFollowSelfForbidden(UserError):
+    def __init__(self, user_id: UUID) -> None:
+        super().__init__(user_id, user_can_not_follow_self(user_id))
+
+
+USER_CAN_NOT_FOLLOW_SELF_PLAYLISTS = "user `{}` can not follow their own playlists"
+user_can_not_follow_self_playlists = USER_CAN_NOT_FOLLOW_SELF_PLAYLISTS.format
+
+
+@default_code(ErrorCode.USER_FOLLOW_SELF_PLAYLISTS_FORBIDDEN)
+@default_status_code(status.HTTP_403_FORBIDDEN)
+class UserFollowSelfPlaylistsForbidden(UserError):
+    def __init__(self, user_id: UUID) -> None:
+        super().__init__(user_id, user_can_not_follow_self_playlists(user_id))
diff --git a/schema/default.esdl b/schema/default.esdl
index b9e9020..c6b987f 100644
--- a/schema/default.esdl
+++ b/schema/default.esdl
@@ -171,6 +171,13 @@ module default {
 
         multi followed_playlists extending with_linked_at: Playlist;
 
+        trigger forbid_follow_self_playlists after insert, update for each do (
+            assert(
+                not exists (__new__.playlists intersect __new__.followed_playlists),
+                message := "users can not follow their own playlists",
+            )
+        );
+
         followed_playlist_count := count(.followed_playlists);
 
         multi streams := .<user[is Stream];