refactor: Enable selection of default single user tenant

cognee/api/v1/permissions/routers/get_permissions_router.py

@@ -1,14 +1,19 @@
 from uuid import UUID
-from typing import List
+from typing import List, Union
 
 from fastapi import APIRouter, Depends
 from fastapi.responses import JSONResponse
 
 from cognee.modules.users.models import User
+from cognee.api.DTO import InDTO
 from cognee.modules.users.methods import get_authenticated_user
 from cognee.shared.utils import send_telemetry
 
 
+class SelectTenantDTO(InDTO):
+    tenant_id: UUID | None = None
+
+
 def get_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)}
         )
 
-    @permissions_router.post("/tenants/{tenant_id}")
-    async def select_tenant(tenant_id: UUID, user: User = Depends(get_authenticated_user)):
+    @permissions_router.post("/tenants/select")
+    async def select_tenant(payload: SelectTenantDTO, user: User = Depends(get_authenticated_user)):
         """
         Select current tenant.
 
@@ -229,8 +234,10 @@ def get_permissions_router() -> APIRouter:
         to organize users and resources in multi-tenant environments, providing
         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
-        - **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
         Returns a success message indicating the tenant was created.
@@ -239,17 +246,18 @@ def get_permissions_router() -> APIRouter:
             "Permissions API Endpoint Invoked",
             user.id,
             additional_properties={
-                "endpoint": f"POST /v1/permissions/tenants/{str(tenant_id)}",
-                "tenant_id": tenant_id,
+                "endpoint": f"POST /v1/permissions/tenants/{str(payload.tenant_id)}",
+                "tenant_id": str(payload.tenant_id),
             },
         )
 
         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(
-            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

cognee/modules/users/tenants/methods/select_tenant.py

@@ -1,4 +1,5 @@
 from uuid import UUID
+from typing import Union
 
 import sqlalchemy.exc
 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
 
 
-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.
+
+        If None tenant_id is provided set current Tenant to the default single user-tenant
     Args:
         user_id: Id of the user.
         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()
     async with db_engine.get_async_session() as session:
         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)
 
         if not user: