feat: remove api brain secrets and schemas on delete (#1621)

Issue: https://github.com/StanGirard/quivr/issues/1573
This commit is contained in:
Mamadou DICKO 2023-11-13 19:08:47 +01:00 committed by GitHub
parent ee864e6441
commit 8ed0adf7b2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
9 changed files with 142 additions and 81 deletions

View File

@ -49,3 +49,10 @@ class PublicBrain(BaseModel):
description: Optional[str] description: Optional[str]
number_of_subscribers: int = 0 number_of_subscribers: int = 0
last_update: str last_update: str
class BrainUser(BaseModel):
id: UUID
user_id: UUID
rights: RoleEnum
default_brain: bool = False

View File

@ -49,40 +49,6 @@ class Brain(BaseModel):
commons=commons, *args, **kwargs # pyright: ignore reportPrivateUsage=none commons=commons, *args, **kwargs # pyright: ignore reportPrivateUsage=none
) # pyright: ignore reportPrivateUsage=none ) # pyright: ignore reportPrivateUsage=none
# TODO: move this to a brand new BrainService
def get_brain_users(self):
response = (
self.supabase_client.table("brains_users")
.select("id:brain_id, *")
.filter("brain_id", "eq", self.id) # type: ignore
.execute()
)
return response.data
# TODO: move this to a brand new BrainService
def delete_user_from_brain(self, user_id):
results = (
self.supabase_client.table("brains_users")
.select("*")
.match({"brain_id": self.id, "user_id": user_id})
.execute()
)
if len(results.data) != 0:
self.supabase_client.table("brains_users").delete().match(
{"brain_id": self.id, "user_id": user_id}
).execute()
def delete_brain(self, user_id):
results = self.supabase_db.delete_brain_user_by_id(user_id, self.id) # type: ignore
if len(results) == 0:
return {"message": "You are not the owner of this brain."}
else:
self.supabase_db.delete_brain_vector(self.id) # type: ignore
self.supabase_db.delete_brain_users(self.id) # type: ignore
self.supabase_db.delete_brain(self.id) # type: ignore
def create_brain_vector(self, vector_id, file_sha1): def create_brain_vector(self, vector_id, file_sha1):
return self.supabase_db.create_brain_vector(self.id, vector_id, file_sha1) # type: ignore return self.supabase_db.create_brain_vector(self.id, vector_id, file_sha1) # type: ignore

View File

