65 lines
2.1 KiB
Python
65 lines
2.1 KiB
Python
from uuid import UUID
|
|
|
|
from sqlalchemy.future import select
|
|
from sqlalchemy import insert
|
|
from sqlalchemy.exc import IntegrityError
|
|
|
|
from cognee.infrastructure.databases.exceptions import EntityAlreadyExistsError
|
|
from cognee.infrastructure.databases.relational import get_relational_engine
|
|
from cognee.modules.users.exceptions import (
|
|
UserNotFoundError,
|
|
RoleNotFoundError,
|
|
TenantNotFoundError,
|
|
PermissionDeniedError,
|
|
)
|
|
from cognee.modules.users.models import (
|
|
User,
|
|
Role,
|
|
Tenant,
|
|
UserRole,
|
|
)
|
|
|
|
|
|
async def add_user_to_role(user_id: UUID, role_id: UUID, owner_id: UUID):
|
|
"""
|
|
Add a user with the given id to the role with the given id.
|
|
Args:
|
|
user_id: Id of the user.
|
|
role_id: Id of the role.
|
|
owner_id: Id of the request owner.
|
|
|
|
Returns:
|
|
None
|
|
|
|
"""
|
|
db_engine = get_relational_engine()
|
|
async with db_engine.get_async_session() as session:
|
|
user = (await session.execute(select(User).where(User.id == user_id))).scalars().first()
|
|
role = (await session.execute(select(Role).where(Role.id == role_id))).scalars().first()
|
|
tenant = (
|
|
(await session.execute(select(Tenant).where(Tenant.id == role.tenant_id)))
|
|
.scalars()
|
|
.first()
|
|
)
|
|
|
|
if not user:
|
|
raise UserNotFoundError
|
|
elif not role:
|
|
raise RoleNotFoundError
|
|
elif user.tenant_id != role.tenant_id:
|
|
raise TenantNotFoundError(
|
|
message="User tenant does not match role tenant. User cannot be added to role."
|
|
)
|
|
elif tenant.owner_id != owner_id:
|
|
raise PermissionDeniedError(
|
|
message="User submitting request does not have permission to add user to role."
|
|
)
|
|
|
|
try:
|
|
# Add association directly to the association table
|
|
create_user_role_statement = insert(UserRole).values(user_id=user_id, role_id=role_id)
|
|
await session.execute(create_user_role_statement)
|
|
except IntegrityError:
|
|
raise EntityAlreadyExistsError(message="User is already part of group.")
|
|
|
|
await session.commit()
|