Emil_Shanaty/open-webui-ai4me
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refac/enh: db session sharing

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2025-12-28 22:00:44 +04:00
parent d4de26bd05
commit 2041ab483e
20 changed files with 600 additions and 562 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/models/folders.py
+29 -27
View File
@@ -7,8 +7,9 @@ import re
from pydantic import BaseModel, ConfigDict
from sqlalchemy import BigInteger, Column, Text, JSON, Boolean, func
from sqlalchemy.orm import Session
from open_webui.internal.db import Base, get_db
from open_webui.internal.db import Base, JSONField, get_db, get_db_context
log = logging.getLogger(__name__)
@@ -83,9 +84,9 @@ class FolderUpdateForm(BaseModel):
class FolderTable:
def insert_new_folder(
self, user_id: str, form_data: FolderForm, parent_id: Optional[str] = None
self, user_id: str, form_data: FolderForm, parent_id: Optional[str] = None, db: Optional[Session] = None
) -> Optional[FolderModel]:
with get_db() as db:
with get_db_context(db) as db:
id = str(uuid.uuid4())
folder = FolderModel(
**{
@@ -111,10 +112,10 @@ class FolderTable:
return None
def get_folder_by_id_and_user_id(
self, id: str, user_id: str
self, id: str, user_id: str, db: Optional[Session] = None
) -> Optional[FolderModel]:
try:
with get_db() as db:
with get_db_context(db) as db:
folder = db.query(Folder).filter_by(id=id, user_id=user_id).first()
if not folder:
@@ -125,15 +126,15 @@ class FolderTable:
return None
def get_children_folders_by_id_and_user_id(
self, id: str, user_id: str
self, id: str, user_id: str, db: Optional[Session] = None
) -> Optional[list[FolderModel]]:
try:
with get_db() as db:
with get_db_context(db) as db:
folders = []
def get_children(folder):
children = self.get_folders_by_parent_id_and_user_id(
folder.id, user_id
folder.id, user_id, db=db
)
for child in children:
get_children(child)
@@ -148,18 +149,18 @@ class FolderTable:
except Exception:
return None
def get_folders_by_user_id(self, user_id: str) -> list[FolderModel]:
with get_db() as db:
def get_folders_by_user_id(self, user_id: str, db: Optional[Session] = None) -> list[FolderModel]:
with get_db_context(db) as db:
return [
FolderModel.model_validate(folder)
for folder in db.query(Folder).filter_by(user_id=user_id).all()
]
def get_folder_by_parent_id_and_user_id_and_name(
self, parent_id: Optional[str], user_id: str, name: str
self, parent_id: Optional[str], user_id: str, name: str, db: Optional[Session] = None
) -> Optional[FolderModel]:
try:
with get_db() as db:
with get_db_context(db) as db:
# Check if folder exists
folder = (
db.query(Folder)
@@ -177,9 +178,9 @@ class FolderTable:
return None
def get_folders_by_parent_id_and_user_id(
self, parent_id: Optional[str], user_id: str
self, parent_id: Optional[str], user_id: str, db: Optional[Session] = None
) -> list[FolderModel]:
with get_db() as db:
with get_db_context(db) as db:
return [
FolderModel.model_validate(folder)
for folder in db.query(Folder)
@@ -192,9 +193,10 @@ class FolderTable:
id: str,
user_id: str,
parent_id: str,
db: Optional[Session] = None,
) -> Optional[FolderModel]:
try:
with get_db() as db:
with get_db_context(db) as db:
folder = db.query(Folder).filter_by(id=id, user_id=user_id).first()
if not folder:
@@ -211,10 +213,10 @@ class FolderTable:
return
def update_folder_by_id_and_user_id(
self, id: str, user_id: str, form_data: FolderUpdateForm
self, id: str, user_id: str, form_data: FolderUpdateForm, db: Optional[Session] = None
) -> Optional[FolderModel]:
try:
with get_db() as db:
with get_db_context(db) as db:
folder = db.query(Folder).filter_by(id=id, user_id=user_id).first()
if not folder:
@@ -257,10 +259,10 @@ class FolderTable:
return
def update_folder_is_expanded_by_id_and_user_id(
self, id: str, user_id: str, is_expanded: bool
self, id: str, user_id: str, is_expanded: bool, db: Optional[Session] = None
) -> Optional[FolderModel]:
try:
with get_db() as db:
with get_db_context(db) as db:
folder = db.query(Folder).filter_by(id=id, user_id=user_id).first()
if not folder:
@@ -276,10 +278,10 @@ class FolderTable:
log.error(f"update_folder: {e}")
return
def delete_folder_by_id_and_user_id(self, id: str, user_id: str) -> list[str]:
def delete_folder_by_id_and_user_id(self, id: str, user_id: str, db: Optional[Session] = None) -> list[str]:
try:
folder_ids = []
with get_db() as db:
with get_db_context(db) as db:
folder = db.query(Folder).filter_by(id=id, user_id=user_id).first()
if not folder:
return folder_ids
@@ -289,7 +291,7 @@ class FolderTable:
# Delete all children folders
def delete_children(folder):
folder_children = self.get_folders_by_parent_id_and_user_id(
folder.id, user_id
folder.id, user_id, db=db
)
for folder_child in folder_children:
@@ -314,7 +316,7 @@ class FolderTable:
return name.strip().lower()
def search_folders_by_names(
self, user_id: str, queries: list[str]
self, user_id: str, queries: list[str], db: Optional[Session] = None
) -> list[FolderModel]:
"""
Search for folders for a user where the name matches any of the queries, treating _ and space as equivalent, case-insensitive.
@@ -324,7 +326,7 @@ class FolderTable:
return []
results = {}
with get_db() as db:
with get_db_context(db) as db:
folders = db.query(Folder).filter_by(user_id=user_id).all()
for folder in folders:
if self.normalize_folder_name(folder.name) in normalized_queries:
@@ -332,7 +334,7 @@ class FolderTable:
# get children folders
children = self.get_children_folders_by_id_and_user_id(
folder.id, user_id
folder.id, user_id, db=db
)
for child in children:
results[child.id] = child
@@ -345,14 +347,14 @@ class FolderTable:
return results
def search_folders_by_name_contains(
self, user_id: str, query: str
self, user_id: str, query: str, db: Optional[Session] = None
) -> list[FolderModel]:
"""
Partial match: normalized name contains (as substring) the normalized query.
"""
normalized_query = self.normalize_folder_name(query)
results = []
with get_db() as db:
with get_db_context(db) as db:
folders = db.query(Folder).filter_by(user_id=user_id).all()
for folder in folders:
norm_name = self.normalize_folder_name(folder.name)