[Fix-17798] Fix assign workergroup to project is incorrect (#17799)

Author: Mrhs121

dolphinscheduler-api/src/main/java/org/apache/dolphinscheduler/api/service/impl/ProjectWorkerGroupRelationServiceImpl.java

@@ -94,7 +94,30 @@ public Result assignWorkerGroupsToProject(User loginUser, Long projectCode, List
             return result;
         }
 
+        Project project = projectMapper.queryByCode(projectCode);
+        if (Objects.isNull(project)) {
+            putMsg(result, Status.PROJECT_NOT_EXIST);
+            return result;
+        }
+
+        /*
+         * Todo : For modification operations on projects, we should acquire project row locks. All project-related
+         * operations and modification/creation actions for workflows/task definitions within the project require
+         * acquiring row locks first
+         */
         if (CollectionUtils.isEmpty(workerGroups)) {
+            Set<String> projectWorkerGroupNames =
+                    projectWorkerGroupDao.queryAssignedWorkerGroupNamesByProjectCode(projectCode);
+            if (CollectionUtils.isNotEmpty(projectWorkerGroupNames)) {
+                Set<String> usedWorkerGroups = getAllUsedWorkerGroups(project);
+                if (CollectionUtils.isNotEmpty(usedWorkerGroups)) {
+                    Set<String> usedInProject = SetUtils.intersection(usedWorkerGroups, projectWorkerGroupNames);
+                    if (!usedInProject.isEmpty()) {
+                        throw new ServiceException(Status.USED_WORKER_GROUP_EXISTS, usedInProject);
+                    }
+                }
+            }
+
             boolean deleted = projectWorkerGroupDao.deleteByProjectCode(projectCode);
             if (deleted) {
                 putMsg(result, Status.SUCCESS);
@@ -104,12 +127,6 @@ public Result assignWorkerGroupsToProject(User loginUser, Long projectCode, List
             return result;
         }
 
-        Project project = projectMapper.queryByCode(projectCode);
-        if (Objects.isNull(project)) {
-            putMsg(result, Status.PROJECT_NOT_EXIST);
-            return result;
-        }
-
         Set<String> allWorkerGroupNames = new HashSet<>(workerGroupDao.queryAllWorkerGroupNames());
         workerGroupService.getConfigWorkerGroupPageDetail().forEach(
                 workerGroupPageDetail -> allWorkerGroupNames.add(workerGroupPageDetail.getName()));
@@ -131,9 +148,11 @@ public Result assignWorkerGroupsToProject(User loginUser, Long projectCode, List
 
         if (CollectionUtils.isNotEmpty(needDeletedWorkerGroups)) {
             Set<String> usedWorkerGroups = getAllUsedWorkerGroups(project);
-            if (CollectionUtils.isNotEmpty(usedWorkerGroups) && usedWorkerGroups.containsAll(needDeletedWorkerGroups)) {
-                throw new ServiceException(Status.USED_WORKER_GROUP_EXISTS,
-                        SetUtils.intersection(usedWorkerGroups, needDeletedWorkerGroups).toSet());
+            if (CollectionUtils.isNotEmpty(usedWorkerGroups) && CollectionUtils.isNotEmpty(needDeletedWorkerGroups)) {
+                Set<String> shouldNotDelete = SetUtils.intersection(usedWorkerGroups, needDeletedWorkerGroups);
+                if (CollectionUtils.isNotEmpty(shouldNotDelete)) {
+                    throw new ServiceException(Status.USED_WORKER_GROUP_EXISTS, shouldNotDelete);
+                }
             }
             boolean deleted =
                     projectWorkerGroupDao.deleteByProjectCodeAndWorkerGroups(projectCode,
@@ -147,6 +166,7 @@ public Result assignWorkerGroupsToProject(User loginUser, Long projectCode, List
                 throw new ServiceException(Status.ASSIGN_WORKER_GROUP_TO_PROJECT_ERROR);
             }
         }
+
         Set<String> needAssignedWorkerGroups =
                 SetUtils.difference(unauthorizedWorkerGroupNames, projectWorkerGroupNames);
         if (CollectionUtils.isNotEmpty(needAssignedWorkerGroups)) {

dolphinscheduler-api/src/test/java/org/apache/dolphinscheduler/api/service/ProjectWorkerGroupRelationServiceTest.java

@@ -138,6 +138,13 @@ public void testAssignWorkerGroupsToProject() {
                 getWorkerGroups());
         Assertions.assertEquals(Status.SUCCESS.getCode(), result.getCode());
 
+        // success when no task referenced any wg
+        Mockito.when(projectWorkerGroupDao.deleteByProjectCode(projectCode))
+                .thenReturn(true);
+        result = projectWorkerGroupRelationService.assignWorkerGroupsToProject(loginUser, projectCode,
+                new ArrayList<>());
+        Assertions.assertEquals(Status.SUCCESS.getCode(), result.getCode());
+
         // db deletion fail
         Mockito.when(projectWorkerGroupDao.deleteByProjectCodeAndWorkerGroups(Mockito.any(), Mockito.any()))
                 .thenReturn(false);
@@ -154,6 +161,26 @@ public void testAssignWorkerGroupsToProject() {
         AssertionsHelper.assertThrowsServiceException(Status.USED_WORKER_GROUP_EXISTS,
                 () -> projectWorkerGroupRelationService.assignWorkerGroupsToProject(loginUser, projectCode,
                         getUnusedWorkerGroups()));
+
+        // test clear all wg and fail when wg is referenced by task definition
+        // test case: project all wg: test, task used wg: test, new wg: null
+        Mockito.when(taskDefinitionDao.queryAllTaskDefinitionWorkerGroups(Mockito.anyLong()))
+                .thenReturn(Collections.singletonList(getProjectWorkerGroup().getWorkerGroup()));
+        Mockito.when(projectWorkerGroupDao.queryAssignedWorkerGroupNamesByProjectCode(Mockito.any()))
+                .thenReturn(Sets.newHashSet(getProjectWorkerGroup().getWorkerGroup()));
+        AssertionsHelper.assertThrowsServiceException(Status.USED_WORKER_GROUP_EXISTS,
+                () -> projectWorkerGroupRelationService.assignWorkerGroupsToProject(loginUser, projectCode,
+                        new ArrayList<>()));
+
+        // test delete superset of the used wg collection and fail when wg is referenced by task definition
+        // test case: project all wg: test,test1,test2. task used wg: test. new wg: test1, delete test2 and test
+        Mockito.when(taskDefinitionDao.queryAllTaskDefinitionWorkerGroups(Mockito.anyLong()))
+                .thenReturn(Collections.singletonList(getProjectWorkerGroup().getWorkerGroup()));
+        Mockito.when(projectWorkerGroupDao.queryAssignedWorkerGroupNamesByProjectCode(Mockito.any()))
+                .thenReturn(Sets.newHashSet("test", "test1", "test2"));
+        AssertionsHelper.assertThrowsServiceException(Status.USED_WORKER_GROUP_EXISTS,
+                () -> projectWorkerGroupRelationService.assignWorkerGroupsToProject(loginUser, projectCode,
+                        getUnusedWorkerGroups()));
     }
 
     @Test