refactor: implement team isolation for task assignments (#240)

* refactor: implement team isolation for task assignments

- Added `original_team_id` field to `TaskAssignmentModel` to track original team context
- Modified `TaskAssignmentRepository.update_assignment()` to set `original_team_id` when reassigning from team to user
- Updated `TaskRepository.list()` and `count()` to include team member tasks with proper team isolation
- Team task lists (`/v1/tasks?teamId=...`) now include tasks assigned to team members
- Tasks are filtered by `original_team_id` to prevent cross-team contamination
- Users in multiple teams only see tasks relevant to each specific team context
- Modified `TaskAssignmentDetailView.patch()` to use `update_assignment()` instead of `update_executor()`
- This ensures `original_team_id` is properly set when UI calls `PATCH /v1/task-assignments/{task_id}`
- Maintains backward compatibility with existing UI workflow
- Added support for both ObjectId and string formats in MongoDB queries
- Prevents data type mismatches in `original_team_id` filtering

* chore: remove unnecessary comments from code

* chore: remove redudant query and added a helper method and calling that helper method in list and count method

* fix: added the originial_team_id in the create task assignment method

* refactor: rename original_team_id to team_id in task assignment models and services

- Updated CreateTaskAssignmentDTO and TaskAssignmentDTO to replace original_team_id with team_id for clarity.
- Adjusted validation methods and repository logic to reflect the new field name.
- Ensured consistency across task assignment service and repository for handling team assignments.

* fix: ensure team_id is stored as ObjectId in task assignments

---------

Co-authored-by: anuj.k <[email protected]>

Author: Achintya-Chatterjee

todo/dto/task_assignment_dto.py

@@ -8,6 +8,7 @@ class CreateTaskAssignmentDTO(BaseModel):
     task_id: str
     assignee_id: str
     user_type: Literal["user", "team"]
+    team_id: Optional[str] = None
 
     @validator("task_id")
     def validate_task_id(cls, value):
@@ -30,6 +31,13 @@ def validate_user_type(cls, value):
             raise ValueError("user_type must be either 'user' or 'team'")
         return value
 
+    @validator("team_id")
+    def validate_team_id(cls, value):
+        """Validate that the original team ID is a valid ObjectId if provided."""
+        if value is not None and not ObjectId.is_valid(value):
+            raise ValueError(f"Invalid original team ID: {value}")
+        return value
+
 
 class TaskAssignmentDTO(BaseModel):
     id: str
@@ -38,6 +46,7 @@ class TaskAssignmentDTO(BaseModel):
     assignee_name: Optional[str] = None
     user_type: Literal["user", "team"]
     executor_id: Optional[str] = None  # User ID executing the task (for team assignments)
+    team_id: Optional[str] = None
     is_active: bool
     created_by: str
     updated_by: Optional[str] = None

todo/models/task_assignment.py

@@ -24,8 +24,9 @@ class TaskAssignmentModel(Document):
     created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
     updated_at: datetime | None = None
     executor_id: PyObjectId | None = None  # User within the team who is executing the task
+    team_id: PyObjectId | None = None  # Track the original team when reassigned from team to user
 
-    @validator("task_id", "assignee_id", "created_by", "updated_by")
+    @validator("task_id", "assignee_id", "created_by", "updated_by", "team_id")
     def validate_object_ids(cls, v):
         if v is None:
             return v

todo/repositories/task_assignment_repository.py

@@ -2,6 +2,7 @@
 from typing import Optional, List
 from bson import ObjectId
 
+from todo.exceptions.task_exceptions import TaskNotFoundException
 from todo.models.task_assignment import TaskAssignmentModel
 from todo.repositories.common.mongo_repository import MongoRepository
 from todo.models.common.pyobjectid import PyObjectId
@@ -72,6 +73,18 @@ def update_assignment(
         """
         collection = cls.get_collection()
         try:
+            current_assignment = cls.get_by_task_id(task_id)
+
+            if not current_assignment:
+                raise TaskNotFoundException(task_id)
+
+            team_id = None
+
+            if user_type == "team":
+                team_id = assignee_id
+            elif user_type == "user" and current_assignment.team_id is not None:
+                team_id = current_assignment.team_id
+
             # Deactivate current assignment if exists (try both ObjectId and string)
             collection.update_many(
                 {"task_id": ObjectId(task_id), "is_active": True},
@@ -102,6 +115,7 @@ def update_assignment(
                 user_type=user_type,
                 created_by=PyObjectId(user_id),
                 updated_by=None,
+                team_id=PyObjectId(team_id),
             )
 
             return cls.create(new_assignment)

todo/repositories/task_repository.py

@@ -2,19 +2,25 @@
 from typing import List
 from bson import ObjectId
 from pymongo import ReturnDocument
-import logging
 
 from todo.exceptions.task_exceptions import TaskNotFoundException
 from todo.models.task import TaskModel
 from todo.repositories.common.mongo_repository import MongoRepository
 from todo.repositories.task_assignment_repository import TaskAssignmentRepository
 from todo.constants.messages import ApiErrors, RepositoryErrors
 from todo.constants.task import SORT_FIELD_PRIORITY, SORT_FIELD_ASSIGNEE, SORT_ORDER_DESC, TaskStatus
+from todo.repositories.team_repository import UserTeamDetailsRepository
 
 
 class TaskRepository(MongoRepository):
     collection_name = TaskModel.collection_name
 
+    @classmethod
+    def _get_team_task_ids(cls, team_id: str) -> List[ObjectId]:
+        team_tasks = TaskAssignmentRepository.get_collection().find({"team_id": team_id, "is_active": True})
+        team_task_ids = [ObjectId(task["task_id"]) for task in team_tasks]
+        return list(set(team_task_ids))
+
     @classmethod
     def _build_status_filter(cls, status_filter: str = None) -> dict:
         now = datetime.now(timezone.utc)
@@ -60,17 +66,12 @@ def list(
         status_filter: str = None,
     ) -> List[TaskModel]:
         tasks_collection = cls.get_collection()
-        logger = logging.getLogger(__name__)
 
         base_filter = cls._build_status_filter(status_filter)
 
         if team_id:
-            logger.debug(f"TaskRepository.list: team_id={team_id}")
-            team_assignments = TaskAssignmentRepository.get_by_assignee_id(team_id, "team")
-            team_task_ids = [assignment.task_id for assignment in team_assignments]
-            logger.debug(f"TaskRepository.list: team_task_ids={team_task_ids}")
-            query_filter = {"$and": [base_filter, {"_id": {"$in": team_task_ids}}]}
-            logger.debug(f"TaskRepository.list: query_filter={query_filter}")
+            all_team_task_ids = cls._get_team_task_ids(team_id)
+            query_filter = {"$and": [base_filter, {"_id": {"$in": all_team_task_ids}}]}
         elif user_id:
             assigned_task_ids = cls._get_assigned_task_ids_for_user(user_id)
             query_filter = {"$and": [base_filter, {"_id": {"$in": assigned_task_ids}}]}
@@ -98,7 +99,7 @@ def _get_assigned_task_ids_for_user(cls, user_id: str) -> List[ObjectId]:
         direct_task_ids = [assignment.task_id for assignment in direct_assignments]
 
         # Get teams where user is a member
-        from todo.repositories.team_repository import UserTeamDetailsRepository, TeamRepository
+        from todo.repositories.team_repository import TeamRepository
 
         user_teams = UserTeamDetailsRepository.get_by_user_id(user_id)
         team_ids = [str(team.team_id) for team in user_teams]
@@ -128,9 +129,9 @@ def count(cls, user_id: str = None, team_id: str = None, status_filter: str = No
         base_filter = cls._build_status_filter(status_filter)
 
         if team_id:
-            team_assignments = TaskAssignmentRepository.get_by_assignee_id(team_id, "team")
-            team_task_ids = [assignment.task_id for assignment in team_assignments]
-            query_filter = {"$and": [base_filter, {"_id": {"$in": team_task_ids}}]}
+            all_team_task_ids = cls._get_team_task_ids(team_id)
+            query_filter = {"$and": [base_filter, {"_id": {"$in": all_team_task_ids}}]}
+
         elif user_id:
             assigned_task_ids = cls._get_assigned_task_ids_for_user(user_id)
             query_filter = {

todo/services/task_assignment_service.py

@@ -66,6 +66,7 @@ def create_task_assignment(cls, dto: CreateTaskAssignmentDTO, user_id: str) -> C
                 user_type=dto.user_type,
                 created_by=PyObjectId(user_id),
                 updated_by=None,
+                team_id=PyObjectId(dto.team_id) if dto.team_id else None,
             )
             assignment = TaskAssignmentRepository.create(task_assignment)
 

todo/services/task_service.py

@@ -235,6 +235,8 @@ def _prepare_assignee_dto(cls, assignee_details: TaskAssignmentModel) -> TaskAss
             assignee_id=assignee_id,
             assignee_name=assignee.name,
             user_type=assignee_details.user_type,
+            executor_id=str(assignee_details.executor_id) if assignee_details.executor_id else None,
+            team_id=str(assignee_details.team_id) if assignee_details.team_id else None,
             is_active=assignee_details.is_active,
             created_by=str(assignee_details.created_by),
             updated_by=str(assignee_details.updated_by) if assignee_details.updated_by else None,
@@ -628,11 +630,16 @@ def create_task(cls, dto: CreateTaskDTO) -> CreateTaskResponse:
             created_task = TaskRepository.create(task)
 
             # Create assignee relationship if assignee is provided
+            team_id = None
+            if dto.assignee and dto.assignee.get("user_type") == "team":
+                team_id = dto.assignee.get("assignee_id")
+
             if dto.assignee:
                 assignee_dto = CreateTaskAssignmentDTO(
                     task_id=str(created_task.id),
                     assignee_id=dto.assignee.get("assignee_id"),
                     user_type=dto.assignee.get("user_type"),
+                    team_id=team_id,
                 )
                 TaskAssignmentService.create_task_assignment(assignee_dto, created_task.createdBy)
 

todo/views/task_assignment.py

@@ -221,11 +221,12 @@ def patch(self, request: Request, task_id: str):
             return Response(
                 {"error": f"User {executor_id} is not a member of the team."}, status=status.HTTP_400_BAD_REQUEST
             )
-
         # Update executor_id
         try:
-            updated = TaskAssignmentRepository.update_executor(task_id, executor_id, user["user_id"])
-            if not updated:
+            updated_assignment = TaskAssignmentRepository.update_assignment(
+                task_id, executor_id, "user", user["user_id"]
+            )
+            if not updated_assignment:
                 # Get more details about why it failed
                 import traceback
 
@@ -234,7 +235,7 @@ def patch(self, request: Request, task_id: str):
                 )
                 print(f"DEBUG: assignment details: {assignment}")
                 return Response(
-                    {"error": "Failed to update executor. Check server logs for details."},
+                    {"error": "Failed to update assignment. Check server logs for details."},
                     status=status.HTTP_500_INTERNAL_SERVER_ERROR,
                 )
         except Exception as e: