From e3b707a0c242fb7268d56d2b485990f120fe0462 Mon Sep 17 00:00:00 2001
From: Igor Ilic <igorilic03@gmail.com>
Date: Tue, 4 Nov 2025 12:20:17 +0100
Subject: [PATCH] refactor: Change variable names, add setting of current
 tenant to be optional for tenant creation

---
 .../users/tenants/methods/add_user_to_tenant.py     |  8 ++++----
 .../modules/users/tenants/methods/create_tenant.py  | 13 +++++++++----
 .../modules/users/tenants/methods/select_tenant.py  |  2 +-
 3 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/cognee/modules/users/tenants/methods/add_user_to_tenant.py b/cognee/modules/users/tenants/methods/add_user_to_tenant.py
index dabab6b6b..edadfe66b 100644
--- a/cognee/modules/users/tenants/methods/add_user_to_tenant.py
+++ b/cognee/modules/users/tenants/methods/add_user_to_tenant.py
@@ -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()
diff --git a/cognee/modules/users/tenants/methods/create_tenant.py b/cognee/modules/users/tenants/methods/create_tenant.py
index 60e10db5c..32baa05fd 100644
--- a/cognee/modules/users/tenants/methods/create_tenant.py
+++ b/cognee/modules/users/tenants/methods/create_tenant.py
@@ -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
diff --git a/cognee/modules/users/tenants/methods/select_tenant.py b/cognee/modules/users/tenants/methods/select_tenant.py
index 732b24858..6e72fea2f 100644
--- a/cognee/modules/users/tenants/methods/select_tenant.py
+++ b/cognee/modules/users/tenants/methods/select_tenant.py
@@ -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