Refactored addOperation method

components/device-mgt/org.wso2.carbon.device.mgt.core/src/main/java/org/wso2/carbon/device/mgt/core/dao/DeviceDAO.java

@@ -94,6 +94,7 @@ public interface DeviceDAO {
 
     /**
      * This method is used to retrieve devices of a given user.
+     *
      * @param username user name.
      * @param tenantId tenant id.
      * @return returns list of devices.
@@ -112,6 +113,7 @@ public interface DeviceDAO {
 
     /**
      * This method is used to retrieve devices of a given device name.
+     *
      * @param deviceName device name.
      * @param tenantId   tenant id.
      * @return returns list of devices.
@@ -187,5 +189,17 @@ public interface DeviceDAO {
      */
     int getEnrolmentByStatus(DeviceIdentifier deviceId, Status status,
                              int tenantId) throws DeviceManagementDAOException;
+
+    /**
+     * This method is used to retrieve the enrollment info of a given list of devices and status.
+     *
+     * @param deviceIds A list of device identifiers.
+     * @param status    enrollment status.
+     * @param tenantId  tenant id.
+     * @return returns a list of enrolment info objects.
+     * @throws DeviceManagementDAOException
+     */
+    List<EnrolmentInfo> getEnrolmentsByStatus(List<DeviceIdentifier> deviceIds, Status status,
+                                              int tenantId) throws DeviceManagementDAOException;
 }
 

components/device-mgt/org.wso2.carbon.device.mgt.core/src/main/java/org/wso2/carbon/device/mgt/core/dao/impl/DeviceDAOImpl.java

@@ -30,6 +30,7 @@ import org.wso2.carbon.device.mgt.core.dao.util.DeviceManagementDAOUtil;
 import java.sql.*;
 import java.util.ArrayList;
 import java.util.Date;
+import java.util.Iterator;
 import java.util.List;
 
 public class DeviceDAOImpl implements DeviceDAO {
@@ -455,6 +456,51 @@ public class DeviceDAOImpl implements DeviceDAO {
         }
     }
 
+    public List<EnrolmentInfo> getEnrolmentsByStatus(List<DeviceIdentifier> deviceIds, Status status,
+                                                     int tenantId) throws DeviceManagementDAOException {
+        Connection conn;
+        PreparedStatement stmt = null;
+        ResultSet rs = null;
+        List<EnrolmentInfo> enrolments = new ArrayList<>();
+        try {
+            conn = this.getConnection();
+            StringBuilder sql = new StringBuilder();
+            sql.append("SELECT e.ID AS ENROLMENT_ID, e.OWNER, e.OWNERSHIP, e.DATE_OF_ENROLMENT, e.DATE_OF_LAST_UPDATE, " +
+                    "e.STATUS FROM DM_ENROLMENT e WHERE e.DEVICE_ID IN (SELECT d.ID FROM DM_DEVICE d " +
+                    "WHERE d.DEVICE_IDENTIFICATION IN (");
+
+            // adding arguments to the sql query
+            Iterator iterator = deviceIds.iterator();
+            while (iterator.hasNext()) {
+                iterator.next();
+                sql.append(" ?");
+                if (iterator.hasNext()) {
+                    sql.append(",");
+                }
+            }
+            sql.append(") AND d.TENANT_ID = ?) AND e.STATUS = ? AND e.TENANT_ID = ?");
+
+            stmt = conn.prepareStatement(sql.toString());
+            int index = 1;
+            for (DeviceIdentifier id : deviceIds) {
+                stmt.setString(index++, id.getId());
+            }
+            stmt.setInt(index++, tenantId);
+            stmt.setString(index++, status.toString());
+            stmt.setInt(index, tenantId);
+            rs = stmt.executeQuery();
+            if (rs.next()) {
+                enrolments.add(this.loadEnrolment(rs));
+            }
+            return enrolments;
+        } catch (SQLException e) {
+            throw new DeviceManagementDAOException("Error occurred while retrieving the enrolment " +
+                    "ids of devices", e);
+        } finally {
+            DeviceManagementDAOUtil.cleanupResources(stmt, rs);
+        }
+    }
+
     private Device loadDevice(ResultSet rs) throws SQLException {
         Device device = new Device();
         device.setId(rs.getInt("DEVICE_ID"));

components/device-mgt/org.wso2.carbon.device.mgt.core/src/main/java/org/wso2/carbon/device/mgt/core/operation/mgt/OperationManagerImpl.java

@@ -80,19 +80,11 @@ public class OperationManagerImpl implements OperationManager {
             }
         }
         try {
-            OperationManagementDAOFactory.beginTransaction();
-
-            org.wso2.carbon.device.mgt.core.dto.operation.mgt.Operation operationDto =
-                    OperationDAOUtil.convertOperation(operation);
-
-            int operationId = this.lookupOperationDAO(operation).addOperation(operationDto);
-
-            int enrolmentId;
             int tenantId = CarbonContext.getThreadLocalCarbonContext().getTenantId();
-            for (DeviceIdentifier deviceId : deviceIds) {
+            List<EnrolmentInfo> enrolments;
             try {
                 DeviceManagementDAOFactory.openConnection();
-                    enrolmentId = deviceDAO.getEnrolmentByStatus(deviceId, EnrolmentInfo.Status.ACTIVE, tenantId);
+                enrolments = deviceDAO.getEnrolmentsByStatus(deviceIds, EnrolmentInfo.Status.ACTIVE, tenantId);
             } catch (SQLException e) {
                 throw new OperationManagementException("Error occurred while opening a connection the data " +
                         "source", e);
@@ -100,14 +92,14 @@ public class OperationManagerImpl implements OperationManager {
                 DeviceManagementDAOFactory.closeConnection();
             }
 
-                if (enrolmentId < 0) {
-                    String errorMsg = "The operation not added for device.The device not found for " +
-                            "device Identifier type -'" + deviceId.getType() + "' and device Id '" +
-                            deviceId.getId();
-                    log.error(errorMsg);
-                } else {
-                    operationMappingDAO.addOperationMapping(operationId, enrolmentId);
-                }
+            OperationManagementDAOFactory.beginTransaction();
+            org.wso2.carbon.device.mgt.core.dto.operation.mgt.Operation operationDto =
+                    OperationDAOUtil.convertOperation(operation);
+
+            int operationId = this.lookupOperationDAO(operation).addOperation(operationDto);
+
+            for (EnrolmentInfo enrolmentInfo : enrolments) {
+                operationMappingDAO.addOperationMapping(operationId, enrolmentInfo.getId());
             }
             OperationManagementDAOFactory.commitTransaction();
             return operationId;

components/device-mgt/org.wso2.carbon.device.mgt.core/src/main/java/org/wso2/carbon/device/mgt/core/service/DeviceManagementProviderServiceImpl.java

@@ -245,6 +245,12 @@ public class DeviceManagementProviderServiceImpl implements DeviceManagementProv
             DeviceManagementDAOFactory.beginTransaction();
 
             Device device = deviceDAO.getDevice(deviceId, tenantId);
+            if (device == null) {
+                if (log.isDebugEnabled()) {
+                    log.debug("Device not found for id '" + deviceId.getId() + "'");
+                }
+                throw new DeviceManagementException("Device not found");
+            }
             DeviceType deviceType = deviceTypeDAO.getDeviceType(device.getType());
 
             device.getEnrolmentInfo().setDateOfLastUpdate(new Date().getTime());