UniqAI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2024-11-16 05:24:02 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

refac: files rbac

Browse Source
This commit is contained in:
Timothy J. Baek 2024-08-23 16:19:04 +02:00
parent 8b3d5e8b80
commit 4519ddd0e9
2 changed files with 15 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
backend/apps/webui/models/files.py
View File

@ -98,6 +98,13 @@ class FilesTable:
return [FileModel.model_validate(file) for file in db.query(File).all()] return [FileModel.model_validate(file) for file in db.query(File).all()]
def get_files_by_user_id(self, user_id: str) -> list[FileModel]:
with get_db() as db:
return [
FileModel.model_validate(file)
for file in db.query(File).filter_by(user_id=user_id).all()
]
def delete_file_by_id(self, id: str) -> bool: def delete_file_by_id(self, id: str) -> bool:
with get_db() as db: with get_db() as db:

14
backend/apps/webui/routers/files.py
View File

@ -106,7 +106,10 @@ def upload_file(file: UploadFile = File(...), user=Depends(get_verified_user)):
@router.get("/", response_model=list[FileModel]) @router.get("/", response_model=list[FileModel])
async def list_files(user=Depends(get_verified_user)): async def list_files(user=Depends(get_verified_user)):
files = Files.get_files() if user.role == "admin":
files = Files.get_files()
else:
files = Files.get_files_by_user_id(user.id)
return files return files
@ -156,7 +159,7 @@ async def delete_all_files(user=Depends(get_admin_user)):
async def get_file_by_id(id: str, user=Depends(get_verified_user)): async def get_file_by_id(id: str, user=Depends(get_verified_user)):
file = Files.get_file_by_id(id) file = Files.get_file_by_id(id)
if file: if file and (file.user_id == user.id or user.role == "admin"):
return file return file
else: else:
raise HTTPException( raise HTTPException(
@ -174,7 +177,7 @@ async def get_file_by_id(id: str, user=Depends(get_verified_user)):
async def get_file_content_by_id(id: str, user=Depends(get_verified_user)): async def get_file_content_by_id(id: str, user=Depends(get_verified_user)):
file = Files.get_file_by_id(id) file = Files.get_file_by_id(id)
if file: if file and (file.user_id == user.id or user.role == "admin"):
file_path = Path(file.meta["path"]) file_path = Path(file.meta["path"])
# Check if the file already exists in the cache # Check if the file already exists in the cache
@ -197,7 +200,7 @@ async def get_file_content_by_id(id: str, user=Depends(get_verified_user)):
async def get_file_content_by_id(id: str, user=Depends(get_verified_user)): async def get_file_content_by_id(id: str, user=Depends(get_verified_user)):
file = Files.get_file_by_id(id) file = Files.get_file_by_id(id)
if file: if file and (file.user_id == user.id or user.role == "admin"):
file_path = Path(file.meta["path"]) file_path = Path(file.meta["path"])
# Check if the file already exists in the cache # Check if the file already exists in the cache
@ -224,8 +227,7 @@ async def get_file_content_by_id(id: str, user=Depends(get_verified_user)):
@router.delete("/{id}") @router.delete("/{id}")
async def delete_file_by_id(id: str, user=Depends(get_verified_user)): async def delete_file_by_id(id: str, user=Depends(get_verified_user)):
file = Files.get_file_by_id(id) file = Files.get_file_by_id(id)
if file and (file.user_id == user.id or user.role == "admin"):
if file:
result = Files.delete_file_by_id(id) result = Files.delete_file_by_id(id)
if result: if result:
return {"message": "File deleted successfully"} return {"message": "File deleted successfully"}