refactor: Update permissions example

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

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

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

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

examples/python/permissions_example.py

@@ -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(
+    await add_user_to_tenant(user_id=user_3.id, tenant_id=tenant_id, owner_id=user_2.id)
-        user_id=user_3.id, tenant_id=tenant_id, owner_id=user_2.id, set_as_active_tenant=True
-    )
 
     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?",