gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor: Change variable names, add setting of current tenant to be optional for tenant creation

Browse source
This commit is contained in:
Igor Ilic 2025-11-04 12:20:17 +01:00
parent 1ab8869af1
commit e3b707a0c2
3 changed files with 14 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
cognee/modules/users/tenants/methods/add_user_to_tenant.py
View file

@ -16,18 +16,18 @@ from cognee.modules.users.exceptions import (
async def add_user_to_tenant(
user_id: UUID, tenant_id: UUID, owner_id: UUID, set_active_tenant: Optional[bool] = True
user_id: UUID, tenant_id: UUID, owner_id: UUID, set_as_active_tenant: Optional[bool] = True
):
"""
Add a user with the given id to the tenant with the given id.
This can only be successful if the request owner with the given id is the tenant owner.
If set_active_tenant is true it will automatically set the users active tenant to provided tenant.
If set_as_active_tenant is true it will automatically set the users active tenant to provided tenant.
Args:
user_id: Id of the user.
tenant_id: Id of the tenant.
owner_id: Id of the request owner.
set_active_tenant: If set_active_tenant is true it will automatically set the users active tenant to provided tenant.
set_as_active_tenant: If set_as_active_tenant is true it will automatically set the users active tenant to provided tenant.
Returns:
None
@ -48,7 +48,7 @@ async def add_user_to_tenant(
message="Only tenant owner can add other users to organization."
)
if set_active_tenant:
if set_as_active_tenant:
user.tenant_id = tenant_id
await session.merge(user)
await session.commit()

13
cognee/modules/users/tenants/methods/create_tenant.py
View file

@ -1,6 +1,7 @@
from uuid import UUID
from sqlalchemy import insert
from sqlalchemy.exc import IntegrityError
from typing import Optional
from cognee.modules.users.models.UserTenant import UserTenant
from cognee.infrastructure.databases.exceptions import EntityAlreadyExistsError
@ -9,13 +10,16 @@ from cognee.modules.users.models import Tenant
from cognee.modules.users.methods import get_user
async def create_tenant(tenant_name: str, user_id: UUID) -> UUID:
async def create_tenant(
tenant_name: str, user_id: UUID, set_as_active_tenant: Optional[bool] = True
) -> UUID:
"""
Create a new tenant with the given name, for the user with the given id.
This user is the owner of the tenant.
Args:
tenant_name: Name of the new tenant.
user_id: Id of the user.
set_as_active_tenant: If true, set the newly created tenant as the active tenant for the user.
Returns:
None
@ -29,9 +33,10 @@ async def create_tenant(tenant_name: str, user_id: UUID) -> UUID:
session.add(tenant)
await session.flush()
user.tenant_id = tenant.id
await session.merge(user)
await session.commit()
if set_as_active_tenant:
user.tenant_id = tenant.id
await session.merge(user)
await session.commit()
try:
# Add association directly to the association table

2
cognee/modules/users/tenants/methods/select_tenant.py
View file

@ -52,7 +52,7 @@ async def select_tenant(user_id: UUID, tenant_id: Union[UUID, None]):
try:
result = result.scalar_one()
except sqlalchemy.exc.NoResultFound as e:
raise TenantNotFoundError("User Tenant relationship not found.") from e
raise TenantNotFoundError("User is not part of the tenant.") from e
if result:
# If user is part of tenant update current tenant of user