refactor: Update permissions example to work with new changes

cognee/api/v1/permissions/routers/get_permissions_router.py

@@ -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,

cognee/modules/users/tenants/methods/select_tenant.py

@@ -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)
         )
 

examples/python/permissions_example.py

@@ -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)
+    await cognee.add([text], dataset_name="QUANTUM_COGNEE_LAB", user=user_2)
-    quantum_cognify_result = await cognee.cognify(["QUANTUM"], 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,
+        user=user_3,
-        dataset_ids=[quantum_dataset_id],
+        dataset_ids=[quantum_cognee_lab_dataset_id],
     )
     for result in search_results:
         print(f"{result}\n")