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

refactor: Update permissions example to work with new changes

Browse source
This commit is contained in:
Igor Ilic 2025-11-04 20:54:00 +01:00
parent f002d3bf0e
commit 7782f246d3
3 changed files with 16 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
cognee/api/v1/permissions/routers/get_permissions_router.py
View file

@ -259,7 +259,7 @@ def get_permissions_router() -> APIRouter:
from cognee.modules.users.tenants.methods import select_tenant as select_tenant_method
await select_tenant_method(user_id=user.id, tenant_id=payload.tenant_id)
await select_tenant_method(user=user, tenant_id=payload.tenant_id)
return JSONResponse(
status_code=200,

9
cognee/modules/users/tenants/methods/select_tenant.py
View file

@ -6,19 +6,18 @@ from sqlalchemy import select
from cognee.infrastructure.databases.relational import get_relational_engine
from cognee.modules.users.models.UserTenant import UserTenant
from cognee.modules.users.methods import get_user
from cognee.modules.users.models.User import User
from cognee.modules.users.permissions.methods import get_tenant
from cognee.modules.users.exceptions import UserNotFoundError, TenantNotFoundError
async def select_tenant(user_id: UUID, tenant_id: Union[UUID, None]) -> User:
async def select_tenant(user: User, tenant_id: Union[UUID, None]) -> User:
"""
Set the users active tenant to provided tenant.
If None tenant_id is provided set current Tenant to the default single user-tenant
Args:
user_id: Id of the user.
user: User object.
tenant_id: Id of the tenant.
Returns:
@ -27,8 +26,6 @@ async def select_tenant(user_id: UUID, tenant_id: Union[UUID, None]) -> User:
"""
db_engine = get_relational_engine()
async with db_engine.get_async_session() as session:
user = await get_user(user_id)
if tenant_id is None:
# If no tenant_id is provided set current Tenant to the single user-tenant
user.tenant_id = None
@ -46,7 +43,7 @@ async def select_tenant(user_id: UUID, tenant_id: Union[UUID, None]) -> User:
# Check if User is part of Tenant
result = await session.execute(
select(UserTenant)
.where(UserTenant.user_id == user_id)
.where(UserTenant.user_id == user.id)
.where(UserTenant.tenant_id == tenant_id)
)

22
examples/python/permissions_example.py
View file

@ -145,7 +145,7 @@ async def main():
tenant_id = await create_tenant("CogneeLab", user_2.id)
print("User 2 is selecting CogneeLab tenant/organization as active tenant")
await select_tenant(user_id=user_2.id, tenant_id=tenant_id)
await select_tenant(user=user_2, tenant_id=tenant_id)
print("\nUser 2 is creating Researcher role")
role_id = await create_role(role_name="Researcher", owner_id=user_2.id)
@ -163,7 +163,7 @@ async def main():
await add_user_to_role(user_id=user_3.id, role_id=role_id, owner_id=user_2.id)
print("\nOperation as user_3 to select CogneeLab tenant/organization as active tenant")
await select_tenant(user_id=user_3.id, tenant_id=tenant_id)
await select_tenant(user=user_3, tenant_id=tenant_id)
print(
"\nOperation started as user_2, with CogneeLab as its active tenant, to give read permission to Researcher role for the dataset QUANTUM owned by user_2"
@ -183,21 +183,23 @@ async def main():
)
print(
"We will now create a new QUANTUM dataset in the CogneeLab tenant so that permissions can be assigned to the Researcher role inside the tenant/organization"
"We will now create a new QUANTUM dataset with the QUANTUM_COGNEE_LAB name in the CogneeLab tenant so that permissions can be assigned to the Researcher role inside the tenant/organization"
)
# Re-create the QUANTUM dataset in the CogneeLab tenant. The old QUANTUM dataset is still owned by user_2 personally
# We can re-create the QUANTUM dataset in the CogneeLab tenant. The old QUANTUM dataset is still owned by user_2 personally
# and can still be accessed by selecting the personal tenant for user 2.
await cognee.add([text], dataset_name="QUANTUM", user=user_2)
quantum_cognify_result = await cognee.cognify(["QUANTUM"], user=user_2)
await cognee.add([text], dataset_name="QUANTUM_COGNEE_LAB", user=user_2)
quantum_cognee_lab_cognify_result = await cognee.cognify(["QUANTUM_COGNEE_LAB"], user=user_2)
# The recreated Quantum dataset will now have a different dataset_id as it's a new dataset in a different organization
quantum_dataset_id_cognee_lab_tenant = extract_dataset_id_from_cognify(quantum_cognify_result)
quantum_cognee_lab_dataset_id = extract_dataset_id_from_cognify(
quantum_cognee_lab_cognify_result
)
print(
"\nOperation started as user_2, with CogneeLab as its active tenant, to give read permission to Researcher role for the dataset QUANTUM owned by the CogneeLab tenant"
)
await authorized_give_permission_on_datasets(
role_id,
[quantum_dataset_id_cognee_lab_tenant],
[quantum_cognee_lab_dataset_id],
"read",
user_2.id,
)
@ -207,8 +209,8 @@ async def main():
search_results = await cognee.search(
query_type=SearchType.GRAPH_COMPLETION,
query_text="What is in the document?",
user=user_1,
dataset_ids=[quantum_dataset_id],
user=user_3,
dataset_ids=[quantum_cognee_lab_dataset_id],
)
for result in search_results:
print(f"{result}\n")