refactor: Change variable names, add setting of current tenant to be optional for tenant creation
This commit is contained in:
parent
1ab8869af1
commit
e3b707a0c2
3 changed files with 14 additions and 9 deletions
|
|
@ -16,18 +16,18 @@ from cognee.modules.users.exceptions import (
|
||||||
|
|
||||||
|
|
||||||
async def add_user_to_tenant(
|
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.
|
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.
|
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:
|
Args:
|
||||||
user_id: Id of the user.
|
user_id: Id of the user.
|
||||||
tenant_id: Id of the tenant.
|
tenant_id: Id of the tenant.
|
||||||
owner_id: Id of the request owner.
|
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:
|
Returns:
|
||||||
None
|
None
|
||||||
|
|
@ -48,7 +48,7 @@ async def add_user_to_tenant(
|
||||||
message="Only tenant owner can add other users to organization."
|
message="Only tenant owner can add other users to organization."
|
||||||
)
|
)
|
||||||
|
|
||||||
if set_active_tenant:
|
if set_as_active_tenant:
|
||||||
user.tenant_id = tenant_id
|
user.tenant_id = tenant_id
|
||||||
await session.merge(user)
|
await session.merge(user)
|
||||||
await session.commit()
|
await session.commit()
|
||||||
|
|
|
||||||
|
|
@ -1,6 +1,7 @@
|
||||||
from uuid import UUID
|
from uuid import UUID
|
||||||
from sqlalchemy import insert
|
from sqlalchemy import insert
|
||||||
from sqlalchemy.exc import IntegrityError
|
from sqlalchemy.exc import IntegrityError
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
from cognee.modules.users.models.UserTenant import UserTenant
|
from cognee.modules.users.models.UserTenant import UserTenant
|
||||||
from cognee.infrastructure.databases.exceptions import EntityAlreadyExistsError
|
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
|
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.
|
Create a new tenant with the given name, for the user with the given id.
|
||||||
This user is the owner of the tenant.
|
This user is the owner of the tenant.
|
||||||
Args:
|
Args:
|
||||||
tenant_name: Name of the new tenant.
|
tenant_name: Name of the new tenant.
|
||||||
user_id: Id of the user.
|
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:
|
Returns:
|
||||||
None
|
None
|
||||||
|
|
@ -29,9 +33,10 @@ async def create_tenant(tenant_name: str, user_id: UUID) -> UUID:
|
||||||
session.add(tenant)
|
session.add(tenant)
|
||||||
await session.flush()
|
await session.flush()
|
||||||
|
|
||||||
user.tenant_id = tenant.id
|
if set_as_active_tenant:
|
||||||
await session.merge(user)
|
user.tenant_id = tenant.id
|
||||||
await session.commit()
|
await session.merge(user)
|
||||||
|
await session.commit()
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# Add association directly to the association table
|
# Add association directly to the association table
|
||||||
|
|
|
||||||
|
|
@ -52,7 +52,7 @@ async def select_tenant(user_id: UUID, tenant_id: Union[UUID, None]):
|
||||||
try:
|
try:
|
||||||
result = result.scalar_one()
|
result = result.scalar_one()
|
||||||
except sqlalchemy.exc.NoResultFound as e:
|
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 result:
|
||||||
# If user is part of tenant update current tenant of user
|
# If user is part of tenant update current tenant of user
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue