diff --git a/backend/open_webui/models/access_grants.py b/backend/open_webui/models/access_grants.py index a15e4f3c5d..227621becd 100644 --- a/backend/open_webui/models/access_grants.py +++ b/backend/open_webui/models/access_grants.py @@ -456,6 +456,31 @@ class AccessGrantsTable: ) return [AccessGrantModel.model_validate(g) for g in grants] + def get_grants_by_resources( + self, + resource_type: str, + resource_ids: list[str], + db: Optional[Session] = None, + ) -> dict[str, list[AccessGrantModel]]: + """Batch-fetch grants for multiple resources. Returns {resource_id: [grants]}.""" + if not resource_ids: + return {} + with get_db_context(db) as db: + grants = ( + db.query(AccessGrant) + .filter( + AccessGrant.resource_type == resource_type, + AccessGrant.resource_id.in_(resource_ids), + ) + .all() + ) + result: dict[str, list[AccessGrantModel]] = { + rid: [] for rid in resource_ids + } + for g in grants: + result[g.resource_id].append(AccessGrantModel.model_validate(g)) + return result + def has_access( self, user_id: str, diff --git a/backend/open_webui/models/channels.py b/backend/open_webui/models/channels.py index 8a55da9345..2253d538c8 100644 --- a/backend/open_webui/models/channels.py +++ b/backend/open_webui/models/channels.py @@ -261,13 +261,19 @@ class ChannelTable: return AccessGrants.get_grants_by_resource("channel", channel_id, db=db) def _to_channel_model( - self, channel: Channel, db: Optional[Session] = None + self, + channel: Channel, + access_grants: Optional[list[AccessGrantModel]] = None, + db: Optional[Session] = None, ) -> ChannelModel: channel_data = ChannelModel.model_validate(channel).model_dump( exclude={"access_grants"} ) - access_grants = self._get_access_grants(channel_data["id"], db=db) - channel_data["access_grants"] = access_grants + channel_data["access_grants"] = ( + access_grants + if access_grants is not None + else self._get_access_grants(channel_data["id"], db=db) + ) return ChannelModel.model_validate(channel_data) def _collect_unique_user_ids( @@ -368,7 +374,18 @@ class ChannelTable: def get_channels(self, db: Optional[Session] = None) -> list[ChannelModel]: with get_db_context(db) as db: channels = db.query(Channel).all() - return [self._to_channel_model(channel, db=db) for channel in channels] + channel_ids = [channel.id for channel in channels] + grants_map = AccessGrants.get_grants_by_resources( + "channel", channel_ids, db=db + ) + return [ + self._to_channel_model( + channel, + access_grants=grants_map.get(channel.id, []), + db=db, + ) + for channel in channels + ] def _has_permission(self, db, query, filter: dict, permission: str = "read"): return AccessGrants.has_permission_filter( @@ -417,7 +434,16 @@ class ChannelTable: standard_channels = query.all() all_channels = membership_channels + standard_channels - return [self._to_channel_model(c, db=db) for c in all_channels] + channel_ids = [c.id for c in all_channels] + grants_map = AccessGrants.get_grants_by_resources( + "channel", channel_ids, db=db + ) + return [ + self._to_channel_model( + c, access_grants=grants_map.get(c.id, []), db=db + ) + for c in all_channels + ] def get_dm_channel_by_user_ids( self, user_ids: list[str], db: Optional[Session] = None @@ -724,7 +750,17 @@ class ChannelTable: ) channel_ids = [cf.channel_id for cf in channel_files] channels = db.query(Channel).filter(Channel.id.in_(channel_ids)).all() - return [self._to_channel_model(channel, db=db) for channel in channels] + grants_map = AccessGrants.get_grants_by_resources( + "channel", channel_ids, db=db + ) + return [ + self._to_channel_model( + channel, + access_grants=grants_map.get(channel.id, []), + db=db, + ) + for channel in channels + ] def get_channels_by_file_id_and_user_id( self, file_id: str, user_id: str, db: Optional[Session] = None diff --git a/backend/open_webui/models/knowledge.py b/backend/open_webui/models/knowledge.py index 1d21d5d910..f9defae6ef 100644 --- a/backend/open_webui/models/knowledge.py +++ b/backend/open_webui/models/knowledge.py @@ -144,13 +144,18 @@ class KnowledgeTable: return AccessGrants.get_grants_by_resource("knowledge", knowledge_id, db=db) def _to_knowledge_model( - self, knowledge: Knowledge, db: Optional[Session] = None + self, + knowledge: Knowledge, + access_grants: Optional[list[AccessGrantModel]] = None, + db: Optional[Session] = None, ) -> KnowledgeModel: knowledge_data = KnowledgeModel.model_validate(knowledge).model_dump( exclude={"access_grants"} ) - knowledge_data["access_grants"] = self._get_access_grants( - knowledge_data["id"], db=db + knowledge_data["access_grants"] = ( + access_grants + if access_grants is not None + else self._get_access_grants(knowledge_data["id"], db=db) ) return KnowledgeModel.model_validate(knowledge_data) @@ -192,9 +197,13 @@ class KnowledgeTable: db.query(Knowledge).order_by(Knowledge.updated_at.desc()).all() ) user_ids = list(set(knowledge.user_id for knowledge in all_knowledge)) + knowledge_ids = [knowledge.id for knowledge in all_knowledge] users = Users.get_users_by_user_ids(user_ids, db=db) if user_ids else [] users_dict = {user.id: user for user in users} + grants_map = AccessGrants.get_grants_by_resources( + "knowledge", knowledge_ids, db=db + ) knowledge_bases = [] for knowledge in all_knowledge: @@ -202,7 +211,11 @@ class KnowledgeTable: knowledge_bases.append( KnowledgeUserModel.model_validate( { - **self._to_knowledge_model(knowledge, db=db).model_dump(), + **self._to_knowledge_model( + knowledge, + access_grants=grants_map.get(knowledge.id, []), + db=db, + ).model_dump(), "user": user.model_dump() if user else None, } ) @@ -261,13 +274,22 @@ class KnowledgeTable: items = query.all() + knowledge_ids = [kb.id for kb, _ in items] + grants_map = AccessGrants.get_grants_by_resources( + "knowledge", knowledge_ids, db=db + ) + knowledge_bases = [] for knowledge_base, user in items: knowledge_bases.append( KnowledgeUserModel.model_validate( { **self._to_knowledge_model( - knowledge_base, db=db + knowledge_base, + access_grants=grants_map.get( + knowledge_base.id, [] + ), + db=db, ).model_dump(), "user": ( UserModel.model_validate(user).model_dump() @@ -440,8 +462,16 @@ class KnowledgeTable: .filter(KnowledgeFile.file_id == file_id) .all() ) + knowledge_ids = [k.id for k in knowledges] + grants_map = AccessGrants.get_grants_by_resources( + "knowledge", knowledge_ids, db=db + ) return [ - self._to_knowledge_model(knowledge, db=db) + self._to_knowledge_model( + knowledge, + access_grants=grants_map.get(knowledge.id, []), + db=db, + ) for knowledge in knowledges ] except Exception: diff --git a/backend/open_webui/models/models.py b/backend/open_webui/models/models.py index cfece00e35..25b135c87f 100755 --- a/backend/open_webui/models/models.py +++ b/backend/open_webui/models/models.py @@ -144,11 +144,20 @@ class ModelsTable: ) -> list[AccessGrantModel]: return AccessGrants.get_grants_by_resource("model", model_id, db=db) - def _to_model_model(self, model: Model, db: Optional[Session] = None) -> ModelModel: + def _to_model_model( + self, + model: Model, + access_grants: Optional[list[AccessGrantModel]] = None, + db: Optional[Session] = None, + ) -> ModelModel: model_data = ModelModel.model_validate(model).model_dump( exclude={"access_grants"} ) - model_data["access_grants"] = self._get_access_grants(model_data["id"], db=db) + model_data["access_grants"] = ( + access_grants + if access_grants is not None + else self._get_access_grants(model_data["id"], db=db) + ) return ModelModel.model_validate(model_data) def insert_new_model( @@ -181,8 +190,14 @@ class ModelsTable: def get_all_models(self, db: Optional[Session] = None) -> list[ModelModel]: with get_db_context(db) as db: + all_models = db.query(Model).all() + model_ids = [model.id for model in all_models] + grants_map = AccessGrants.get_grants_by_resources("model", model_ids, db=db) return [ - self._to_model_model(model, db=db) for model in db.query(Model).all() + self._to_model_model( + model, access_grants=grants_map.get(model.id, []), db=db + ) + for model in all_models ] def get_models(self, db: Optional[Session] = None) -> list[ModelUserResponse]: @@ -190,9 +205,11 @@ class ModelsTable: all_models = db.query(Model).filter(Model.base_model_id != None).all() user_ids = list(set(model.user_id for model in all_models)) + model_ids = [model.id for model in all_models] users = Users.get_users_by_user_ids(user_ids, db=db) if user_ids else [] users_dict = {user.id: user for user in users} + grants_map = AccessGrants.get_grants_by_resources("model", model_ids, db=db) models = [] for model in all_models: @@ -200,7 +217,11 @@ class ModelsTable: models.append( ModelUserResponse.model_validate( { - **self._to_model_model(model, db=db).model_dump(), + **self._to_model_model( + model, + access_grants=grants_map.get(model.id, []), + db=db, + ).model_dump(), "user": user.model_dump() if user else None, } ) @@ -209,9 +230,16 @@ class ModelsTable: def get_base_models(self, db: Optional[Session] = None) -> list[ModelModel]: with get_db_context(db) as db: + all_models = ( + db.query(Model).filter(Model.base_model_id == None).all() + ) + model_ids = [model.id for model in all_models] + grants_map = AccessGrants.get_grants_by_resources("model", model_ids, db=db) return [ - self._to_model_model(model, db=db) - for model in db.query(Model).filter(Model.base_model_id == None).all() + self._to_model_model( + model, access_grants=grants_map.get(model.id, []), db=db + ) + for model in all_models ] def get_models_by_user_id( @@ -325,11 +353,18 @@ class ModelsTable: items = query.all() + model_ids = [model.id for model, _ in items] + grants_map = AccessGrants.get_grants_by_resources("model", model_ids, db=db) + models = [] for model, user in items: models.append( ModelUserResponse( - **self._to_model_model(model, db=db).model_dump(), + **self._to_model_model( + model, + access_grants=grants_map.get(model.id, []), + db=db, + ).model_dump(), user=( UserResponse(**UserModel.model_validate(user).model_dump()) if user @@ -356,7 +391,18 @@ class ModelsTable: try: with get_db_context(db) as db: models = db.query(Model).filter(Model.id.in_(ids)).all() - return [self._to_model_model(model, db=db) for model in models] + model_ids = [model.id for model in models] + grants_map = AccessGrants.get_grants_by_resources( + "model", model_ids, db=db + ) + return [ + self._to_model_model( + model, + access_grants=grants_map.get(model.id, []), + db=db, + ) + for model in models + ] except Exception: return [] @@ -465,9 +511,18 @@ class ModelsTable: db.commit() + all_models = db.query(Model).all() + model_ids = [model.id for model in all_models] + grants_map = AccessGrants.get_grants_by_resources( + "model", model_ids, db=db + ) return [ - self._to_model_model(model, db=db) - for model in db.query(Model).all() + self._to_model_model( + model, + access_grants=grants_map.get(model.id, []), + db=db, + ) + for model in all_models ] except Exception as e: log.exception(f"Error syncing models for user {user_id}: {e}") diff --git a/backend/open_webui/models/notes.py b/backend/open_webui/models/notes.py index d17c749d1c..ff8a3ac635 100644 --- a/backend/open_webui/models/notes.py +++ b/backend/open_webui/models/notes.py @@ -93,9 +93,18 @@ class NoteTable: ) -> list[AccessGrantModel]: return AccessGrants.get_grants_by_resource("note", note_id, db=db) - def _to_note_model(self, note: Note, db: Optional[Session] = None) -> NoteModel: + def _to_note_model( + self, + note: Note, + access_grants: Optional[list[AccessGrantModel]] = None, + db: Optional[Session] = None, + ) -> NoteModel: note_data = NoteModel.model_validate(note).model_dump(exclude={"access_grants"}) - note_data["access_grants"] = self._get_access_grants(note_data["id"], db=db) + note_data["access_grants"] = ( + access_grants + if access_grants is not None + else self._get_access_grants(note_data["id"], db=db) + ) return NoteModel.model_validate(note_data) def _has_permission(self, db, query, filter: dict, permission: str = "read"): @@ -142,7 +151,14 @@ class NoteTable: if limit is not None: query = query.limit(limit) notes = query.all() - return [self._to_note_model(note, db=db) for note in notes] + note_ids = [note.id for note in notes] + grants_map = AccessGrants.get_grants_by_resources("note", note_ids, db=db) + return [ + self._to_note_model( + note, access_grants=grants_map.get(note.id, []), db=db + ) + for note in notes + ] def search_notes( self, @@ -227,11 +243,18 @@ class NoteTable: items = query.all() + note_ids = [note.id for note, _ in items] + grants_map = AccessGrants.get_grants_by_resources("note", note_ids, db=db) + notes = [] for note, user in items: notes.append( NoteUserResponse( - **self._to_note_model(note, db=db).model_dump(), + **self._to_note_model( + note, + access_grants=grants_map.get(note.id, []), + db=db, + ).model_dump(), user=( UserResponse(**UserModel.model_validate(user).model_dump()) if user @@ -266,7 +289,14 @@ class NoteTable: query = query.limit(limit) notes = query.all() - return [self._to_note_model(note, db=db) for note in notes] + note_ids = [note.id for note in notes] + grants_map = AccessGrants.get_grants_by_resources("note", note_ids, db=db) + return [ + self._to_note_model( + note, access_grants=grants_map.get(note.id, []), db=db + ) + for note in notes + ] def get_note_by_id( self, id: str, db: Optional[Session] = None diff --git a/backend/open_webui/models/prompts.py b/backend/open_webui/models/prompts.py index 80b7856f0f..dc83ee51dc 100644 --- a/backend/open_webui/models/prompts.py +++ b/backend/open_webui/models/prompts.py @@ -97,12 +97,19 @@ class PromptsTable: return AccessGrants.get_grants_by_resource("prompt", prompt_id, db=db) def _to_prompt_model( - self, prompt: Prompt, db: Optional[Session] = None + self, + prompt: Prompt, + access_grants: Optional[list[AccessGrantModel]] = None, + db: Optional[Session] = None, ) -> PromptModel: prompt_data = PromptModel.model_validate(prompt).model_dump( exclude={"access_grants"} ) - prompt_data["access_grants"] = self._get_access_grants(prompt_data["id"], db=db) + prompt_data["access_grants"] = ( + access_grants + if access_grants is not None + else self._get_access_grants(prompt_data["id"], db=db) + ) return PromptModel.model_validate(prompt_data) def insert_new_prompt( @@ -206,9 +213,11 @@ class PromptsTable: ) user_ids = list(set(prompt.user_id for prompt in all_prompts)) + prompt_ids = [prompt.id for prompt in all_prompts] users = Users.get_users_by_user_ids(user_ids, db=db) if user_ids else [] users_dict = {user.id: user for user in users} + grants_map = AccessGrants.get_grants_by_resources("prompt", prompt_ids, db=db) prompts = [] for prompt in all_prompts: @@ -216,7 +225,11 @@ class PromptsTable: prompts.append( PromptUserResponse.model_validate( { - **self._to_prompt_model(prompt, db=db).model_dump(), + **self._to_prompt_model( + prompt, + access_grants=grants_map.get(prompt.id, []), + db=db, + ).model_dump(), "user": user.model_dump() if user else None, } ) @@ -329,11 +342,18 @@ class PromptsTable: items = query.all() + prompt_ids = [prompt.id for prompt, _ in items] + grants_map = AccessGrants.get_grants_by_resources("prompt", prompt_ids, db=db) + prompts = [] for prompt, user in items: prompts.append( PromptUserResponse( - **self._to_prompt_model(prompt, db=db).model_dump(), + **self._to_prompt_model( + prompt, + access_grants=grants_map.get(prompt.id, []), + db=db, + ).model_dump(), user=( UserResponse(**UserModel.model_validate(user).model_dump()) if user diff --git a/backend/open_webui/models/skills.py b/backend/open_webui/models/skills.py index 1262830153..13e62c6f19 100644 --- a/backend/open_webui/models/skills.py +++ b/backend/open_webui/models/skills.py @@ -110,11 +110,20 @@ class SkillsTable: ) -> list[AccessGrantModel]: return AccessGrants.get_grants_by_resource("skill", skill_id, db=db) - def _to_skill_model(self, skill: Skill, db: Optional[Session] = None) -> SkillModel: + def _to_skill_model( + self, + skill: Skill, + access_grants: Optional[list[AccessGrantModel]] = None, + db: Optional[Session] = None, + ) -> SkillModel: skill_data = SkillModel.model_validate(skill).model_dump( exclude={"access_grants"} ) - skill_data["access_grants"] = self._get_access_grants(skill_data["id"], db=db) + skill_data["access_grants"] = ( + access_grants + if access_grants is not None + else self._get_access_grants(skill_data["id"], db=db) + ) return SkillModel.model_validate(skill_data) def insert_new_skill( @@ -172,9 +181,11 @@ class SkillsTable: all_skills = db.query(Skill).order_by(Skill.updated_at.desc()).all() user_ids = list(set(skill.user_id for skill in all_skills)) + skill_ids = [skill.id for skill in all_skills] users = Users.get_users_by_user_ids(user_ids, db=db) if user_ids else [] users_dict = {user.id: user for user in users} + grants_map = AccessGrants.get_grants_by_resources("skill", skill_ids, db=db) skills = [] for skill in all_skills: @@ -182,7 +193,11 @@ class SkillsTable: skills.append( SkillUserModel.model_validate( { - **self._to_skill_model(skill, db=db).model_dump(), + **self._to_skill_model( + skill, + access_grants=grants_map.get(skill.id, []), + db=db, + ).model_dump(), "user": user.model_dump() if user else None, } ) @@ -267,11 +282,18 @@ class SkillsTable: items = query.all() + skill_ids = [skill.id for skill, _ in items] + grants_map = AccessGrants.get_grants_by_resources("skill", skill_ids, db=db) + skills = [] for skill, user in items: skills.append( SkillUserResponse( - **self._to_skill_model(skill, db=db).model_dump(), + **self._to_skill_model( + skill, + access_grants=grants_map.get(skill.id, []), + db=db, + ).model_dump(), user=( UserResponse( **UserModel.model_validate(user).model_dump() diff --git a/backend/open_webui/models/tools.py b/backend/open_webui/models/tools.py index eaac4c385d..7b66bb7b40 100644 --- a/backend/open_webui/models/tools.py +++ b/backend/open_webui/models/tools.py @@ -100,9 +100,18 @@ class ToolsTable: ) -> list[AccessGrantModel]: return AccessGrants.get_grants_by_resource("tool", tool_id, db=db) - def _to_tool_model(self, tool: Tool, db: Optional[Session] = None) -> ToolModel: + def _to_tool_model( + self, + tool: Tool, + access_grants: Optional[list[AccessGrantModel]] = None, + db: Optional[Session] = None, + ) -> ToolModel: tool_data = ToolModel.model_validate(tool).model_dump(exclude={"access_grants"}) - tool_data["access_grants"] = self._get_access_grants(tool_data["id"], db=db) + tool_data["access_grants"] = ( + access_grants + if access_grants is not None + else self._get_access_grants(tool_data["id"], db=db) + ) return ToolModel.model_validate(tool_data) def insert_new_tool( @@ -152,9 +161,11 @@ class ToolsTable: all_tools = db.query(Tool).order_by(Tool.updated_at.desc()).all() user_ids = list(set(tool.user_id for tool in all_tools)) + tool_ids = [tool.id for tool in all_tools] users = Users.get_users_by_user_ids(user_ids, db=db) if user_ids else [] users_dict = {user.id: user for user in users} + grants_map = AccessGrants.get_grants_by_resources("tool", tool_ids, db=db) tools = [] for tool in all_tools: @@ -162,7 +173,11 @@ class ToolsTable: tools.append( ToolUserModel.model_validate( { - **self._to_tool_model(tool, db=db).model_dump(), + **self._to_tool_model( + tool, + access_grants=grants_map.get(tool.id, []), + db=db, + ).model_dump(), "user": user.model_dump() if user else None, } )