feat: Add tenant select method/endpoint for users

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

@@ -220,4 +220,36 @@ def get_permissions_router() -> APIRouter:
             status_code=200, content={"message": "Tenant created.", "tenant_id": str(tenant_id)}
         )
 
+    @permissions_router.post("/tenants/{tenant_id}")
+    async def select_tenant(tenant_id: UUID, user: User = Depends(get_authenticated_user)):
+        """
+        Select current tenant.
+
+        This endpoint selects a tenant with the specified UUID. Tenants are used
+        to organize users and resources in multi-tenant environments, providing
+        isolation and access control between different groups or organizations.
+
+        ## Request Parameters
+        - **tenant_id** (UUID): UUID of the tenant to create
+
+        ## Response
+        Returns a success message indicating the tenant was created.
+        """
+        send_telemetry(
+            "Permissions API Endpoint Invoked",
+            user.id,
+            additional_properties={
+                "endpoint": f"POST /v1/permissions/tenants/{str(tenant_id)}",
+                "tenant_id": tenant_id,
+            },
+        )
+
+        from cognee.modules.users.tenants.methods import select_tenant as select_tenant_method
+
+        await select_tenant_method(user_id=user.id, tenant_id=tenant_id)
+
+        return JSONResponse(
+            status_code=200, content={"message": "Tenant selected.", "tenant_id": str(tenant_id)}
+        )
+
     return permissions_router

cognee/modules/users/permissions/methods/get_all_user_permission_datasets.py

@@ -41,5 +41,4 @@ async def get_all_user_permission_datasets(user: User, permission_type: str) ->
         unique.setdefault(dataset.id, dataset)
 
     # TODO: Add filtering out of datasets that aren't currently selected tenant of user (currently selected tenant is the tenant_id value in the User model)
-    # TODO: Add endpoint/method to select current Tenant for a user (This UUID value should be stored in tenant_id of User model)
     return list(unique.values())

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

@@ -1,2 +1,3 @@
 from .create_tenant import create_tenant
 from .add_user_to_tenant import add_user_to_tenant
+from .select_tenant import select_tenant

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

@@ -30,6 +30,8 @@ async def create_tenant(tenant_name: str, user_id: UUID) -> UUID:
             await session.flush()
 
             user.tenant_id = tenant.id
+            await session.merge(user)
+            await session.commit()
 
             try:
                 # Add association directly to the association table
@@ -37,11 +39,10 @@ async def create_tenant(tenant_name: str, user_id: UUID) -> UUID:
                     user_id=user_id, tenant_id=tenant.id
                 )
                 await session.execute(create_user_tenant_statement)
+                await session.commit()
             except IntegrityError:
-                raise EntityAlreadyExistsError(message="User is already part of group.")
+                raise EntityAlreadyExistsError(message="User is already part of tenant.")
 
-            await session.merge(user)
-            await session.commit()
             return tenant.id
         except IntegrityError as e:
             raise EntityAlreadyExistsError(message="Tenant already exists.") from e

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

@@ -0,0 +1,50 @@
+from uuid import UUID
+
+import sqlalchemy.exc
+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.permissions.methods import get_tenant
+from cognee.modules.users.exceptions import UserNotFoundError, TenantNotFoundError
+
+
+async def select_tenant(user_id: UUID, tenant_id: UUID):
+    """
+        Set the users active tenant to provided tenant.
+    Args:
+        user_id: Id of the user.
+        tenant_id: Id of the tenant.
+
+    Returns:
+        None
+
+    """
+    db_engine = get_relational_engine()
+    async with db_engine.get_async_session() as session:
+        user = await get_user(user_id)
+        tenant = await get_tenant(tenant_id)
+
+        if not user:
+            raise UserNotFoundError
+        elif not tenant:
+            raise TenantNotFoundError
+
+        # Check if User is part of Tenant
+        result = await session.execute(
+            select(UserTenant)
+            .where(UserTenant.user_id == user_id)
+            .where(UserTenant.tenant_id == tenant_id)
+        )
+
+        try:
+            result = result.scalar_one()
+        except sqlalchemy.exc.NoResultFound as e:
+            raise TenantNotFoundError("User Tenant relationship not found.") from e
+
+        if result:
+            # If user is part of tenant update current tenant of user
+            user.tenant_id = tenant_id
+            await session.merge(user)
+            await session.commit()