refactor: Update permissions example
This commit is contained in:
parent
db2a32dd17
commit
f002d3bf0e
3 changed files with 44 additions and 9 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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?",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue