refactor: Enable selection of default single user tenant
This commit is contained in:
parent
4f874deace
commit
6934692e1b
2 changed files with 28 additions and 9 deletions
|
|
@ -1,14 +1,19 @@
|
||||||
from uuid import UUID
|
from uuid import UUID
|
||||||
from typing import List
|
from typing import List, Union
|
||||||
|
|
||||||
from fastapi import APIRouter, Depends
|
from fastapi import APIRouter, Depends
|
||||||
from fastapi.responses import JSONResponse
|
from fastapi.responses import JSONResponse
|
||||||
|
|
||||||
from cognee.modules.users.models import User
|
from cognee.modules.users.models import User
|
||||||
|
from cognee.api.DTO import InDTO
|
||||||
from cognee.modules.users.methods import get_authenticated_user
|
from cognee.modules.users.methods import get_authenticated_user
|
||||||
from cognee.shared.utils import send_telemetry
|
from cognee.shared.utils import send_telemetry
|
||||||
|
|
||||||
|
|
||||||
|
class SelectTenantDTO(InDTO):
|
||||||
|
tenant_id: UUID | None = None
|
||||||
|
|
||||||
|
|
||||||
def get_permissions_router() -> APIRouter:
|
def get_permissions_router() -> APIRouter:
|
||||||
permissions_router = APIRouter()
|
permissions_router = APIRouter()
|
||||||
|
|
||||||
|
|
@ -220,8 +225,8 @@ def get_permissions_router() -> APIRouter:
|
||||||
status_code=200, content={"message": "Tenant created.", "tenant_id": str(tenant_id)}
|
status_code=200, content={"message": "Tenant created.", "tenant_id": str(tenant_id)}
|
||||||
)
|
)
|
||||||
|
|
||||||
@permissions_router.post("/tenants/{tenant_id}")
|
@permissions_router.post("/tenants/select")
|
||||||
async def select_tenant(tenant_id: UUID, user: User = Depends(get_authenticated_user)):
|
async def select_tenant(payload: SelectTenantDTO, user: User = Depends(get_authenticated_user)):
|
||||||
"""
|
"""
|
||||||
Select current tenant.
|
Select current tenant.
|
||||||
|
|
||||||
|
|
@ -229,8 +234,10 @@ def get_permissions_router() -> APIRouter:
|
||||||
to organize users and resources in multi-tenant environments, providing
|
to organize users and resources in multi-tenant environments, providing
|
||||||
isolation and access control between different groups or organizations.
|
isolation and access control between different groups or organizations.
|
||||||
|
|
||||||
|
Sending a null/None value as tenant_id selects his default single user tenant
|
||||||
|
|
||||||
## Request Parameters
|
## Request Parameters
|
||||||
- **tenant_id** (UUID): UUID of the tenant to create
|
- **tenant_id** (Union[UUID, None]): UUID of the tenant to select, If null/None is provided use the default single user tenant
|
||||||
|
|
||||||
## Response
|
## Response
|
||||||
Returns a success message indicating the tenant was created.
|
Returns a success message indicating the tenant was created.
|
||||||
|
|
@ -239,17 +246,18 @@ def get_permissions_router() -> APIRouter:
|
||||||
"Permissions API Endpoint Invoked",
|
"Permissions API Endpoint Invoked",
|
||||||
user.id,
|
user.id,
|
||||||
additional_properties={
|
additional_properties={
|
||||||
"endpoint": f"POST /v1/permissions/tenants/{str(tenant_id)}",
|
"endpoint": f"POST /v1/permissions/tenants/{str(payload.tenant_id)}",
|
||||||
"tenant_id": tenant_id,
|
"tenant_id": str(payload.tenant_id),
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
from cognee.modules.users.tenants.methods import select_tenant as select_tenant_method
|
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)
|
await select_tenant_method(user_id=user.id, tenant_id=payload.tenant_id)
|
||||||
|
|
||||||
return JSONResponse(
|
return JSONResponse(
|
||||||
status_code=200, content={"message": "Tenant selected.", "tenant_id": str(tenant_id)}
|
status_code=200,
|
||||||
|
content={"message": "Tenant selected.", "tenant_id": str(payload.tenant_id)},
|
||||||
)
|
)
|
||||||
|
|
||||||
return permissions_router
|
return permissions_router
|
||||||
|
|
|
||||||
|
|
@ -1,4 +1,5 @@
|
||||||
from uuid import UUID
|
from uuid import UUID
|
||||||
|
from typing import Union
|
||||||
|
|
||||||
import sqlalchemy.exc
|
import sqlalchemy.exc
|
||||||
from sqlalchemy import select
|
from sqlalchemy import select
|
||||||
|
|
@ -10,9 +11,11 @@ from cognee.modules.users.permissions.methods import get_tenant
|
||||||
from cognee.modules.users.exceptions import UserNotFoundError, TenantNotFoundError
|
from cognee.modules.users.exceptions import UserNotFoundError, TenantNotFoundError
|
||||||
|
|
||||||
|
|
||||||
async def select_tenant(user_id: UUID, tenant_id: UUID):
|
async def select_tenant(user_id: UUID, tenant_id: Union[UUID, None]):
|
||||||
"""
|
"""
|
||||||
Set the users active tenant to provided tenant.
|
Set the users active tenant to provided tenant.
|
||||||
|
|
||||||
|
If None tenant_id is provided set current Tenant to the default single user-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.
|
||||||
|
|
@ -24,6 +27,14 @@ async def select_tenant(user_id: UUID, tenant_id: UUID):
|
||||||
db_engine = get_relational_engine()
|
db_engine = get_relational_engine()
|
||||||
async with db_engine.get_async_session() as session:
|
async with db_engine.get_async_session() as session:
|
||||||
user = await get_user(user_id)
|
user = await get_user(user_id)
|
||||||
|
|
||||||
|
if tenant_id is None:
|
||||||
|
# If no tenant_id is provided set current Tenant to the single user-tenant
|
||||||
|
user.tenant_id = None
|
||||||
|
await session.merge(user)
|
||||||
|
await session.commit()
|
||||||
|
return
|
||||||
|
|
||||||
tenant = await get_tenant(tenant_id)
|
tenant = await get_tenant(tenant_id)
|
||||||
|
|
||||||
if not user:
|
if not user:
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue