refactor: Update permissions example

This commit is contained in:
Igor Ilic 2025-11-04 20:24:16 +01:00
parent db2a32dd17
commit f002d3bf0e
3 changed files with 44 additions and 9 deletions

View file

@ -16,7 +16,7 @@ from cognee.modules.users.exceptions import (
async def add_user_to_tenant(
user_id: UUID, tenant_id: UUID, owner_id: UUID, set_as_active_tenant: Optional[bool] = True
user_id: UUID, tenant_id: UUID, owner_id: UUID, set_as_active_tenant: Optional[bool] = False
):
"""
Add a user with the given id to the tenant with the given id.

View file

@ -7,11 +7,12 @@ 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]):
async def select_tenant(user_id: UUID, tenant_id: Union[UUID, None]) -> User:
"""
Set the users active tenant to provided tenant.
@ -33,7 +34,7 @@ async def select_tenant(user_id: UUID, tenant_id: Union[UUID, None]):
user.tenant_id = None
await session.merge(user)
await session.commit()
return
return user
tenant = await get_tenant(tenant_id)
@ -59,3 +60,4 @@ async def select_tenant(user_id: UUID, tenant_id: Union[UUID, None]):
user.tenant_id = tenant_id
await session.merge(user)
await session.commit()
return user

View file

@ -3,6 +3,7 @@ import cognee
import pathlib
from cognee.modules.users.exceptions import PermissionDeniedError
from cognee.modules.users.tenants.methods import select_tenant
from cognee.shared.logging_utils import get_logger
from cognee.modules.search.types import SearchType
from cognee.modules.users.methods import create_user
@ -116,6 +117,7 @@ async def main():
print(
"\nOperation started as user_2 to give read permission to user_1 for the dataset owned by user_2"
)
await authorized_give_permission_on_datasets(
user_1.id,
[quantum_dataset_id],
@ -142,6 +144,9 @@ async def main():
print("User 2 is creating CogneeLab tenant/organization")
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)
print("\nUser 2 is creating Researcher role")
role_id = await create_role(role_name="Researcher", owner_id=user_2.id)
@ -150,27 +155,55 @@ async def main():
# To add a user to a role he must be part of the same tenant/organization
print("\nOperation started as user_2 to add user_3 to CogneeLab tenant/organization")
await add_user_to_tenant(
user_id=user_3.id, tenant_id=tenant_id, owner_id=user_2.id, set_as_active_tenant=True
)
await add_user_to_tenant(user_id=user_3.id, tenant_id=tenant_id, owner_id=user_2.id)
print(
"\nOperation started by user_2, as tenant owner, to add user_3 to Researcher role inside the tenant/organization"
)
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)
print(
"\nOperation started as user_2 to give read permission to Researcher role for the dataset owned by user_2"
"\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"
)
# Even though the dataset owner is user_2, the dataset doesn't belong to the tenant/organization CogneeLab.
# So we can't assign permissions to it when we're acting in the CogneeLab tenant.
try:
await authorized_give_permission_on_datasets(
role_id,
[quantum_dataset_id],
"read",
user_2.id,
)
except PermissionDeniedError:
print(
"User 2 could not give permission to the role as the QUANTUM dataset is not part of the CogneeLab tenant"
)
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"
)
# 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)
# 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)
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],
[quantum_dataset_id_cognee_lab_tenant],
"read",
user_2.id,
)
# Now user_3 can read from QUANTUM dataset as part of the Researcher role after proper permissions have been assigned by the QUANTUM dataset owner, user_2.
print("\nSearch result as user_3 on the dataset owned by user_2:")
print("\nSearch result as user_3 on the QUANTUM dataset owned by the CogneeLab organization:")
search_results = await cognee.search(
query_type=SearchType.GRAPH_COMPLETION,
query_text="What is in the document?",