gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fixes to the sqlalchemy adapter

Browse source
This commit is contained in:
Vasilije 2024-08-05 11:59:36 +02:00
parent 9a2cde95d0
commit 6035010036
3 changed files with 12 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
cognee/modules/users/methods/get_default_user.py
View file

@ -1,11 +1,12 @@
from sqlalchemy.orm import joinedload
from cognee.modules.users.models import User from cognee.modules.users.models import User
from cognee.infrastructure.databases.relational import get_relational_engine from cognee.infrastructure.databases.relational import get_relational_engine
from sqlalchemy.future import select from sqlalchemy.future import select
async def get_default_user(session): async def get_default_user(session):
stmt = select(User).options(joinedload(User.groups)).where(User.email == "default_user@example.com")
stmt = select(User).where(User.email == "default_user@example.com")
result = await session.execute(stmt) result = await session.execute(stmt)
user = result.scalars().first() user = result.scalars().first()
return user return user

2
cognee/modules/users/models/User.py
View file

@ -24,7 +24,7 @@ class User(SQLAlchemyBaseUserTableUUID, Principal):
from fastapi_users import schemas from fastapi_users import schemas
class UserRead(schemas.BaseUser[uuid_UUID]): class UserRead(schemas.BaseUser[uuid_UUID]):
pass groups: list[uuid_UUID] # Add groups attribute
class UserCreate(schemas.BaseUserCreate): class UserCreate(schemas.BaseUserCreate):
pass pass

16
cognee/modules/users/permissions/methods/check_permissions_on_documents.py
View file

@ -15,18 +15,18 @@ class PermissionDeniedException(Exception):
async def check_permissions_on_documents(user: User, permission_type: str, document_ids: list[str], session): async def check_permissions_on_documents(user: User, permission_type: str, document_ids: list[str], session):
logging.info("This is the user: %s", user.__dict__)
try: try:
user_group_ids = [group.id for group in user.groups] user_group_ids = [group.id for group in user.groups]
result = await session.execute( acls = await session.execute(
select(ACL).filter( select(ACL)
ACL.principal_id.in_([user.id, *user_group_ids]), .join(ACL.permission)
ACL.permission.name == permission_type .where(ACL.principal_id.in_([user.id, *user_group_ids]))
) .where(ACL.permission.has(name=permission_type))
) )
acls = result.scalars().all() resource_ids = [resource.resource_id for acl in acls.scalars().all() for resource in acl.resources]
resource_ids = [resource.resource_id for acl in acls for resource in acl.resources]
has_permissions = all(document_id in resource_ids for document_id in document_ids) has_permissions = all(document_id in resource_ids for document_id in document_ids)
if not has_permissions: if not has_permissions: