server,engine-schema: unlink userdata during acount cleanup

Fixes #9477

Signed-off-by: Abhishek Kumar <abhishek.mrt22@gmail.com>

Author: shwstppr

engine/components-api/src/main/java/com/cloud/template/TemplateManager.java

@@ -149,4 +149,6 @@ public interface TemplateManager {
     static Boolean getValidateUrlIsResolvableBeforeRegisteringTemplateValue() {
         return ValidateUrlIsResolvableBeforeRegisteringTemplate.value();
     }
+
+    void unlinkAllAndCleanupUserDataFromAccount(long accountId);
 }

engine/schema/src/main/java/com/cloud/storage/dao/VMTemplateDao.java

@@ -103,4 +103,6 @@ public interface VMTemplateDao extends GenericDao<VMTemplateVO, Long>, StateDao<
     List<Long> listIdsByTemplateTag(String tag);
 
     List<Long> listIdsByExtensionId(long extensionId);
+
+    void unlinkUserdataFromTemplate(List<Long> userdataIds);
 }

engine/schema/src/main/java/com/cloud/storage/dao/VMTemplateDaoImpl.java

@@ -917,4 +917,20 @@ public boolean updateState(
         }
         return rows > 0;
     }
+
+    @Override
+    public void unlinkUserdataFromTemplate(List<Long> userdataIds) {
+        if (CollectionUtils.isEmpty(userdataIds)) {
+            return;
+        }
+        SearchBuilder<VMTemplateVO> sb = createSearchBuilder();
+        sb.and("userDataId", sb.entity().getUserDataId(), SearchCriteria.Op.IN);
+        sb.done();
+        SearchCriteria<VMTemplateVO> sc = sb.create();
+        sc.setParameters("userDataId", userdataIds.toArray());
+        VMTemplateVO vo = createForUpdate();
+        vo.setUserDataId(null);
+        vo.setUserDataLinkPolicy(null);
+        update(vo, sc);
+    }
 }

engine/schema/src/main/java/com/cloud/user/dao/UserDataDao.java

@@ -16,6 +16,8 @@
 // under the License.
 package com.cloud.user.dao;
 
+import java.util.List;
+
 import com.cloud.user.UserDataVO;
 import com.cloud.utils.db.GenericDao;
 
@@ -25,6 +27,8 @@ public interface UserDataDao extends GenericDao<UserDataVO, Long> {
 
     public UserDataVO findByName(long accountId, long domainId, String name);
 
-    int removeByAccountId(long accountId);
+    List<Long> listIdsByAccountId(long accountId);
+
+    int removeAllIds(List<Long> ids);
 
 }

engine/schema/src/main/java/com/cloud/user/dao/UserDataDaoImpl.java

@@ -16,10 +16,15 @@
 // under the License.
 package com.cloud.user.dao;
 
+import java.util.List;
+
 import com.cloud.user.UserDataVO;
 import com.cloud.utils.db.GenericDaoBase;
+import com.cloud.utils.db.GenericSearchBuilder;
 import com.cloud.utils.db.SearchBuilder;
 import com.cloud.utils.db.SearchCriteria;
+
+import org.apache.commons.collections.CollectionUtils;
 import org.springframework.stereotype.Component;
 
 @Component
@@ -65,9 +70,26 @@ public UserDataVO findByName(long accountId, long domainId, String name) {
     }
 
     @Override
-    public int removeByAccountId(long accountId) {
-        SearchCriteria<UserDataVO> sc = userdataSearch.create();
+    public List<Long> listIdsByAccountId(long accountId) {
+        GenericSearchBuilder<UserDataVO, Long> sb = createSearchBuilder(Long.class);
+        sb.selectFields(sb.entity().getId());
+        sb.and("accountId", sb.entity().getAccountId(), SearchCriteria.Op.EQ);
+        sb.done();
+        SearchCriteria<Long> sc = sb.create();
         sc.setParameters("accountId", accountId);
+        return customSearch(sc, null);
+    }
+
+    @Override
+    public int removeAllIds(List<Long> ids) {
+        if (CollectionUtils.isEmpty(ids)) {
+            return 0;
+        }
+        SearchBuilder<UserDataVO> sb = createSearchBuilder();
+        sb.and("idIn", sb.entity().getId(), SearchCriteria.Op.IN);
+        sb.done();
+        SearchCriteria<UserDataVO> sc = sb.create();
+        sc.setParameters("idIn", ids.toArray());
         return remove(sc);
     }
 }

