From fc76f9d46696545736c3391776daa55cbd2625a5 Mon Sep 17 00:00:00 2001
From: Mamadou DICKO <63923024+mamadoudicko@users.noreply.github.com>
Date: Thu, 13 Jul 2023 18:05:36 +0200
Subject: [PATCH] feat(shareBrain): add unsubscribe endpoint (#633)

---
 backend/models/brains.py              | 31 ++++++++++++++-
 backend/routes/subscription_routes.py | 54 +++++++++++++++++++++++++--
 2 files changed, 79 insertions(+), 6 deletions(-)

diff --git a/backend/models/brains.py b/backend/models/brains.py
index 2378f3dc3..e0e1f341d 100644
--- a/backend/models/brains.py
+++ b/backend/models/brains.py
@@ -3,11 +3,12 @@ from typing import Any, List, Optional
 from uuid import UUID
 
 from logger import get_logger
-from models.settings import CommonsDep, common_dependencies
-from models.users import User
 from pydantic import BaseModel
 from utils.vectors import get_unique_files_from_vector_ids
 
+from models.settings import CommonsDep, common_dependencies
+from models.users import User
+
 logger = get_logger(__name__)
 
 
@@ -49,6 +50,32 @@ class Brain(BaseModel):
             commons=commons, *args, **kwargs  # pyright: ignore reportPrivateUsage=none
         )  # pyright: ignore reportPrivateUsage=none
 
+    # TODO: move this to a brand new BrainService
+    def get_brain_users(self):
+        response = (
+            self.commons["supabase"]
+            .table("brains_users")
+            .select("id:brain_id, *")
+            .filter("brain_id", "eq", self.id)
+            .execute()
+        )
+        return response.data
+
+    # TODO: move this to a brand new BrainService
+    def delete_user_from_brain(self, user_id):
+        results = (
+            self.commons["supabase"]
+            .table("brains_users")
+            .select("*")
+            .match({"brain_id": self.id, "user_id": user_id})
+            .execute()
+        )
+
+        if len(results.data) != 0:
+            self.commons["supabase"].table("brains_users").delete().match(
+                {"brain_id": self.id, "user_id": user_id}
+            ).execute()
+
     def get_user_brains(self, user_id):
         response = (
             self.commons["supabase"]
diff --git a/backend/routes/subscription_routes.py b/backend/routes/subscription_routes.py
index 78c8a4eec..5c21ac510 100644
--- a/backend/routes/subscription_routes.py
+++ b/backend/routes/subscription_routes.py
@@ -3,6 +3,7 @@ from uuid import UUID
 
 from auth.auth_bearer import get_current_user
 from fastapi import APIRouter, Depends, HTTPException
+from models.brains import Brain
 from models.brains_subscription_invitations import BrainSubscription
 from models.users import User
 
@@ -10,12 +11,19 @@ subscription_router = APIRouter()
 
 
 @subscription_router.post("/brain/{brain_id}/subscription")
-async def invite_user_to_brain(brain_id: UUID, users: List[dict], current_user: User = Depends(get_current_user)):
+async def invite_user_to_brain(
+    brain_id: UUID, users: List[dict], current_user: User = Depends(get_current_user)
+):
     # TODO: Ensure the current user has permissions to invite users to this brain
-    
+
     for user in users:
-        subscription = BrainSubscription(brain_id=brain_id, email=user['email'], rights=user['rights'], inviter_email=current_user.email or "Quivr")
-        
+        subscription = BrainSubscription(
+            brain_id=brain_id,
+            email=user["email"],
+            rights=user["rights"],
+            inviter_email=current_user.email or "Quivr",
+        )
+
         try:
             subscription.create_or_update_subscription_invitation()
             subscription.resend_invitation_email()
@@ -23,3 +31,41 @@ async def invite_user_to_brain(brain_id: UUID, users: List[dict], current_user:
             raise HTTPException(status_code=400, detail=f"Error inviting user: {e}")
 
     return {"message": "Invitations sent successfully"}
+
+
+@subscription_router.delete(
+    "/brain/{brain_id}/subscription",
+)
+async def remove_user_subscription(
+    brain_id: UUID, current_user: User = Depends(get_current_user)
+):
+    """
+    Remove a user's subscription to a brain
+    """
+    brain = Brain(
+        id=brain_id,
+    )
+    user_brain = brain.get_brain_for_user(current_user.id)
+    if user_brain is None:
+        raise HTTPException(
+            status_code=403,
+            detail="You don't have permission for this brain",
+        )
+
+    if user_brain.get("rights") != "Owner":
+        brain.delete_user_from_brain(current_user.id)
+    else:
+        brain_other_users = brain.get_brain_users()
+        brain_other_owners = [
+            brain
+            for brain in brain_other_users
+            if brain["rights"] == "Owner"
+            and str(brain["user_id"]) != str(current_user.id)
+        ]
+
+        if len(brain_other_owners) == 0:
+            brain.delete_brain(current_user.id)
+        else:
+            brain.delete_user_from_brain(current_user.id)
+
+    return {"message": f"Subscription removed successfully from brain {brain_id}"}