mirror of
https://github.com/open-webui/open-webui.git
synced 2026-08-23 14:04:49 -06:00
Classic298
6db64c4855
Opening the shared folder list fetched every shared folder in its own query, fetched a chunk of them a second time to walk their children, and looked up each distinct owner separately. With forty folders shared with a user that is over a hundred queries before any subtree work starts. The folders and their owners now come back in one query each, and the inheritance pass reuses the rows already in hand. Both folder listings also gained an explicit order: the sidebar merges shared subfolders in response order without sorting them, and neither query had an ORDER BY, so on Postgres a folder rename could reshuffle its siblings. Measured with forty shared folders and no subtrees: 181 queries and ~105 ms before, 92 and ~66 ms after. With subtrees attached, 203 folders in total, it is 341 queries before against 252 after; the remainder is the recursive child walk, which this change deliberately leaves alone. The returned set, permissions and owner names are unchanged, including for a grant pointing at a deleted folder row, a folder the caller owns that is also shared with them, a folder whose owner record is gone, and a child folder that is itself directly shared.
464 lines
16 KiB
Python
464 lines
16 KiB
Python
import logging
|
|
import re
|
|
import time
|
|
import uuid
|
|
from typing import Optional
|
|
|
|
from open_webui.internal.db import Base, JSONField, get_async_db_context
|
|
from pydantic import BaseModel, ConfigDict
|
|
from sqlalchemy import JSON, BigInteger, Boolean, Column, Text, delete, func, select, or_, and_
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
####################
|
|
# Folder DB Schema
|
|
# Let every room in this house shelter someone who needs it,
|
|
# and let no chamber stand empty while there is want.
|
|
####################
|
|
|
|
|
|
class Folder(Base):
|
|
__tablename__ = 'folder'
|
|
id = Column(Text, primary_key=True, unique=True)
|
|
parent_id = Column(Text, nullable=True)
|
|
user_id = Column(Text)
|
|
name = Column(Text)
|
|
items = Column(JSON, nullable=True)
|
|
meta = Column(JSON, nullable=True)
|
|
data = Column(JSON, nullable=True)
|
|
is_expanded = Column(Boolean, default=False)
|
|
created_at = Column(BigInteger)
|
|
updated_at = Column(BigInteger)
|
|
|
|
|
|
class FolderModel(BaseModel):
|
|
id: str
|
|
parent_id: Optional[str] = None
|
|
user_id: str
|
|
name: str
|
|
items: Optional[dict] = None
|
|
meta: Optional[dict] = None
|
|
data: Optional[dict] = None
|
|
is_expanded: bool = False
|
|
created_at: int
|
|
updated_at: int
|
|
|
|
model_config = ConfigDict(from_attributes=True)
|
|
|
|
|
|
class FolderMetadataResponse(BaseModel):
|
|
icon: Optional[str] = None
|
|
|
|
|
|
class FolderNameIdResponse(BaseModel):
|
|
id: str
|
|
name: str
|
|
meta: Optional[FolderMetadataResponse] = None
|
|
parent_id: Optional[str] = None
|
|
is_expanded: bool = False
|
|
unread_count: int = 0
|
|
created_at: int
|
|
updated_at: int
|
|
|
|
|
|
class SharedFolderResponse(BaseModel):
|
|
id: str
|
|
name: str
|
|
parent_id: Optional[str] = None
|
|
user_id: str
|
|
owner_name: Optional[str] = None
|
|
permission: str = 'read'
|
|
access_grants: list = []
|
|
is_expanded: bool = False
|
|
meta: Optional[dict] = None
|
|
created_at: int
|
|
updated_at: int
|
|
|
|
|
|
####################
|
|
# Forms
|
|
####################
|
|
|
|
|
|
class FolderForm(BaseModel):
|
|
name: str
|
|
data: Optional[dict] = None
|
|
meta: Optional[dict] = None
|
|
parent_id: Optional[str] = None
|
|
model_config = ConfigDict(extra='forbid')
|
|
|
|
|
|
class FolderUpdateForm(BaseModel):
|
|
name: Optional[str] = None
|
|
data: Optional[dict] = None
|
|
meta: Optional[dict] = None
|
|
model_config = ConfigDict(extra='forbid')
|
|
|
|
|
|
class FolderTable:
|
|
async def insert_new_folder(
|
|
self,
|
|
user_id: str,
|
|
form_data: FolderForm,
|
|
parent_id: Optional[str] = None,
|
|
db: Optional[AsyncSession] = None,
|
|
) -> Optional[FolderModel]:
|
|
async with get_async_db_context(db) as db:
|
|
id = str(uuid.uuid4())
|
|
folder = FolderModel(
|
|
**{
|
|
'id': id,
|
|
'user_id': user_id,
|
|
**(form_data.model_dump(exclude_unset=True) or {}),
|
|
'parent_id': parent_id,
|
|
'created_at': int(time.time()),
|
|
'updated_at': int(time.time()),
|
|
}
|
|
)
|
|
try:
|
|
result = Folder(**folder.model_dump())
|
|
db.add(result)
|
|
await db.commit()
|
|
await db.refresh(result)
|
|
if result:
|
|
return FolderModel.model_validate(result)
|
|
else:
|
|
return None
|
|
except Exception as e:
|
|
log.exception(f'Error inserting a new folder: {e}')
|
|
return None
|
|
|
|
async def get_folder_by_id_and_user_id(
|
|
self, id: str, user_id: str, db: Optional[AsyncSession] = None
|
|
) -> Optional[FolderModel]:
|
|
try:
|
|
async with get_async_db_context(db) as db:
|
|
result = await db.execute(select(Folder).filter_by(id=id, user_id=user_id))
|
|
folder = result.scalars().first()
|
|
|
|
if not folder:
|
|
return None
|
|
|
|
return FolderModel.model_validate(folder)
|
|
except Exception:
|
|
return None
|
|
|
|
async def get_folder_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[FolderModel]:
|
|
"""Fetch folder by ID only (no user_id filter). Used for shared access."""
|
|
try:
|
|
async with get_async_db_context(db) as db:
|
|
result = await db.execute(select(Folder).filter_by(id=id))
|
|
folder = result.scalars().first()
|
|
if not folder:
|
|
return None
|
|
return FolderModel.model_validate(folder)
|
|
except Exception:
|
|
return None
|
|
|
|
async def get_folders_by_ids(self, ids: list[str], db: AsyncSession | None = None) -> list[FolderModel]:
|
|
async with get_async_db_context(db) as db:
|
|
result = await db.execute(select(Folder).filter(Folder.id.in_(ids)).order_by(Folder.updated_at.desc()))
|
|
return [FolderModel.model_validate(folder) for folder in result.scalars().all()]
|
|
|
|
async def get_shared_folder_ids_for_user(
|
|
self, user_id: str, user_group_ids: set[str], db: Optional[AsyncSession] = None
|
|
) -> dict[str, str]:
|
|
"""
|
|
Returns {folder_id: highest_permission} for all folders shared with user.
|
|
Checks direct user grants, group grants, and public (user:*) grants.
|
|
"""
|
|
from open_webui.models.access_grants import AccessGrant
|
|
|
|
async with get_async_db_context(db) as db:
|
|
conditions = [
|
|
and_(AccessGrant.principal_type == 'user', AccessGrant.principal_id == '*'),
|
|
and_(AccessGrant.principal_type == 'user', AccessGrant.principal_id == user_id),
|
|
]
|
|
if user_group_ids:
|
|
conditions.append(
|
|
and_(AccessGrant.principal_type == 'group', AccessGrant.principal_id.in_(user_group_ids))
|
|
)
|
|
result = await db.execute(
|
|
select(AccessGrant).filter(
|
|
AccessGrant.resource_type == 'folder',
|
|
or_(*conditions),
|
|
)
|
|
)
|
|
grants = result.scalars().all()
|
|
|
|
# Build {folder_id: highest_permission} ('write' > 'read')
|
|
folder_perms = {}
|
|
for g in grants:
|
|
existing = folder_perms.get(g.resource_id)
|
|
if existing != 'write':
|
|
folder_perms[g.resource_id] = g.permission
|
|
return folder_perms
|
|
|
|
async def get_children_folders_by_id_and_user_id(
|
|
self, id: str, user_id: str, db: Optional[AsyncSession] = None
|
|
) -> Optional[list[FolderModel]]:
|
|
try:
|
|
async with get_async_db_context(db) as db:
|
|
folders = []
|
|
|
|
async def get_children(folder):
|
|
children = await self.get_folders_by_parent_id_and_user_id(folder.id, user_id, db=db)
|
|
for child in children:
|
|
await get_children(child)
|
|
folders.append(child)
|
|
|
|
result = await db.execute(select(Folder).filter_by(id=id, user_id=user_id))
|
|
folder = result.scalars().first()
|
|
if not folder:
|
|
return None
|
|
|
|
await get_children(folder)
|
|
return folders
|
|
except Exception:
|
|
return None
|
|
|
|
async def get_folders_by_user_id(self, user_id: str, db: Optional[AsyncSession] = None) -> list[FolderModel]:
|
|
async with get_async_db_context(db) as db:
|
|
result = await db.execute(select(Folder).filter_by(user_id=user_id))
|
|
return [FolderModel.model_validate(folder) for folder in result.scalars().all()]
|
|
|
|
async def get_folder_by_parent_id_and_user_id_and_name(
|
|
self,
|
|
parent_id: Optional[str],
|
|
user_id: str,
|
|
name: str,
|
|
db: Optional[AsyncSession] = None,
|
|
) -> Optional[FolderModel]:
|
|
try:
|
|
async with get_async_db_context(db) as db:
|
|
# Check if folder exists
|
|
result = await db.execute(
|
|
select(Folder)
|
|
.filter_by(parent_id=parent_id, user_id=user_id)
|
|
.filter(func.lower(Folder.name) == func.lower(name))
|
|
)
|
|
folder = result.scalars().first()
|
|
|
|
if not folder:
|
|
return None
|
|
|
|
return FolderModel.model_validate(folder)
|
|
except Exception as e:
|
|
log.error(f'get_folder_by_parent_id_and_user_id_and_name: {e}')
|
|
return None
|
|
|
|
async def get_folders_by_parent_id_and_user_id(
|
|
self, parent_id: Optional[str], user_id: str, db: Optional[AsyncSession] = None
|
|
) -> list[FolderModel]:
|
|
async with get_async_db_context(db) as db:
|
|
result = await db.execute(
|
|
select(Folder).filter_by(parent_id=parent_id, user_id=user_id).order_by(Folder.updated_at.desc())
|
|
)
|
|
return [FolderModel.model_validate(folder) for folder in result.scalars().all()]
|
|
|
|
async def get_folder_ids_by_id_and_user_id_in_subtree(
|
|
self, id: str, user_id: str, db: Optional[AsyncSession] = None
|
|
) -> list[str]:
|
|
async with get_async_db_context(db) as db:
|
|
result = await db.execute(select(Folder).filter_by(id=id, user_id=user_id))
|
|
folder = result.scalars().first()
|
|
if not folder:
|
|
return []
|
|
|
|
folder_ids = [folder.id]
|
|
folders = [FolderModel.model_validate(folder)]
|
|
while folders:
|
|
current_folder = folders.pop()
|
|
children = await self.get_folders_by_parent_id_and_user_id(current_folder.id, user_id, db=db)
|
|
folder_ids.extend(child.id for child in children)
|
|
folders.extend(children)
|
|
|
|
return folder_ids
|
|
|
|
async def update_folder_parent_id_by_id_and_user_id(
|
|
self,
|
|
id: str,
|
|
user_id: str,
|
|
parent_id: str,
|
|
db: Optional[AsyncSession] = None,
|
|
) -> Optional[FolderModel]:
|
|
try:
|
|
async with get_async_db_context(db) as db:
|
|
result = await db.execute(select(Folder).filter_by(id=id, user_id=user_id))
|
|
folder = result.scalars().first()
|
|
|
|
if not folder:
|
|
return None
|
|
|
|
folder.parent_id = parent_id
|
|
folder.updated_at = int(time.time())
|
|
|
|
await db.commit()
|
|
|
|
return FolderModel.model_validate(folder)
|
|
except Exception as e:
|
|
log.error(f'update_folder: {e}')
|
|
return
|
|
|
|
async def update_folder_by_id_and_user_id(
|
|
self,
|
|
id: str,
|
|
user_id: str,
|
|
form_data: FolderUpdateForm,
|
|
db: Optional[AsyncSession] = None,
|
|
) -> Optional[FolderModel]:
|
|
try:
|
|
async with get_async_db_context(db) as db:
|
|
result = await db.execute(select(Folder).filter_by(id=id, user_id=user_id))
|
|
folder = result.scalars().first()
|
|
|
|
if not folder:
|
|
return None
|
|
|
|
form_data = form_data.model_dump(exclude_unset=True)
|
|
|
|
existing_result = await db.execute(
|
|
select(Folder).filter_by(
|
|
name=form_data.get('name'),
|
|
parent_id=folder.parent_id,
|
|
user_id=user_id,
|
|
)
|
|
)
|
|
existing_folder = existing_result.scalars().first()
|
|
|
|
if existing_folder and existing_folder.id != id:
|
|
return None
|
|
|
|
folder.name = form_data.get('name', folder.name)
|
|
if 'data' in form_data:
|
|
folder.data = {
|
|
**(folder.data or {}),
|
|
**form_data['data'],
|
|
}
|
|
|
|
if 'meta' in form_data:
|
|
folder.meta = {
|
|
**(folder.meta or {}),
|
|
**form_data['meta'],
|
|
}
|
|
|
|
folder.updated_at = int(time.time())
|
|
await db.commit()
|
|
|
|
return FolderModel.model_validate(folder)
|
|
except Exception as e:
|
|
log.error(f'update_folder: {e}')
|
|
return
|
|
|
|
async def update_folder_is_expanded_by_id_and_user_id(
|
|
self, id: str, user_id: str, is_expanded: bool, db: Optional[AsyncSession] = None
|
|
) -> Optional[FolderModel]:
|
|
try:
|
|
async with get_async_db_context(db) as db:
|
|
result = await db.execute(select(Folder).filter_by(id=id, user_id=user_id))
|
|
folder = result.scalars().first()
|
|
|
|
if not folder:
|
|
return None
|
|
|
|
folder.is_expanded = is_expanded
|
|
folder.updated_at = int(time.time())
|
|
|
|
await db.commit()
|
|
|
|
return FolderModel.model_validate(folder)
|
|
except Exception as e:
|
|
log.error(f'update_folder: {e}')
|
|
return
|
|
|
|
async def delete_folder_by_id_and_user_id(
|
|
self, id: str, user_id: str, db: Optional[AsyncSession] = None
|
|
) -> list[str]:
|
|
try:
|
|
folder_ids = []
|
|
async with get_async_db_context(db) as db:
|
|
result = await db.execute(select(Folder).filter_by(id=id, user_id=user_id))
|
|
folder = result.scalars().first()
|
|
if not folder:
|
|
return folder_ids
|
|
|
|
folder_ids.append(folder.id)
|
|
|
|
# Delete all children folders
|
|
async def delete_children(folder):
|
|
folder_children = await self.get_folders_by_parent_id_and_user_id(folder.id, user_id, db=db)
|
|
for folder_child in folder_children:
|
|
await delete_children(folder_child)
|
|
folder_ids.append(folder_child.id)
|
|
|
|
child_result = await db.execute(select(Folder).filter_by(id=folder_child.id))
|
|
child_folder = child_result.scalars().first()
|
|
await db.delete(child_folder)
|
|
await db.commit()
|
|
|
|
await delete_children(folder)
|
|
await db.delete(folder)
|
|
await db.commit()
|
|
return folder_ids
|
|
except Exception as e:
|
|
log.error(f'delete_folder: {e}')
|
|
return []
|
|
|
|
def normalize_folder_name(self, name: str) -> str:
|
|
# Replace _ and space with a single space, lower case, collapse multiple spaces
|
|
name = re.sub(r'[\s_]+', ' ', name)
|
|
return name.strip().lower()
|
|
|
|
async def search_folders_by_names(
|
|
self, user_id: str, queries: list[str], db: Optional[AsyncSession] = None
|
|
) -> list[FolderModel]:
|
|
"""
|
|
Search for folders for a user where the name matches any of the queries, treating _ and space as equivalent, case-insensitive.
|
|
"""
|
|
normalized_queries = [self.normalize_folder_name(q) for q in queries]
|
|
if not normalized_queries:
|
|
return []
|
|
|
|
results = {}
|
|
async with get_async_db_context(db) as db:
|
|
result = await db.execute(select(Folder).filter_by(user_id=user_id))
|
|
folders = result.scalars().all()
|
|
for folder in folders:
|
|
if self.normalize_folder_name(folder.name) in normalized_queries:
|
|
results[folder.id] = FolderModel.model_validate(folder)
|
|
|
|
# get children folders
|
|
children = await self.get_children_folders_by_id_and_user_id(folder.id, user_id, db=db)
|
|
if children:
|
|
for child in children:
|
|
results[child.id] = child
|
|
|
|
# Return the results as a list
|
|
if not results:
|
|
return []
|
|
else:
|
|
results = list(results.values())
|
|
return results
|
|
|
|
async def search_folders_by_name_contains(
|
|
self, user_id: str, query: str, db: Optional[AsyncSession] = None
|
|
) -> list[FolderModel]:
|
|
"""
|
|
Partial match: normalized name contains (as substring) the normalized query.
|
|
"""
|
|
normalized_query = self.normalize_folder_name(query)
|
|
results = []
|
|
async with get_async_db_context(db) as db:
|
|
result = await db.execute(select(Folder).filter_by(user_id=user_id))
|
|
folders = result.scalars().all()
|
|
for folder in folders:
|
|
norm_name = self.normalize_folder_name(folder.name)
|
|
if normalized_query in norm_name:
|
|
results.append(FolderModel.model_validate(folder))
|
|
return results
|
|
|
|
|
|
Folders = FolderTable()
|