server/src/main/java/com/cloud/template/TemplateManagerImpl.java

@@ -193,6 +193,7 @@
 import com.cloud.user.User;
 import com.cloud.user.UserData;
 import com.cloud.user.dao.AccountDao;
+import com.cloud.user.dao.UserDataDao;
 import com.cloud.uservm.UserVm;
 import com.cloud.utils.DateUtil;
 import com.cloud.utils.EncryptionUtil;
@@ -301,6 +302,8 @@ public class TemplateManagerImpl extends ManagerBase implements TemplateManager,
     private VMTemplateDetailsDao _tmpltDetailsDao;
     @Inject
     private HypervisorGuruManager _hvGuruMgr;
+    @Inject
+    UserDataDao userDataDao;
 
     private List<TemplateAdapter> _adapters;
 
@@ -2495,4 +2498,14 @@ public VirtualMachineTemplate linkUserDataToTemplate(LinkUserDataToTemplateCmd c
 
         return _tmpltDao.findById(template.getId());
     }
+
+    @Override
+    public void unlinkAllAndCleanupUserDataFromAccount(long accountId) {
+        List<Long> userDataIds = userDataDao.listIdsByAccountId(accountId);
+        if (CollectionUtils.isEmpty(userDataIds)) {
+            return;
+        }
+        _tmpltDao.unlinkUserdataFromTemplate(userDataIds);
+        userDataDao.removeAllIds(userDataIds);
+    }
 }

server/src/main/java/com/cloud/user/AccountManagerImpl.java

@@ -1201,7 +1201,7 @@ public int compare(NetworkVO network1, NetworkVO network2) {
             }
 
             // Delete registered UserData
-            userDataDao.removeByAccountId(accountId);
+            _tmpltMgr.unlinkAllAndCleanupUserDataFromAccount(accountId);
 
             // Delete Webhooks
             deleteWebhooksForAccount(accountId);

server/src/test/java/com/cloud/template/TemplateManagerImplTest.java

@@ -63,6 +63,7 @@
 import com.cloud.user.UserData;
 import com.cloud.user.UserVO;
 import com.cloud.user.dao.AccountDao;
+import com.cloud.user.dao.UserDataDao;
 import com.cloud.utils.component.ComponentContext;
 import com.cloud.utils.concurrency.NamedThreadFactory;
 import com.cloud.utils.exception.CloudRuntimeException;
@@ -203,14 +204,19 @@ public class TemplateManagerImplTest {
 
     @Inject
     AccountManager _accountMgr;
+
     @Inject
     VnfTemplateManager vnfTemplateManager;
+
     @Inject
     SnapshotJoinDao snapshotJoinDao;
 
     @Inject
     HeuristicRuleHelper heuristicRuleHelperMock;
 
+    @Inject
+    UserDataDao userDataDao;
+
     public class CustomThreadPoolExecutor extends ThreadPoolExecutor {
         AtomicInteger ai = new AtomicInteger(0);
         public CustomThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit,
@@ -978,11 +984,17 @@ public SecondaryStorageHeuristicDao secondaryStorageHeuristicDao() {
         public HeuristicRuleHelper heuristicRuleHelper() {
             return Mockito.mock(HeuristicRuleHelper.class);
         }
+
         @Bean
         public SnapshotJoinDao snapshotJoinDao() {
             return Mockito.mock(SnapshotJoinDao.class);
         }
 
+        @Bean
+        public UserDataDao userDataDao() {
+            return Mockito.mock(UserDataDao.class);
+        }
+
 
         public static class Library implements TypeFilter {
             @Override

server/src/test/java/com/cloud/user/AccountManagerImplTest.java

@@ -200,7 +200,6 @@ public void deleteUserAccount() {
         sshkeyList.add(sshkey);
         Mockito.when(_sshKeyPairDao.listKeyPairs(Mockito.anyLong(), Mockito.anyLong())).thenReturn(sshkeyList);
         Mockito.when(_sshKeyPairDao.remove(Mockito.anyLong())).thenReturn(true);
-        Mockito.when(userDataDao.removeByAccountId(Mockito.anyLong())).thenReturn(222);
         Mockito.doNothing().when(accountManagerImpl).deleteWebhooksForAccount(Mockito.anyLong());
         Mockito.doNothing().when(accountManagerImpl).verifyCallerPrivilegeForUserOrAccountOperations((Account) any());