Repository: topoteretes/cognee
File: cognee/modules/users/tenants/methods/remove_user_from_tenant.py (Python)

Code before:

    Only the tenant owner can remove users. The tenant owner cannot be removed
    from their own tenant. Removes the user from all roles in the tenant and
    revokes their permissions on datasets belonging to the tenant. Data owned
    by the removed user within the tenant (e.g. datasets they created) remains
    in the tenant; only their membership and direct permissions are removed.

    Args:
        user_id: Id of the user to remove.
        tenant_id: Id of the tenant.
        owner_id: Id of the requester (must be tenant owner).

    Raises:
        PermissionDeniedError: If requester is not the tenant owner.
        TenantNotFoundError: If the tenant does not exist.
        EntityNotFoundError: If the user does not exist or is not in the tenant.
        CogneeValidationError: If attempting to remove the tenant owner.
    """
    await has_user_management_permission(requester_id=owner_id, tenant_id=tenant_id)
    tenant = await get_tenant(tenant_id)

    if tenant.owner_id == user_id:
        raise CogneeValidationError(
            message="Cannot remove the tenant owner from their own tenant.",
            name="CogneeValidationError",
            status_code=status.HTTP_400_BAD_REQUEST,
        )

    await get_user(user_id)

    db_engine = get_relational_engine()
    async with db_engine.get_async_session() as session:
        # Check user is in tenant
        user_tenant_result = await session.execute(
            select(UserTenant).where(
                UserTenant.user_id == user_id,
                UserTenant.tenant_id == tenant_id,
            )
        )
        if user_tenant_result.scalars().first() is None:
            raise EntityNotFoundError(message="User not found in this tenant.")

        # Subquery for role ids in this tenant
        role_ids_in_tenant = select(Role.id).where(Role.tenant_id == tenant_id)
        # Remove user from all roles in this tenant
        await session.execute(
            delete(UserRole).where(
                UserRole.user_id == user_id,
                UserRole.role_id.in_(role_ids_in_tenant),
            )
        )

Review comment (bug):
@dexters1 
Hi! I saw your comment about the edge case where an admin could remove the tenant owner. I understand the concern, but I’m not able to find that specific comment on the PR anymore.

Just to confirm, I’ve already added an explicit guard to prevent this scenario. Even if an admin has user-management permissions, they still cannot remove the tenant owner.

This ensures the tenant owner is always protected, regardless of role permissions.
Please let me know if you’d like this handled in a different way.

Code after:

    The tenant owner or any user with user management permission in the tenant
    (e.g. users in the Admin role) can remove users. The tenant owner cannot
    be removed from their own tenant. Removes the user from all roles in the
    tenant and revokes their permissions on datasets belonging to the tenant.
    Data owned by the removed user within the tenant (e.g. datasets they
    created) remains in the tenant; only their membership and direct
    permissions are removed.

    Args:
        user_id: Id of the user to remove.
        tenant_id: Id of the tenant.
        owner_id: Id of the requester (must be tenant owner or have user
            management permission in the tenant).

    Raises:
        PermissionDeniedError: If requester is not the tenant owner and does
            not have user management permission in the tenant.
        TenantNotFoundError: If the tenant does not exist.
        EntityNotFoundError: If the user does not exist or is not in the tenant.
        CogneeValidationError: If attempting to remove the tenant owner.
    """
    await has_user_management_permission(requester_id=owner_id, tenant_id=tenant_id)
    tenant = await get_tenant(tenant_id)

    if tenant.owner_id == user_id:
        raise CogneeValidationError(
            message="Cannot remove the tenant owner from their own tenant.",
            name="CogneeValidationError",
            status_code=status.HTTP_400_BAD_REQUEST,
        )

    await get_user(user_id)

    db_engine = get_relational_engine()
    async with db_engine.get_async_session() as session:
        # Check user is in tenant
        user_tenant_result = await session.execute(
            select(UserTenant).where(
                UserTenant.user_id == user_id,
                UserTenant.tenant_id == tenant_id,
            )
        )
        if user_tenant_result.scalars().first() is None:
            raise EntityNotFoundError(message="User not found in this tenant.")

        # Subquery for role ids in this tenant
        role_ids_in_tenant = select(Role.id).where(Role.tenant_id == tenant_id)
        # Remove user from all roles in this tenant
        await session.execute(
            delete(UserRole).where(