@ -2,7 +2,13 @@ from typing import Optional
from uuid import UUID from uuid import UUID
from logger import get_logger from logger import get_logger
from models.brain_entity import BrainEntity, BrainType, MinimalBrainEntity, PublicBrain from models.brain_entity import (
BrainEntity,
BrainType,
BrainUser,
MinimalBrainEntity,
PublicBrain,
)
from models.databases.repository import Repository from models.databases.repository import Repository
from models.databases.supabase.api_brain_definition import ( from models.databases.supabase.api_brain_definition import (
CreateApiBrainDefinition, CreateApiBrainDefinition,
@ -337,3 +343,13 @@ class Brain(Repository):
if len(response) == 0: if len(response) == 0:
raise ValueError(f"Brain with id {brain_id} does not exist.") raise ValueError(f"Brain with id {brain_id} does not exist.")
return response[0]["count"] return response[0]["count"]
def get_brain_users(self, brain_id: UUID) -> list[BrainUser]:
response = (
self.db.table("brains_users")
.select("id:brain_id, *")
.filter("brain_id", "eq", str(brain_id))
.execute()
)
return [BrainUser(**item) for item in response.data]

View File

@ -0,0 +1,36 @@
from uuid import UUID
from fastapi import HTTPException
from models.brain_entity import BrainType
from models.settings import get_supabase_db
from repository.api_brain_definition.delete_api_brain_definition import (
delete_api_brain_definition,
)
from repository.brain import get_brain_by_id
from repository.brain.delete_brain_secrets import delete_brain_secrets
from repository.knowledge.remove_brain_all_knowledge import (
remove_brain_all_knowledge,
)
def delete_brain(brain_id: UUID) -> dict[str, str]:
supabase_db = get_supabase_db()
brain_to_delete = get_brain_by_id(brain_id=brain_id)
if brain_to_delete is None:
raise HTTPException(status_code=404, detail="Brain not found.")
if brain_to_delete.brain_type == BrainType.API:
delete_brain_secrets(
brain_id=brain_id,
)
delete_api_brain_definition(brain_id=brain_id)
else:
remove_brain_all_knowledge(brain_id)
supabase_db.delete_brain_vector(str(brain_id))
supabase_db.delete_brain_users(str(brain_id))
supabase_db.delete_brain(str(brain_id))
return {"message": "Brain deleted."}

View File

@ -0,0 +1,30 @@
from uuid import UUID
from fastapi import HTTPException
from models.settings import get_supabase_db
from repository.api_brain_definition.get_api_brain_definition import (
get_api_brain_definition,
)
from repository.external_api_secret import delete_secret
def delete_brain_secrets(brain_id: UUID) -> None:
supabase_db = get_supabase_db()
brain_definition = get_api_brain_definition(brain_id=brain_id)
if brain_definition is None:
raise HTTPException(status_code=404, detail="Brain definition not found.")
secrets = brain_definition.secrets
if len(secrets) > 0:
brain_users = supabase_db.get_brain_users(brain_id=brain_id)
for user in brain_users:
for secret in secrets:
delete_secret(
user_id=user.user_id,
brain_id=brain_id,
secret_name=secret.name,
)

View File

@ -1,10 +1,34 @@
from uuid import UUID from uuid import UUID
from fastapi import HTTPException
from models.brain_entity import BrainType
from models.settings import get_supabase_db from models.settings import get_supabase_db
from repository.api_brain_definition.get_api_brain_definition import (
get_api_brain_definition,
)
from repository.brain.get_brain_by_id import get_brain_by_id
from repository.external_api_secret.delete_secret import delete_secret
def delete_brain_user(user_id: UUID, brain_id: UUID) -> None: def delete_brain_user(user_id: UUID, brain_id: UUID) -> None:
supabase_db = get_supabase_db() supabase_db = get_supabase_db()
brain_to_delete_user_from = get_brain_by_id(brain_id=brain_id)
if brain_to_delete_user_from is None:
raise HTTPException(status_code=404, detail="Brain not found.")
if brain_to_delete_user_from.brain_type == BrainType.API:
brain_definition = get_api_brain_definition(brain_id=brain_id)
if brain_definition is None:
raise HTTPException(status_code=404, detail="Brain definition not found.")
secrets = brain_definition.secrets
for secret in secrets:
delete_secret(
user_id=user_id,
brain_id=brain_id,
secret_name=secret.name,
)
supabase_db.delete_brain_user_by_id( supabase_db.delete_brain_user_by_id(
user_id=user_id, user_id=user_id,
brain_id=brain_id, brain_id=brain_id,

View File

@ -0,0 +1,10 @@
from uuid import UUID
from models.brain_entity import BrainUser
from models.settings import get_supabase_db
def get_brain_users(brain_id: UUID) -> list[BrainUser]:
supabase_db = get_supabase_db()
return supabase_db.get_brain_users(brain_id)

View File

@ -10,7 +10,7 @@ def delete_secret(user_id: UUID, brain_id: UUID, secret_name: str) -> bool:
response = supabase_client.rpc( response = supabase_client.rpc(
"delete_secret", "delete_secret",
{ {
"name": build_secret_unique_name( "secret_name": build_secret_unique_name(
user_id=user_id, brain_id=brain_id, secret_name=secret_name user_id=user_id, brain_id=brain_id, secret_name=secret_name
), ),
}, },

View File

@ -3,7 +3,7 @@ from uuid import UUID
from auth.auth_bearer import AuthBearer, get_current_user from auth.auth_bearer import AuthBearer, get_current_user
from fastapi import APIRouter, Depends, HTTPException from fastapi import APIRouter, Depends, HTTPException
from models import Brain, BrainSubscription, PromptStatusEnum, UserIdentity from models import BrainSubscription, PromptStatusEnum, UserIdentity
from pydantic import BaseModel from pydantic import BaseModel
from repository.brain import ( from repository.brain import (
create_brain_user, create_brain_user,
@ -12,16 +12,15 @@ from repository.brain import (
get_brain_for_user, get_brain_for_user,
update_brain_user_rights, update_brain_user_rights,
) )
from repository.brain.delete_brain import delete_brain
from repository.brain.delete_brain_user import delete_brain_user from repository.brain.delete_brain_user import delete_brain_user
from repository.brain.get_brain_users import get_brain_users
from repository.brain_subscription import ( from repository.brain_subscription import (
SubscriptionInvitationService, SubscriptionInvitationService,
resend_invitation_email, resend_invitation_email,
) )
from repository.knowledge.remove_brain_all_knowledge import (
remove_brain_all_knowledge,
)
from repository.prompt import delete_prompt_by_id, get_prompt_by_id from repository.prompt import delete_prompt_by_id, get_prompt_by_id
from repository.user import get_user_email_by_user_id, get_user_id_by_user_email from repository.user import get_user_id_by_user_email
from routes.authorizations.brain_authorization import ( from routes.authorizations.brain_authorization import (
RoleEnum, RoleEnum,
@ -99,29 +98,6 @@ def invite_users_to_brain(
Depends(has_brain_authorization([RoleEnum.Owner, RoleEnum.Editor])), Depends(has_brain_authorization([RoleEnum.Owner, RoleEnum.Editor])),
], ],
) )
def get_brain_users(
brain_id: UUID,
):
"""
Get all users for a brain
"""
brain = Brain(
id=brain_id,
)
brain_users = brain.get_brain_users()
brain_access_list = []
for brain_user in brain_users:
brain_access = {}
# TODO: find a way to fetch user email concurrently
brain_access["email"] = get_user_email_by_user_id(brain_user["user_id"])
brain_access["rights"] = brain_user["rights"]
brain_access_list.append(brain_access)
return brain_access_list
@subscription_router.delete( @subscription_router.delete(
"/brains/{brain_id}/subscription", "/brains/{brain_id}/subscription",
) )
@ -139,9 +115,6 @@ async def remove_user_subscription(
detail="Brain not found while trying to delete", detail="Brain not found while trying to delete",
) )
brain = Brain(
id=brain_id,
)
user_brain = get_brain_for_user(current_user.id, brain_id) user_brain = get_brain_for_user(current_user.id, brain_id)
if user_brain is None: if user_brain is None:
raise HTTPException( raise HTTPException(
@ -150,21 +123,21 @@ async def remove_user_subscription(
) )
if user_brain.rights != "Owner": if user_brain.rights != "Owner":
brain.delete_user_from_brain(current_user.id) delete_brain_user(current_user.id, brain_id)
else: else:
brain_users = brain.get_brain_users() brain_users = get_brain_users(
brain_id=brain_id,
)
brain_other_owners = [ brain_other_owners = [
brain brain
for brain in brain_users for brain in brain_users
if brain["rights"] == "Owner" if brain.rights == "Owner" and str(brain.user_id) != str(current_user.id)
and str(brain["user_id"]) != str(current_user.id)
] ]
if len(brain_other_owners) == 0: if len(brain_other_owners) == 0:
# Delete its prompt if it's private delete_brain(
brain_id=brain_id,
remove_brain_all_knowledge(brain_id) )
brain.delete_brain(current_user.id)
if targeted_brain.prompt_id: if targeted_brain.prompt_id:
brain_to_delete_prompt = get_prompt_by_id(targeted_brain.prompt_id) brain_to_delete_prompt = get_prompt_by_id(targeted_brain.prompt_id)
if brain_to_delete_prompt is not None and ( if brain_to_delete_prompt is not None and (
@ -173,7 +146,10 @@ async def remove_user_subscription(
delete_prompt_by_id(targeted_brain.prompt_id) delete_prompt_by_id(targeted_brain.prompt_id)
else: else:
brain.delete_user_from_brain(current_user.id) delete_brain_user(
current_user.id,
brain_id,
)
return {"message": f"Subscription removed successfully from brain {brain_id}"} return {"message": f"Subscription removed successfully from brain {brain_id}"}
@ -321,10 +297,6 @@ def update_brain_subscription(
detail="User not found", detail="User not found",
) )
brain = Brain(
id=brain_id,
)
# check if user is an editor but trying to give high level permissions # check if user is an editor but trying to give high level permissions
if subscription.rights == "Owner": if subscription.rights == "Owner":
try: try:
@ -363,7 +335,7 @@ def update_brain_subscription(
current_user.id, current_user.id,
RoleEnum.Owner, RoleEnum.Owner,
) )
brain.delete_user_from_brain(user_id) delete_brain_user(user_id, brain_id)
except HTTPException: except HTTPException:
raise HTTPException( raise HTTPException(
status_code=403, status_code=403,