refac
This commit is contained in:
@@ -15,9 +15,10 @@ from open_webui.models.files import (
|
||||
)
|
||||
from open_webui.models.groups import Groups
|
||||
from open_webui.models.users import User, UserModel, Users, UserResponse
|
||||
from open_webui.models.access_grants import AccessGrantModel, AccessGrants
|
||||
|
||||
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
from sqlalchemy import (
|
||||
BigInteger,
|
||||
Column,
|
||||
@@ -29,9 +30,6 @@ from sqlalchemy import (
|
||||
or_,
|
||||
)
|
||||
|
||||
from open_webui.utils.access_control import has_access
|
||||
from open_webui.utils.db.access_control import has_permission
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
@@ -50,22 +48,6 @@ class Knowledge(Base):
|
||||
description = Column(Text)
|
||||
|
||||
meta = Column(JSON, nullable=True)
|
||||
access_control = Column(JSON, nullable=True) # Controls data access levels.
|
||||
# Defines access control rules for this entry.
|
||||
# - `None`: Public access, available to all users with the "user" role.
|
||||
# - `{}`: Private access, restricted exclusively to the owner.
|
||||
# - Custom permissions: Specific access control for reading and writing;
|
||||
# Can specify group or user-level restrictions:
|
||||
# {
|
||||
# "read": {
|
||||
# "group_ids": ["group_id1", "group_id2"],
|
||||
# "user_ids": ["user_id1", "user_id2"]
|
||||
# },
|
||||
# "write": {
|
||||
# "group_ids": ["group_id1", "group_id2"],
|
||||
# "user_ids": ["user_id1", "user_id2"]
|
||||
# }
|
||||
# }
|
||||
|
||||
created_at = Column(BigInteger)
|
||||
updated_at = Column(BigInteger)
|
||||
@@ -82,7 +64,7 @@ class KnowledgeModel(BaseModel):
|
||||
|
||||
meta: Optional[dict] = None
|
||||
|
||||
access_control: Optional[dict] = None
|
||||
access_grants: list[AccessGrantModel] = Field(default_factory=list)
|
||||
|
||||
created_at: int # timestamp in epoch
|
||||
updated_at: int # timestamp in epoch
|
||||
@@ -139,7 +121,7 @@ class KnowledgeUserResponse(KnowledgeUserModel):
|
||||
class KnowledgeForm(BaseModel):
|
||||
name: str
|
||||
description: str
|
||||
access_control: Optional[dict] = None
|
||||
access_grants: Optional[list[dict]] = None
|
||||
|
||||
|
||||
class FileUserResponse(FileModelResponse):
|
||||
@@ -157,27 +139,47 @@ class KnowledgeFileListResponse(BaseModel):
|
||||
|
||||
|
||||
class KnowledgeTable:
|
||||
def _get_access_grants(
|
||||
self, knowledge_id: str, db: Optional[Session] = None
|
||||
) -> list[AccessGrantModel]:
|
||||
return AccessGrants.get_grants_by_resource("knowledge", knowledge_id, db=db)
|
||||
|
||||
def _to_knowledge_model(
|
||||
self, knowledge: Knowledge, 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
|
||||
)
|
||||
return KnowledgeModel.model_validate(knowledge_data)
|
||||
|
||||
def insert_new_knowledge(
|
||||
self, user_id: str, form_data: KnowledgeForm, db: Optional[Session] = None
|
||||
) -> Optional[KnowledgeModel]:
|
||||
with get_db_context(db) as db:
|
||||
knowledge = KnowledgeModel(
|
||||
**{
|
||||
**form_data.model_dump(),
|
||||
**form_data.model_dump(exclude={"access_grants"}),
|
||||
"id": str(uuid.uuid4()),
|
||||
"user_id": user_id,
|
||||
"created_at": int(time.time()),
|
||||
"updated_at": int(time.time()),
|
||||
"access_grants": [],
|
||||
}
|
||||
)
|
||||
|
||||
try:
|
||||
result = Knowledge(**knowledge.model_dump())
|
||||
result = Knowledge(**knowledge.model_dump(exclude={"access_grants"}))
|
||||
db.add(result)
|
||||
db.commit()
|
||||
db.refresh(result)
|
||||
AccessGrants.set_access_grants(
|
||||
"knowledge", result.id, form_data.access_grants, db=db
|
||||
)
|
||||
if result:
|
||||
return KnowledgeModel.model_validate(result)
|
||||
return self._to_knowledge_model(result, db=db)
|
||||
else:
|
||||
return None
|
||||
except Exception:
|
||||
@@ -201,7 +203,7 @@ class KnowledgeTable:
|
||||
knowledge_bases.append(
|
||||
KnowledgeUserModel.model_validate(
|
||||
{
|
||||
**KnowledgeModel.model_validate(knowledge).model_dump(),
|
||||
**self._to_knowledge_model(knowledge, db=db).model_dump(),
|
||||
"user": user.model_dump() if user else None,
|
||||
}
|
||||
)
|
||||
@@ -241,7 +243,14 @@ class KnowledgeTable:
|
||||
elif view_option == "shared":
|
||||
query = query.filter(Knowledge.user_id != user_id)
|
||||
|
||||
query = has_permission(db, Knowledge, query, filter)
|
||||
query = AccessGrants.has_permission_filter(
|
||||
db=db,
|
||||
query=query,
|
||||
DocumentModel=Knowledge,
|
||||
filter=filter,
|
||||
resource_type="knowledge",
|
||||
permission="read",
|
||||
)
|
||||
|
||||
query = query.order_by(Knowledge.updated_at.desc(), Knowledge.id.asc())
|
||||
|
||||
@@ -258,8 +267,8 @@ class KnowledgeTable:
|
||||
knowledge_bases.append(
|
||||
KnowledgeUserModel.model_validate(
|
||||
{
|
||||
**KnowledgeModel.model_validate(
|
||||
knowledge_base
|
||||
**self._to_knowledge_model(
|
||||
knowledge_base, db=db
|
||||
).model_dump(),
|
||||
"user": (
|
||||
UserModel.model_validate(user).model_dump()
|
||||
@@ -294,7 +303,14 @@ class KnowledgeTable:
|
||||
|
||||
# Apply access-control directly to the joined query
|
||||
# This makes the database handle filtering, even with 10k+ KBs
|
||||
query = has_permission(db, Knowledge, query, filter)
|
||||
query = AccessGrants.has_permission_filter(
|
||||
db=db,
|
||||
query=query,
|
||||
DocumentModel=Knowledge,
|
||||
filter=filter,
|
||||
resource_type="knowledge",
|
||||
permission="read",
|
||||
)
|
||||
|
||||
# Apply filename search
|
||||
if filter:
|
||||
@@ -327,8 +343,8 @@ class KnowledgeTable:
|
||||
if user
|
||||
else None
|
||||
),
|
||||
collection=KnowledgeModel.model_validate(
|
||||
knowledge
|
||||
collection=self._to_knowledge_model(
|
||||
knowledge, db=db
|
||||
).model_dump(),
|
||||
)
|
||||
)
|
||||
@@ -350,7 +366,14 @@ class KnowledgeTable:
|
||||
user_group_ids = {
|
||||
group.id for group in Groups.get_groups_by_member_id(user_id, db=db)
|
||||
}
|
||||
return has_access(user_id, permission, knowledge.access_control, user_group_ids)
|
||||
return AccessGrants.has_access(
|
||||
user_id=user_id,
|
||||
resource_type="knowledge",
|
||||
resource_id=knowledge.id,
|
||||
permission=permission,
|
||||
user_group_ids=user_group_ids,
|
||||
db=db,
|
||||
)
|
||||
|
||||
def get_knowledge_bases_by_user_id(
|
||||
self, user_id: str, permission: str = "write", db: Optional[Session] = None
|
||||
@@ -363,8 +386,13 @@ class KnowledgeTable:
|
||||
knowledge_base
|
||||
for knowledge_base in knowledge_bases
|
||||
if knowledge_base.user_id == user_id
|
||||
or has_access(
|
||||
user_id, permission, knowledge_base.access_control, user_group_ids
|
||||
or AccessGrants.has_access(
|
||||
user_id=user_id,
|
||||
resource_type="knowledge",
|
||||
resource_id=knowledge_base.id,
|
||||
permission=permission,
|
||||
user_group_ids=user_group_ids,
|
||||
db=db,
|
||||
)
|
||||
]
|
||||
|
||||
@@ -374,7 +402,9 @@ class KnowledgeTable:
|
||||
try:
|
||||
with get_db_context(db) as db:
|
||||
knowledge = db.query(Knowledge).filter_by(id=id).first()
|
||||
return KnowledgeModel.model_validate(knowledge) if knowledge else None
|
||||
return (
|
||||
self._to_knowledge_model(knowledge, db=db) if knowledge else None
|
||||
)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
@@ -391,7 +421,14 @@ class KnowledgeTable:
|
||||
user_group_ids = {
|
||||
group.id for group in Groups.get_groups_by_member_id(user_id, db=db)
|
||||
}
|
||||
if has_access(user_id, "write", knowledge.access_control, user_group_ids):
|
||||
if AccessGrants.has_access(
|
||||
user_id=user_id,
|
||||
resource_type="knowledge",
|
||||
resource_id=knowledge.id,
|
||||
permission="write",
|
||||
user_group_ids=user_group_ids,
|
||||
db=db,
|
||||
):
|
||||
return knowledge
|
||||
return None
|
||||
|
||||
@@ -406,9 +443,7 @@ class KnowledgeTable:
|
||||
.filter(KnowledgeFile.file_id == file_id)
|
||||
.all()
|
||||
)
|
||||
return [
|
||||
KnowledgeModel.model_validate(knowledge) for knowledge in knowledges
|
||||
]
|
||||
return [self._to_knowledge_model(knowledge, db=db) for knowledge in knowledges]
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
@@ -591,11 +626,15 @@ class KnowledgeTable:
|
||||
knowledge = self.get_knowledge_by_id(id=id, db=db)
|
||||
db.query(Knowledge).filter_by(id=id).update(
|
||||
{
|
||||
**form_data.model_dump(),
|
||||
**form_data.model_dump(exclude={"access_grants"}),
|
||||
"updated_at": int(time.time()),
|
||||
}
|
||||
)
|
||||
db.commit()
|
||||
if form_data.access_grants is not None:
|
||||
AccessGrants.set_access_grants(
|
||||
"knowledge", id, form_data.access_grants, db=db
|
||||
)
|
||||
return self.get_knowledge_by_id(id=id, db=db)
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
@@ -622,6 +661,7 @@ class KnowledgeTable:
|
||||
def delete_knowledge_by_id(self, id: str, db: Optional[Session] = None) -> bool:
|
||||
try:
|
||||
with get_db_context(db) as db:
|
||||
AccessGrants.revoke_all_access("knowledge", id, db=db)
|
||||
db.query(Knowledge).filter_by(id=id).delete()
|
||||
db.commit()
|
||||
return True
|
||||
@@ -631,6 +671,9 @@ class KnowledgeTable:
|
||||
def delete_all_knowledge(self, db: Optional[Session] = None) -> bool:
|
||||
with get_db_context(db) as db:
|
||||
try:
|
||||
knowledge_ids = [row[0] for row in db.query(Knowledge.id).all()]
|
||||
for knowledge_id in knowledge_ids:
|
||||
AccessGrants.revoke_all_access("knowledge", knowledge_id, db=db)
|
||||
db.query(Knowledge).delete()
|
||||
db.commit()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user