refac/enh: db session sharing
This commit is contained in:
@@ -22,6 +22,8 @@ from open_webui.models.knowledge import Knowledges
|
||||
|
||||
from open_webui.config import UPLOAD_DIR
|
||||
from open_webui.constants import ERROR_MESSAGES
|
||||
from open_webui.internal.db import get_session
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
|
||||
from fastapi import APIRouter, Depends, File, HTTPException, UploadFile, status, Request
|
||||
@@ -44,7 +46,11 @@ router = APIRouter()
|
||||
|
||||
|
||||
@router.get("/", response_model=list[FolderNameIdResponse])
|
||||
async def get_folders(request: Request, user=Depends(get_verified_user)):
|
||||
async def get_folders(
|
||||
request: Request,
|
||||
user=Depends(get_verified_user),
|
||||
db: Session = Depends(get_session),
|
||||
):
|
||||
if request.app.state.config.ENABLE_FOLDERS is False:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
@@ -55,22 +61,23 @@ async def get_folders(request: Request, user=Depends(get_verified_user)):
|
||||
user.id,
|
||||
"features.folders",
|
||||
request.app.state.config.USER_PERMISSIONS,
|
||||
db=db,
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
)
|
||||
|
||||
folders = Folders.get_folders_by_user_id(user.id)
|
||||
folders = Folders.get_folders_by_user_id(user.id, db=db)
|
||||
|
||||
# Verify folder data integrity
|
||||
folder_list = []
|
||||
for folder in folders:
|
||||
if folder.parent_id and not Folders.get_folder_by_id_and_user_id(
|
||||
folder.parent_id, user.id
|
||||
folder.parent_id, user.id, db=db
|
||||
):
|
||||
folder = Folders.update_folder_parent_id_by_id_and_user_id(
|
||||
folder.id, user.id, None
|
||||
folder.id, user.id, None, db=db
|
||||
)
|
||||
|
||||
if folder.data:
|
||||
@@ -80,12 +87,12 @@ async def get_folders(request: Request, user=Depends(get_verified_user)):
|
||||
|
||||
if file.get("type") == "file":
|
||||
if Files.check_access_by_user_id(
|
||||
file.get("id"), user.id, "read"
|
||||
file.get("id"), user.id, "read", db=db
|
||||
):
|
||||
valid_files.append(file)
|
||||
elif file.get("type") == "collection":
|
||||
if Knowledges.check_access_by_user_id(
|
||||
file.get("id"), user.id, "read"
|
||||
file.get("id"), user.id, "read", db=db
|
||||
):
|
||||
valid_files.append(file)
|
||||
else:
|
||||
@@ -93,7 +100,7 @@ async def get_folders(request: Request, user=Depends(get_verified_user)):
|
||||
|
||||
folder.data["files"] = valid_files
|
||||
Folders.update_folder_by_id_and_user_id(
|
||||
folder.id, user.id, FolderUpdateForm(data=folder.data)
|
||||
folder.id, user.id, FolderUpdateForm(data=folder.data), db=db
|
||||
)
|
||||
|
||||
folder_list.append(FolderNameIdResponse(**folder.model_dump()))
|
||||
@@ -107,9 +114,13 @@ async def get_folders(request: Request, user=Depends(get_verified_user)):
|
||||
|
||||
|
||||
@router.post("/")
|
||||
def create_folder(form_data: FolderForm, user=Depends(get_verified_user)):
|
||||
def create_folder(
|
||||
form_data: FolderForm,
|
||||
user=Depends(get_verified_user),
|
||||
db: Session = Depends(get_session),
|
||||
):
|
||||
folder = Folders.get_folder_by_parent_id_and_user_id_and_name(
|
||||
None, user.id, form_data.name
|
||||
None, user.id, form_data.name, db=db
|
||||
)
|
||||
|
||||
if folder:
|
||||
@@ -119,7 +130,7 @@ def create_folder(form_data: FolderForm, user=Depends(get_verified_user)):
|
||||
)
|
||||
|
||||
try:
|
||||
folder = Folders.insert_new_folder(user.id, form_data)
|
||||
folder = Folders.insert_new_folder(user.id, form_data, db=db)
|
||||
return folder
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
@@ -136,8 +147,10 @@ def create_folder(form_data: FolderForm, user=Depends(get_verified_user)):
|
||||
|
||||
|
||||
@router.get("/{id}", response_model=Optional[FolderModel])
|
||||
async def get_folder_by_id(id: str, user=Depends(get_verified_user)):
|
||||
folder = Folders.get_folder_by_id_and_user_id(id, user.id)
|
||||
async def get_folder_by_id(
|
||||
id: str, user=Depends(get_verified_user), db: Session = Depends(get_session)
|
||||
):
|
||||
folder = Folders.get_folder_by_id_and_user_id(id, user.id, db=db)
|
||||
if folder:
|
||||
return folder
|
||||
else:
|
||||
@@ -154,15 +167,18 @@ async def get_folder_by_id(id: str, user=Depends(get_verified_user)):
|
||||
|
||||
@router.post("/{id}/update")
|
||||
async def update_folder_name_by_id(
|
||||
id: str, form_data: FolderUpdateForm, user=Depends(get_verified_user)
|
||||
id: str,
|
||||
form_data: FolderUpdateForm,
|
||||
user=Depends(get_verified_user),
|
||||
db: Session = Depends(get_session),
|
||||
):
|
||||
folder = Folders.get_folder_by_id_and_user_id(id, user.id)
|
||||
folder = Folders.get_folder_by_id_and_user_id(id, user.id, db=db)
|
||||
if folder:
|
||||
|
||||
if form_data.name is not None:
|
||||
# Check if folder with same name exists
|
||||
existing_folder = Folders.get_folder_by_parent_id_and_user_id_and_name(
|
||||
folder.parent_id, user.id, form_data.name
|
||||
folder.parent_id, user.id, form_data.name, db=db
|
||||
)
|
||||
if existing_folder and existing_folder.id != id:
|
||||
raise HTTPException(
|
||||
@@ -171,7 +187,9 @@ async def update_folder_name_by_id(
|
||||
)
|
||||
|
||||
try:
|
||||
folder = Folders.update_folder_by_id_and_user_id(id, user.id, form_data)
|
||||
folder = Folders.update_folder_by_id_and_user_id(
|
||||
id, user.id, form_data, db=db
|
||||
)
|
||||
return folder
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
@@ -198,12 +216,15 @@ class FolderParentIdForm(BaseModel):
|
||||
|
||||
@router.post("/{id}/update/parent")
|
||||
async def update_folder_parent_id_by_id(
|
||||
id: str, form_data: FolderParentIdForm, user=Depends(get_verified_user)
|
||||
id: str,
|
||||
form_data: FolderParentIdForm,
|
||||
user=Depends(get_verified_user),
|
||||
db: Session = Depends(get_session),
|
||||
):
|
||||
folder = Folders.get_folder_by_id_and_user_id(id, user.id)
|
||||
folder = Folders.get_folder_by_id_and_user_id(id, user.id, db=db)
|
||||
if folder:
|
||||
existing_folder = Folders.get_folder_by_parent_id_and_user_id_and_name(
|
||||
form_data.parent_id, user.id, folder.name
|
||||
form_data.parent_id, user.id, folder.name, db=db
|
||||
)
|
||||
|
||||
if existing_folder:
|
||||
@@ -214,7 +235,7 @@ async def update_folder_parent_id_by_id(
|
||||
|
||||
try:
|
||||
folder = Folders.update_folder_parent_id_by_id_and_user_id(
|
||||
id, user.id, form_data.parent_id
|
||||
id, user.id, form_data.parent_id, db=db
|
||||
)
|
||||
return folder
|
||||
except Exception as e:
|
||||
@@ -242,13 +263,16 @@ class FolderIsExpandedForm(BaseModel):
|
||||
|
||||
@router.post("/{id}/update/expanded")
|
||||
async def update_folder_is_expanded_by_id(
|
||||
id: str, form_data: FolderIsExpandedForm, user=Depends(get_verified_user)
|
||||
id: str,
|
||||
form_data: FolderIsExpandedForm,
|
||||
user=Depends(get_verified_user),
|
||||
db: Session = Depends(get_session),
|
||||
):
|
||||
folder = Folders.get_folder_by_id_and_user_id(id, user.id)
|
||||
folder = Folders.get_folder_by_id_and_user_id(id, user.id, db=db)
|
||||
if folder:
|
||||
try:
|
||||
folder = Folders.update_folder_is_expanded_by_id_and_user_id(
|
||||
id, user.id, form_data.is_expanded
|
||||
id, user.id, form_data.is_expanded, db=db
|
||||
)
|
||||
return folder
|
||||
except Exception as e:
|
||||
@@ -276,10 +300,11 @@ async def delete_folder_by_id(
|
||||
id: str,
|
||||
delete_contents: Optional[bool] = True,
|
||||
user=Depends(get_verified_user),
|
||||
db: Session = Depends(get_session),
|
||||
):
|
||||
if Chats.count_chats_by_folder_id_and_user_id(id, user.id):
|
||||
if Chats.count_chats_by_folder_id_and_user_id(id, user.id, db=db):
|
||||
chat_delete_permission = has_permission(
|
||||
user.id, "chat.delete", request.app.state.config.USER_PERMISSIONS
|
||||
user.id, "chat.delete", request.app.state.config.USER_PERMISSIONS, db=db
|
||||
)
|
||||
if user.role != "admin" and not chat_delete_permission:
|
||||
raise HTTPException(
|
||||
@@ -288,19 +313,21 @@ async def delete_folder_by_id(
|
||||
)
|
||||
|
||||
folders = []
|
||||
folders.append(Folders.get_folder_by_id_and_user_id(id, user.id))
|
||||
folders.append(Folders.get_folder_by_id_and_user_id(id, user.id, db=db))
|
||||
while folders:
|
||||
folder = folders.pop()
|
||||
if folder:
|
||||
try:
|
||||
folder_ids = Folders.delete_folder_by_id_and_user_id(id, user.id)
|
||||
folder_ids = Folders.delete_folder_by_id_and_user_id(id, user.id, db=db)
|
||||
|
||||
for folder_id in folder_ids:
|
||||
if delete_contents:
|
||||
Chats.delete_chats_by_user_id_and_folder_id(user.id, folder_id)
|
||||
Chats.delete_chats_by_user_id_and_folder_id(
|
||||
user.id, folder_id, db=db
|
||||
)
|
||||
else:
|
||||
Chats.move_chats_by_user_id_and_folder_id(
|
||||
user.id, folder_id, None
|
||||
user.id, folder_id, None, db=db
|
||||
)
|
||||
|
||||
return True
|
||||
@@ -314,7 +341,7 @@ async def delete_folder_by_id(
|
||||
finally:
|
||||
# Get all subfolders
|
||||
subfolders = Folders.get_folders_by_parent_id_and_user_id(
|
||||
folder.id, user.id
|
||||
folder.id, user.id, db=db
|
||||
)
|
||||
folders.extend(subfolders)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user