Fixing the roles and users adding issue

components/policy-mgt/org.wso2.carbon.policy.mgt.core/src/main/java/org/wso2/carbon/policy/mgt/core/dao/PolicyDAO.java

@@ -41,6 +41,8 @@ public interface PolicyDAO {
      */
     Policy addPolicyToRole(List<String> roleNames, Policy policy) throws PolicyManagerDAOException;
 
+     Policy updateRolesOfPolicy(List<String> rolesToAdd, Policy policy) throws PolicyManagerDAOException;
+
     /**
      * This method is used to add/update the users associated with the policy.
      * @param usernameList - List of the users that needs to be applied
@@ -50,6 +52,8 @@ public interface PolicyDAO {
      */
     Policy addPolicyToUser(List<String> usernameList, Policy policy) throws PolicyManagerDAOException;
 
+    Policy updateUserOfPolicy(List<String> usersToAdd, Policy policy) throws PolicyManagerDAOException;
+
     Policy addPolicyToDevice(List<Device> devices, Policy policy) throws PolicyManagerDAOException;
 
     boolean updatePolicyPriorities(List<Policy> policies) throws PolicyManagerDAOException;
@@ -114,6 +118,8 @@ public interface PolicyDAO {
 
     boolean deleteAllPolicyRelatedConfigs(int policyId) throws PolicyManagerDAOException;
 
+    boolean deleteCriteriaAndDeviceRelatedConfigs(int policyId) throws PolicyManagerDAOException;
+
     List<String> getPolicyAppliedRoles(int policyId) throws PolicyManagerDAOException;
 
     List<String> getPolicyAppliedUsers(int policyId) throws PolicyManagerDAOException;

components/policy-mgt/org.wso2.carbon.policy.mgt.core/src/main/java/org/wso2/carbon/policy/mgt/core/dao/impl/PolicyDAOImpl.java

@@ -71,8 +71,52 @@ public class PolicyDAOImpl implements PolicyDAO {
     public Policy addPolicyToRole(List<String> rolesToAdd, Policy policy) throws PolicyManagerDAOException {
         Connection conn;
         PreparedStatement insertStmt = null;
+//        PreparedStatement deleteStmt = null;
+//        final List<String> currentRoles = this.getPolicy(policy.getId()).getRoles();
+//
+//        SetReferenceTransformer<String> transformer = new SetReferenceTransformer<String>();
+//
+//        transformer.transform(currentRoles, rolesToAdd);
+//        rolesToAdd = transformer.getObjectsToAdd();
+//        List<String> rolesToDelete = transformer.getObjectsToRemove();
+        try {
+            conn = this.getConnection();
+            if (rolesToAdd.size() > 0) {
+                String query = "INSERT INTO DM_ROLE_POLICY (ROLE_NAME, POLICY_ID) VALUES (?, ?)";
+                insertStmt = conn.prepareStatement(query);
+                for (String role : rolesToAdd) {
+                    insertStmt.setString(1, role);
+                    insertStmt.setInt(2, policy.getId());
+                    insertStmt.addBatch();
+                }
+                insertStmt.executeBatch();
+            }
+//            if (rolesToDelete.size() > 0){
+//                String deleteQuery = "DELETE FROM DM_ROLE_POLICY WHERE ROLE_NAME=? AND POLICY_ID=?";
+//                deleteStmt = conn.prepareStatement(deleteQuery);
+//                for (String role : rolesToDelete) {
+//                    deleteStmt.setString(1, role);
+//                    deleteStmt.setInt(2, policy.getId());
+//                    deleteStmt.addBatch();
+//                }
+//                deleteStmt.executeBatch();
+//            }
+        } catch (SQLException e) {
+            throw new PolicyManagerDAOException("Error occurred while adding the role name with policy to database", e);
+        } finally {
+            PolicyManagementDAOUtil.cleanupResources(insertStmt, null);
+        }
+        return policy;
+    }
+
+
+    @Override
+    public Policy updateRolesOfPolicy(List<String> rolesToAdd, Policy previousPolicy) throws PolicyManagerDAOException {
+        Connection conn;
+        PreparedStatement insertStmt = null;
         PreparedStatement deleteStmt = null;
-        final List<String> currentRoles = policy.getRoles();
+
+        final List<String> currentRoles = previousPolicy.getRoles();
 
         SetReferenceTransformer<String> transformer = new SetReferenceTransformer<String>();
 
@@ -86,17 +130,17 @@ public class PolicyDAOImpl implements PolicyDAO {
                 insertStmt = conn.prepareStatement(query);
                 for (String role : rolesToAdd) {
                     insertStmt.setString(1, role);
-                    insertStmt.setInt(2, policy.getId());
+                    insertStmt.setInt(2, previousPolicy.getId());
                     insertStmt.addBatch();
                 }
                 insertStmt.executeBatch();
             }
-            if (rolesToAdd.size() > 0){
+            if (rolesToDelete.size() > 0) {
                 String deleteQuery = "DELETE FROM DM_ROLE_POLICY WHERE ROLE_NAME=? AND POLICY_ID=?";
                 deleteStmt = conn.prepareStatement(deleteQuery);
                 for (String role : rolesToDelete) {
                     deleteStmt.setString(1, role);
-                    deleteStmt.setInt(2, policy.getId());
+                    deleteStmt.setInt(2, previousPolicy.getId());
                     deleteStmt.addBatch();
                 }
                 deleteStmt.executeBatch();
@@ -105,14 +149,60 @@ public class PolicyDAOImpl implements PolicyDAO {
             throw new PolicyManagerDAOException("Error occurred while adding the role name with policy to database", e);
         } finally {
             PolicyManagementDAOUtil.cleanupResources(insertStmt, null);
+            PolicyManagementDAOUtil.cleanupResources(deleteStmt, null);
         }
-        return policy;
+        return previousPolicy;
     }
 
     @Override
     public Policy addPolicyToUser(List<String> usersToAdd, Policy policy) throws PolicyManagerDAOException {
         Connection conn;
         PreparedStatement insertStmt = null;
+//        PreparedStatement deleteStmt = null;
+//        final List<String> currentUsers = this.getPolicy(policy.getId()).getUsers();
+//
+//        SetReferenceTransformer<String> transformer = new SetReferenceTransformer<String>();
+//
+//        transformer.transform(currentUsers, usersToAdd);
+//        usersToAdd = transformer.getObjectsToAdd();
+//        List<String> usersToDelete = transformer.getObjectsToRemove();
+        try {
+            conn = this.getConnection();
+            if (usersToAdd.size() > 0) {
+                String query = "INSERT INTO DM_USER_POLICY (POLICY_ID, USERNAME) VALUES (?, ?)";
+                insertStmt = conn.prepareStatement(query);
+                for (String username : usersToAdd) {
+                    insertStmt.setInt(1, policy.getId());
+                    insertStmt.setString(2, username);
+                    insertStmt.addBatch();
+                }
+                insertStmt.executeBatch();
+            }
+//            if (usersToDelete.size() > 0){
+//                String deleteQuery = "DELETE FROM DM_USER_POLICY WHERE USERNAME=? AND POLICY_ID=?";
+//                deleteStmt = conn.prepareStatement(deleteQuery);
+//                for (String username : usersToDelete) {
+//                    deleteStmt.setString(1, username);
+//                    deleteStmt.setInt(2, policy.getId());
+//                    deleteStmt.addBatch();
+//                }
+//                deleteStmt.executeBatch();
+//            }
+
+        } catch (SQLException e) {
+            throw new PolicyManagerDAOException("Error occurred while adding the user name with policy to database", e);
+        } finally {
+            PolicyManagementDAOUtil.cleanupResources(insertStmt, null);
+//            PolicyManagementDAOUtil.cleanupResources(deleteStmt, null);
+        }
+        return policy;
+    }
+
+
+    @Override
+    public Policy updateUserOfPolicy(List<String> usersToAdd, Policy policy) throws PolicyManagerDAOException {
+        Connection conn;
+        PreparedStatement insertStmt = null;
         PreparedStatement deleteStmt = null;
         final List<String> currentUsers = policy.getUsers();
 
@@ -153,6 +243,7 @@ public class PolicyDAOImpl implements PolicyDAO {
         return policy;
     }
 
+
     @Override
     public Policy addPolicyToDevice(List<Device> devices, Policy policy) throws PolicyManagerDAOException {
         Connection conn;
@@ -1202,6 +1293,46 @@ public class PolicyDAOImpl implements PolicyDAO {
         }
     }
 
+
+    @Override
+    public boolean deleteCriteriaAndDeviceRelatedConfigs(int policyId) throws PolicyManagerDAOException {
+        Connection conn;
+        PreparedStatement stmt = null;
+        try {
+            conn = this.getConnection();
+
+//            String userPolicy = "DELETE FROM DM_USER_POLICY WHERE POLICY_ID = ?";
+//            stmt = conn.prepareStatement(userPolicy);
+//            stmt.setInt(1, policyId);
+//            stmt.executeUpdate();
+//
+//            String rolePolicy = "DELETE FROM DM_ROLE_POLICY WHERE POLICY_ID = ?";
+//            stmt = conn.prepareStatement(rolePolicy);
+//            stmt.setInt(1, policyId);
+//            stmt.executeUpdate();
+
+            String devicePolicy = "DELETE FROM DM_DEVICE_POLICY WHERE POLICY_ID = ?";
+            stmt = conn.prepareStatement(devicePolicy);
+            stmt.setInt(1, policyId);
+            stmt.executeUpdate();
+
+            String deleteCriteria = "DELETE FROM DM_POLICY_CRITERIA WHERE POLICY_ID = ?";
+            stmt = conn.prepareStatement(deleteCriteria);
+            stmt.setInt(1, policyId);
+            stmt.executeUpdate();
+
+            if (log.isDebugEnabled()) {
+                log.debug("Policy (" + policyId + ") related configs deleted from database.");
+            }
+            return true;
+        } catch (SQLException e) {
+            throw new PolicyManagerDAOException("Unable to delete the policy (" + policyId +
+                    ") related configs from database", e);
+        } finally {
+            PolicyManagementDAOUtil.cleanupResources(stmt, null);
+        }
+    }
+
     private Connection getConnection() throws PolicyManagerDAOException {
         return PolicyManagementDAOFactory.getConnection();
     }
@@ -1345,7 +1476,8 @@ public class PolicyDAOImpl implements PolicyDAO {
         int tenantId = PrivilegedCarbonContext.getThreadLocalCarbonContext().getTenantId();
         try {
             conn = this.getConnection();
-            String query = "SELECT * FROM DM_DEVICE_POLICY_APPLIED WHERE DEVICE_ID = ? AND TENANT_ID = ? AND ENROLMENT_ID = ?";
+            String query = "SELECT * FROM DM_DEVICE_POLICY_APPLIED WHERE DEVICE_ID = ? AND TENANT_ID = ? AND " +
+                    "ENROLMENT_ID = ?";
             stmt = conn.prepareStatement(query);
             stmt.setInt(1, deviceId);
             stmt.setInt(2, tenantId);

components/policy-mgt/org.wso2.carbon.policy.mgt.core/src/main/java/org/wso2/carbon/policy/mgt/core/mgt/impl/PolicyManagerImpl.java

@@ -180,15 +180,15 @@ public class PolicyManagerImpl implements PolicyManager {
             if (!newFeaturesList.isEmpty()) {
                 featureDAO.addProfileFeatures(newFeaturesList, profileId);
             }
-            policyDAO.deleteAllPolicyRelatedConfigs(policy.getId());
+            policyDAO.deleteCriteriaAndDeviceRelatedConfigs(policy.getId());
 
 
             if (policy.getUsers() != null) {
-                policyDAO.addPolicyToUser(policy.getUsers(), previousPolicy);
+                policyDAO.updateUserOfPolicy(policy.getUsers(), previousPolicy);
             }
 
             if (policy.getRoles() != null) {
-                policyDAO.addPolicyToRole(policy.getRoles(), previousPolicy);
+                policyDAO.updateRolesOfPolicy(policy.getRoles(), previousPolicy);
             }
 
             if (policy.getDevices() != null) {