Fixing the feature remove from the policy

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

@@ -123,4 +123,6 @@ public interface FeatureDAO {
      */
     boolean deleteFeaturesOfProfile(int profileId) throws FeatureManagerDAOException;
 
+    boolean deleteProfileFeatures(int featureId) throws FeatureManagerDAOException;
+
 }

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

@@ -20,7 +20,6 @@ package org.wso2.carbon.policy.mgt.core.dao.impl;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
-
 import org.wso2.carbon.context.PrivilegedCarbonContext;
 import org.wso2.carbon.device.mgt.common.Feature;
 import org.wso2.carbon.policy.mgt.common.Profile;
@@ -28,14 +27,16 @@ import org.wso2.carbon.policy.mgt.common.ProfileFeature;
 import org.wso2.carbon.policy.mgt.core.dao.FeatureDAO;
 import org.wso2.carbon.policy.mgt.core.dao.FeatureManagerDAOException;
 import org.wso2.carbon.policy.mgt.core.dao.PolicyManagementDAOFactory;
-import org.wso2.carbon.policy.mgt.core.dao.PolicyManagerDAOException;
 import org.wso2.carbon.policy.mgt.core.dao.util.PolicyManagementDAOUtil;
 import org.wso2.carbon.policy.mgt.core.util.PolicyManagerUtil;
 
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.io.ObjectInputStream;
-import java.sql.*;
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
 import java.util.ArrayList;
 import java.util.List;
 
@@ -177,6 +178,29 @@ public class FeatureDAOImpl implements FeatureDAO {
         }
     }
 
+
+    @Override
+    public boolean deleteProfileFeatures(int featureId) throws FeatureManagerDAOException {
+        Connection conn;
+        PreparedStatement stmt = null;
+        int tenantId = PrivilegedCarbonContext.getThreadLocalCarbonContext().getTenantId();
+        try {
+            conn = this.getConnection();
+            String query = "DELETE FROM DM_PROFILE_FEATURES WHERE ID = ? AND TENANT_ID = ?";
+            stmt = conn.prepareStatement(query);
+            stmt.setInt(1, featureId);
+            stmt.setInt(2, tenantId);
+            if (stmt.executeUpdate() > 0) {
+                return true;
+            }
+            return false;
+        } catch (SQLException e) {
+            throw new FeatureManagerDAOException("Error occurred while deleting the feature related to a profile.", e);
+        } finally {
+            PolicyManagementDAOUtil.cleanupResources(stmt, null);
+        }
+    }
+
     @Override
     public List<ProfileFeature> getAllProfileFeatures() throws FeatureManagerDAOException {
         Connection conn;

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

@@ -134,7 +134,7 @@ public class PolicyManagerImpl implements PolicyManager {
     public Policy updatePolicy(Policy policy) throws PolicyManagementException {
 
         try {
-            // Previous policy needs to be obtained before begining the transaction
+            // Previous policy needs to be obtained before beginning the transaction
             Policy previousPolicy = this.getPolicy(policy.getId());
 
             PolicyManagementDAOFactory.beginTransaction();
@@ -144,7 +144,9 @@ public class PolicyManagerImpl implements PolicyManager {
 
             List<ProfileFeature> existingFeaturesList = new ArrayList<>();
             List<ProfileFeature> newFeaturesList = new ArrayList<>();
+            List<ProfileFeature> feturesToDelete = new ArrayList<>();
             List<String> temp = new ArrayList<>();
+            List<String> updateDFes = new ArrayList<>();
 
             List<ProfileFeature> updatedFeatureList = policy.getProfile().getProfileFeaturesList();
 
@@ -158,6 +160,14 @@ public class PolicyManagerImpl implements PolicyManager {
                         temp.add(feature.getFeatureCode());
                     }
                 }
+                updateDFes.add(feature.getFeatureCode());
+            }
+
+            // Check for the features to delete
+            for(ProfileFeature feature : existingProfileFeaturesList) {
+                if(!updateDFes.contains(feature.getFeatureCode())){
+                    feturesToDelete.add(feature);
+                }
             }
 
             // Checks for the new features
@@ -180,6 +190,12 @@ public class PolicyManagerImpl implements PolicyManager {
             if (!newFeaturesList.isEmpty()) {
                 featureDAO.addProfileFeatures(newFeaturesList, profileId);
             }
+
+            if(!feturesToDelete.isEmpty()){
+                for (ProfileFeature pf : feturesToDelete)
+                featureDAO.deleteProfileFeatures(pf.getId());
+            }
+
             policyDAO.deleteCriteriaAndDeviceRelatedConfigs(policy.getId());