feat: restrict member removal to admins and reassign tasks (#255)

Hariom01010 · web-flow · commit e58e9eca0efa · 2025-09-19T23:54:08.000+05:30
* feat(team): restrict member removal to admins and reassign tasks

- enforce auth so only admins can remove another member
- automatically remove user roles on removal
- reassign all non-DONE tasks of the removed user to the team
- add audit log entry for task reassignment

* feat(remove-member): resolve comments by bot

* feat(remove-member): instantiate the exception raised

- fix cannot remove error being shown for correct exception

* feat(remove-from-team): sync changes for task assignment to postgres

- update postgres model for audit log to include task_count

* feat(remove-member): update cannot remove POC exception message
- fix formatting

* feat(remove-member-from-team): handle mixed ObjectId and string types in task reassignment

- Moves service imports to the top level in `team_service.py` for better code organization.
- Modifies the `update_many` call in `task_assignment_repository.py` to use the `$in` operator. This ensures that `assignee_id` and `team_id` can be matched regardless of whether they are stored as ObjectIds or strings, preventing potential issues when deactivating task assignments during member removal.

* feat(remove-member-from-team): move UserTeamDetailsRepository to the
function

* feat(remove-from-team): set updated_by to None in the ops for task
assignment creation in reassign_tasks_from_user_to_team

- fix formatting issues

* feat(remove-member): move task_reassignment function in transaction

- add team exception in views
- add assigned_from and assigned_to to audit log

* feat(remove-member): fix formatting issues

* feat(remove-from-team): include is_active in unique constraint to prevent sync failures

* feat(remove-from-team): remove unecessary comments

- fix bug where done tasks were also being reassigned
- fix bug where in postgres_user_role where data sync failed if a user
  was removed and added more than once
- updated reassigned tasks status and add sync to postgres db

* feat(remove-from-member): add serializer for path params validation

* feat(remove-from-team): removed aggregation pipeline and use normal db
queries

* feat(remove-from-todo): remove unecessary migration files

* feat(remove-from-team): remove task_count field from audit log model

* feat(remove-from-team): fix bug where reassigned todos not showing on
poc dashboard

* feat(remove-from-team): move remove role and validation to individual
function

* feat(remove-from-team): fix formatting issues

* feat(remove-from-team): use bool return value instead of count for task reassignment function
diff --git a/todo/constants/messages.py b/todo/constants/messages.py
@@ -54,6 +54,8 @@ class ApiErrors:
     USER_NOT_FOUND_GENERIC = "User not found."
     SEARCH_QUERY_EMPTY = "Search query cannot be empty"
     TASK_ALREADY_IN_WATCHLIST = "Task is already in the watchlist"
+    CANNOT_REMOVE_OWNER = "Owner cannot be removed from the team"
+    CANNOT_REMOVE_POC = "POC cannot be removed from a team. Reassign the POC first."
 
 
 # Validation error messages
diff --git a/todo/exceptions/team_exceptions.py b/todo/exceptions/team_exceptions.py
@@ -0,0 +1,22 @@
+from todo.constants.messages import ApiErrors
+
+
+class BaseTeamException(Exception):
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(self.message)
+
+
+class CannotRemoveOwnerException(BaseTeamException):
+    def __init__(self, message=ApiErrors.CANNOT_REMOVE_OWNER):
+        super().__init__(message)
+
+
+class NotTeamAdminException(BaseTeamException):
+    def __init__(self, message=ApiErrors.UNAUTHORIZED_TITLE):
+        super().__init__(message)
+
+
+class CannotRemoveTeamPOCException(BaseTeamException):
+    def __init__(self, message=ApiErrors.CANNOT_REMOVE_POC):
+        super().__init__(message)
diff --git a/todo/migrations/0003_alter_postgresuserrole_unique_together_and_more.py b/todo/migrations/0003_alter_postgresuserrole_unique_together_and_more.py
@@ -0,0 +1,24 @@
+# Generated by Django 5.1.5 on 2025-09-09 14:49
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("todo", "0002_rename_postgres_ta_assignee_95ca3b_idx_postgres_ta_assigne_f1c6e7_idx_and_more"),
+    ]
+
+    operations = [
+        migrations.AlterUniqueTogether(
+            name="postgresuserrole",
+            unique_together=set(),
+        ),
+        migrations.AddConstraint(
+            model_name="postgresuserrole",
+            constraint=models.UniqueConstraint(
+                condition=models.Q(("is_active", True)),
+                fields=("user_id", "role_name", "scope", "team_id"),
+                name="unique_active_user_team_role",
+            ),
+        ),
+    ]
diff --git a/todo/models/postgres/user_role.py b/todo/models/postgres/user_role.py
@@ -27,7 +27,13 @@ class PostgresUserRole(models.Model):
 
     class Meta:
         db_table = "postgres_user_roles"
-        unique_together = ["user_id", "role_name", "scope", "team_id"]
+        constraints = [
+            models.UniqueConstraint(
+                fields=["user_id", "role_name", "scope", "team_id"],
+                condition=models.Q(is_active=True),
+                name="unique_active_user_team_role",
+            )
+        ]
         indexes = [
             models.Index(fields=["mongo_id"]),
             models.Index(fields=["user_id"]),
diff --git a/todo/repositories/task_assignment_repository.py b/todo/repositories/task_assignment_repository.py
@@ -6,7 +6,9 @@
 from todo.models.task_assignment import TaskAssignmentModel
 from todo.repositories.common.mongo_repository import MongoRepository
 from todo.models.common.pyobjectid import PyObjectId
+from todo.constants.task import TaskStatus
 from todo.services.enhanced_dual_write_service import EnhancedDualWriteService
+from todo.repositories.audit_log_repository import AuditLogRepository, AuditLogModel
 
 
 class TaskAssignmentRepository(MongoRepository):
@@ -355,3 +357,167 @@ def deactivate_by_task_id(cls, task_id: str, user_id: str) -> bool:
             return result.modified_count > 0
         except Exception:
             return False
+
+    @classmethod
+    def reassign_tasks_from_user_to_team(cls, user_id: str, team_id: str, performed_by_user_id: str) -> bool:
+        """
+        Reassign all tasks of user to team
+        """
+        collection = cls.get_collection()
+        client = cls.get_client()
+        with client.start_session() as session:
+            try:
+                with session.start_transaction():
+                    now = datetime.now(timezone.utc)
+                    user_task_assignments = list(
+                        collection.find(
+                            {
+                                "$and": [
+                                    {"is_active": True},
+                                    {
+                                        "$or": [{"assignee_id": user_id}, {"assignee_id": ObjectId(user_id)}],
+                                    },
+                                    {"$or": [{"team_id": team_id}, {"team_id": ObjectId(team_id)}]},
+                                ]
+                            },
+                            session=session,
+                        )
+                    )
+                    if not user_task_assignments:
+                        return 0
+                    active_user_task_assignments_ids = [
+                        ObjectId(assignment["task_id"]) for assignment in user_task_assignments
+                    ]
+
+                    from todo.repositories.task_repository import TaskRepository
+
+                    tasks_collection = TaskRepository.get_collection()
+                    active_tasks = list(
+                        tasks_collection.find(
+                            {
+                                "_id": {"$in": active_user_task_assignments_ids},
+                                "status": {"$ne": TaskStatus.DONE.value},
+                            },
+                            session=session,
+                        )
+                    )
+                    not_done_tasks_ids = [str(tasks["_id"]) for tasks in active_tasks]
+                    tasks_to_reset_status_ids = []
+                    tasks_to_clear_deferred_ids = []
+                    for tasks in active_tasks:
+                        if tasks["status"] == TaskStatus.IN_PROGRESS.value:
+                            tasks_to_reset_status_ids.append(tasks["_id"])
+                        elif tasks.get("deferredDetails") is not None:
+                            tasks_to_clear_deferred_ids.append(tasks["_id"])
+
+                    collection.update_many(
+                        {
+                            "task_id": {"$in": not_done_tasks_ids},
+                        },
+                        {
+                            "$set": {
+                                "assignee_id": team_id,
+                                "user_type": "team",
+                                "updated_at": now,
+                                "updated_by": ObjectId(performed_by_user_id),
+                            }
+                        },
+                        session=session,
+                    )
+
+                    for assignment in user_task_assignments:
+                        AuditLogRepository.create(
+                            AuditLogModel(
+                                task_id=PyObjectId(assignment["task_id"]),
+                                team_id=PyObjectId(team_id),
+                                action="assigned_to_team",
+                                performed_by=PyObjectId(performed_by_user_id),
+                            )
+                        )
+
+                    tasks_collection.update_many(
+                        {"_id": {"$in": tasks_to_reset_status_ids}},
+                        {
+                            "$set": {
+                                "status": TaskStatus.TODO.value,
+                                "updated_at": now,
+                                "updated_by": ObjectId(performed_by_user_id),
+                            }
+                        },
+                        session=session,
+                    )
+                    tasks_collection.update_many(
+                        {"_id": {"$in": tasks_to_clear_deferred_ids}},
+                        {
+                            "$set": {
+                                "status": TaskStatus.TODO.value,
+                                "deferredDetails": None,
+                                "updated_at": now,
+                                "updated_by": ObjectId(performed_by_user_id),
+                            }
+                        },
+                        session=session,
+                    )
+
+                    tasks_by_id = {task["_id"]: task for task in active_tasks}
+                    operations = []
+                    dual_write_service = EnhancedDualWriteService()
+                    for assignment in user_task_assignments:
+                        operations.append(
+                            {
+                                "collection_name": "task_assignments",
+                                "operation": "update",
+                                "mongo_id": assignment["_id"],
+                                "data": {
+                                    "task_mongo_id": str(assignment["task_id"]),
+                                    "assignee_id": str(assignment["team_id"]),
+                                    "user_type": "team",
+                                    "team_id": str(assignment["team_id"]),
+                                    "is_active": True,
+                                    "created_at": assignment["created_at"],
+                                    "created_by": str(assignment["created_by"]),
+                                    "updated_at": datetime.now(timezone.utc),
+                                    "updated_by": str(performed_by_user_id),
+                                },
+                            }
+                        )
+                        if (
+                            assignment["task_id"] in tasks_to_clear_deferred_ids
+                            or assignment["task_id"] in tasks_to_reset_status_ids
+                        ):
+                            task = tasks_by_id[assignment["task_id"]]
+                            operations.append(
+                                {
+                                    "collection_name": "tasks",
+                                    "operation": "update",
+                                    "mongo_id": assignment["task_id"],
+                                    "data": {
+                                        "title": task.get("title"),
+                                        "description": task.get("description"),
+                                        "priority": task.get("priority"),
+                                        "status": TaskStatus.TODO.value,
+                                        "displayId": task.get("displayId"),
+                                        "deferredDetails": None,
+                                        "isAcknowledged": task.get("isAcknowledged", False),
+                                        "isDeleted": task.get("isDeleted", False),
+                                        "startedAt": task.get("startedAt"),
+                                        "dueAt": task.get("dueAt"),
+                                        "createdAt": task.get("createdAt"),
+                                        "createdBy": str(task.get("createdBy")),
+                                        "updatedAt": datetime.now(timezone.utc),
+                                        "updated_by": str(performed_by_user_id),
+                                    },
+                                }
+                            )
+
+                    dual_write_success = dual_write_service.batch_operations(operations)
+                    if not dual_write_success:
+                        import logging
+
+                        logger = logging.getLogger(__name__)
+                        logger.warning("Failed to sync task reassignments to Postgres")
+
+                        return False
+                return True
+            except Exception:
+                return False
diff --git a/todo/serializers/remove_from_team_serializer.py b/todo/serializers/remove_from_team_serializer.py
@@ -0,0 +1,18 @@
+from rest_framework import serializers
+from bson import ObjectId
+from todo.constants.messages import ValidationErrors
+
+
+class RemoveFromTeamSerializer(serializers.Serializer):
+    team_id = serializers.CharField()
+    user_id = serializers.CharField()
+
+    def validate_team_id(self, value):
+        if not ObjectId.is_valid(value):
+            raise serializers.ValidationError(ValidationErrors.INVALID_OBJECT_ID.format(value))
+        return value
+
+    def validate_user_id(self, value):
+        if not ObjectId.is_valid(value):
+            raise serializers.ValidationError(ValidationErrors.INVALID_OBJECT_ID.format(value))
+        return value
diff --git a/todo/services/task_assignment_service.py b/todo/services/task_assignment_service.py
@@ -150,3 +150,10 @@ def delete_task_assignment(cls, task_id: str, user_id: str) -> bool:
         Delete task assignment by task ID.
         """
         return TaskAssignmentRepository.delete_assignment(task_id, user_id)
+
+    @classmethod
+    def reassign_tasks_from_user_to_team(cls, user_id: str, team_id: str, performed_by_user_id: str):
+        """
+        Reassign all tasks of user to team
+        """
+        return TaskAssignmentRepository.reassign_tasks_from_user_to_team(user_id, team_id, performed_by_user_id)
diff --git a/todo/services/team_service.py b/todo/services/team_service.py
@@ -13,6 +13,15 @@
 from todo.models.audit_log import AuditLogModel
 from todo.repositories.audit_log_repository import AuditLogRepository
 from todo.dto.responses.error_response import ApiErrorResponse, ApiErrorDetail
+from todo.constants.role import RoleScope
+from todo.services.user_role_service import UserRoleService
+from todo.services.task_assignment_service import TaskAssignmentService
+from todo.services.user_service import UserService
+from todo.exceptions.team_exceptions import (
+    CannotRemoveOwnerException,
+    NotTeamAdminException,
+    CannotRemoveTeamPOCException,
+)
 
 DEFAULT_ROLE_ID = "1"
 
@@ -483,7 +492,25 @@ class TeamOrUserNotFound(Exception):
         pass
 
     @classmethod
-    def remove_member_from_team(cls, user_id: str, team_id: str, removed_by_user_id: str = None):
+    def _validate_remove_member_permissions(cls, user_id: str, team_id: str, removed_by_user_id: str):
+        team = TeamService.get_team_by_id(team_id)
+        team_members = UserService.get_users_by_team_id(team_id)
+        team_member_ids = [user.id for user in team_members]
+
+        if user_id not in team_member_ids:
+            raise cls.TeamOrUserNotFound
+        if user_id == team.created_by:
+            raise CannotRemoveOwnerException()
+        if user_id == team.poc_id:
+            raise CannotRemoveTeamPOCException()
+        if user_id != removed_by_user_id:
+            if not UserRoleService.has_role(removed_by_user_id, RoleName.ADMIN.value, RoleScope.TEAM.value, team_id):
+                raise NotTeamAdminException()
+
+    @classmethod
+    def remove_member_from_team(cls, user_id: str, team_id: str, removed_by_user_id: str):
+        cls._validate_remove_member_permissions(user_id, team_id, removed_by_user_id)
+
         from todo.repositories.user_team_details_repository import UserTeamDetailsRepository
 
         success = UserTeamDetailsRepository.remove_member_from_team(user_id=user_id, team_id=team_id)
@@ -494,9 +521,12 @@ def remove_member_from_team(cls, user_id: str, team_id: str, removed_by_user_id:
         AuditLogRepository.create(
             AuditLogModel(
                 team_id=PyObjectId(team_id),
-                action="member_removed_from_team",
+                action="member_removed_from_team" if user_id != removed_by_user_id else "member_left_team",
                 performed_by=PyObjectId(removed_by_user_id) if removed_by_user_id else PyObjectId(user_id),
             )
         )
 
+        UserRoleService.remove_all_user_roles_for_team(user_id, team_id)
+        TaskAssignmentService.reassign_tasks_from_user_to_team(user_id, team_id, removed_by_user_id)
+
         return True
diff --git a/todo/services/user_role_service.py b/todo/services/user_role_service.py
@@ -129,3 +129,16 @@ def get_team_users_with_roles(cls, team_id: str) -> List[Dict[str, Any]]:
         except Exception as e:
             logger.error(f"Failed to get team users with roles: {str(e)}")
             return []
+
+    @classmethod
+    def remove_all_user_roles_for_team(cls, user_id: str, team_id: str) -> bool:
+        """Remove all roles for a user within a specific team."""
+        try:
+            user_roles = cls.get_user_roles(user_id, RoleScope.TEAM.value, team_id)
+            user_role_ids = [roles["role_id"] for roles in user_roles]
+            for role_id in user_role_ids:
+                cls.remove_role_by_id(user_id, role_id, RoleScope.TEAM.value, team_id)
+            return True
+        except Exception as e:
+            logger.error(f"Failed to remove roles of user: {str(e)}")
+            return False
diff --git a/todo/views/team.py b/todo/views/team.py
