Setting maintenance user when carbon context user is null + fixing the function which sets statuses of all devices belonging to a user

components/device-mgt/org.wso2.carbon.device.mgt.common/src/main/java/org/wso2/carbon/device/mgt/common/DeviceManagementConstants.java

@@ -142,4 +142,9 @@ public final class DeviceManagementConstants {
         public static final String GEOFENCE = "GEOFENCE";
     }
 
+    public static final class MaintenanceProperties {
+        private MaintenanceProperties() {throw new AssertionError();}
+        public static final String MAINTENANCE_USER = "system";
+    }
+
 }

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

@@ -34,8 +34,11 @@ public interface EnrollmentDAO {
 
     int removeEnrollment(int deviceId, String currentOwner, int tenantId) throws DeviceManagementDAOException;
 
+    @Deprecated
     boolean setStatus(String currentOwner, Status status, int tenantId) throws DeviceManagementDAOException;
 
+    boolean setStatusAllDevices(String currentOwner, Status status, int tenantId) throws DeviceManagementDAOException;
+
     boolean setStatus(int enrolmentId, Status status, int tenantId) throws DeviceManagementDAOException;
 
     Status getStatus(int deviceId, String currentOwner, int tenantId) throws DeviceManagementDAOException;

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

@@ -20,10 +20,10 @@ package org.wso2.carbon.device.mgt.core.dao.impl;
 
 import org.wso2.carbon.context.PrivilegedCarbonContext;
 import org.wso2.carbon.device.mgt.common.Device;
+import org.wso2.carbon.device.mgt.common.DeviceManagementConstants;
 import org.wso2.carbon.device.mgt.common.EnrolmentInfo;
 import org.wso2.carbon.device.mgt.core.dao.DeviceManagementDAOException;
 import org.wso2.carbon.device.mgt.core.dao.DeviceManagementDAOFactory;
-import org.wso2.carbon.device.mgt.core.dao.DeviceStatusDAO;
 import org.wso2.carbon.device.mgt.core.dao.EnrollmentDAO;
 import org.wso2.carbon.device.mgt.core.dao.util.DeviceManagementDAOUtil;
 
@@ -200,29 +200,34 @@ public class EnrollmentDAOImpl implements EnrollmentDAO {
     @Override
     public boolean setStatus(String currentOwner, EnrolmentInfo.Status status,
                              int tenantId) throws DeviceManagementDAOException {
+        return setStatusAllDevices(currentOwner, status, tenantId);
+    }
+
+    @Override
+    public boolean setStatusAllDevices(String currentOwner, EnrolmentInfo.Status status, int tenantId)
+            throws DeviceManagementDAOException{
         Connection conn;
         PreparedStatement stmt = null;
+        Timestamp updateTime = new Timestamp(new Date().getTime());
         if(getCountOfDevicesOfOwner(currentOwner, tenantId) > 0){
             try {
                 conn = this.getConnection();
-                // TODO add DATE_OF_LAST_UPDATE
+                String sql = "UPDATE DM_ENROLMENT SET STATUS = ?, DATE_OF_LAST_UPDATE = ? WHERE OWNER = ? AND TENANT_ID = ?";
-                String sql = "UPDATE DM_ENROLMENT SET STATUS = ? WHERE OWNER = ? AND TENANT_ID = ?";
                 stmt = conn.prepareStatement(sql);
                 stmt.setString(1, status.toString());
-                stmt.setString(2, currentOwner);
+                stmt.setTimestamp(2, updateTime);
-                stmt.setInt(3, tenantId);
+                stmt.setString(3, currentOwner);
+                stmt.setInt(4, tenantId);
                 stmt.executeUpdate();
             } catch (SQLException e) {
                 throw new DeviceManagementDAOException("Error occurred while setting the status of device enrolment", e);
             } finally {
                 DeviceManagementDAOUtil.cleanupResources(stmt, null);
             }
-            return true;
+            return addDeviceStatus(currentOwner, status, tenantId);
         } else {
             return false;
         }
-        // TODO: Needs device Id since having owner id doesn't necessary make it unique?
-        //getDeviceStatusDAO().updateStatus(deviceId, status);
     }
 
     @Override
@@ -257,9 +262,75 @@ public class EnrollmentDAOImpl implements EnrollmentDAO {
         return addDeviceStatus(config.getId(), config.getStatus());
     }
 
+    private boolean addDeviceStatus(String currentOwner, EnrolmentInfo.Status status, int tenantId) throws DeviceManagementDAOException {
+        Connection conn;
+        String changedBy = PrivilegedCarbonContext.getThreadLocalCarbonContext().getUsername();
+        if (changedBy == null){
+            changedBy = DeviceManagementConstants.MaintenanceProperties.MAINTENANCE_USER;
+        }
+        PreparedStatement stmt = null;
+        ResultSet rs = null;
+        List<int[]> enrolmentInfoList = new ArrayList<>();
+        try {
+            conn = this.getConnection();
+            String sql = "SELECT ID, DEVICE_ID, OWNER, OWNERSHIP, STATUS, IS_TRANSFERRED, DATE_OF_ENROLMENT, " +
+                    "DATE_OF_LAST_UPDATE, TENANT_ID FROM DM_ENROLMENT WHERE OWNER = ? AND TENANT_ID = ?";
+            stmt = conn.prepareStatement(sql);
+            stmt.setString(1, currentOwner);
+            stmt.setInt(2, tenantId);
+            rs = stmt.executeQuery();
+            while (rs.next()) {
+                int enrolmentId = rs.getInt("ID");
+                int deviceId = rs.getInt("DEVICE_ID");
+                enrolmentInfoList.add(new int[]{enrolmentId, deviceId});
+            }
+            DeviceManagementDAOUtil.cleanupResources(stmt, rs);
+            Timestamp updateTime = new Timestamp(new Date().getTime());
+            sql = "INSERT INTO DM_DEVICE_STATUS (ENROLMENT_ID, DEVICE_ID, STATUS, UPDATE_TIME, CHANGED_BY) VALUES(?, ?, ?, ?, ?)";
+            try (PreparedStatement ps = conn.prepareStatement(sql)) {
+                if (conn.getMetaData().supportsBatchUpdates()) {
+                    for(int[] info: enrolmentInfoList){
+                        ps.setInt(1, info[0]);
+                        ps.setInt(2, info[1]);
+                        ps.setString(3, status.toString());
+                        ps.setTimestamp(4, updateTime);
+                        ps.setString(5, changedBy);
+                        ps.addBatch();
+                    }
+                    int[] batchResult = ps.executeBatch();
+                    for (int i : batchResult) {
+                        if (i == 0 || i == Statement.SUCCESS_NO_INFO || i == Statement.EXECUTE_FAILED) {
+                            return false;
+                        }
+                    }
+                } else {
+                    for(int[] info: enrolmentInfoList){
+                        ps.setInt(1, info[0]);
+                        ps.setInt(2, info[1]);
+                        ps.setString(3, status.toString());
+                        ps.setTimestamp(4, updateTime);
+                        ps.setString(5, changedBy);
+                        ps.execute();
+                    }
+
+                }
+            }
+
+        } catch (SQLException e) {
+            throw new DeviceManagementDAOException("Error occurred while retrieving the enrolments " +
+                    "information of owner '" + currentOwner + "'", e);
+        } finally {
+            DeviceManagementDAOUtil.cleanupResources(stmt, rs);
+        }
+        return true;
+    }
+
     private boolean addDeviceStatus(int enrolmentId, EnrolmentInfo.Status status) throws DeviceManagementDAOException {
         Connection conn;
         String changedBy = PrivilegedCarbonContext.getThreadLocalCarbonContext().getUsername();
+        if (changedBy == null){
+            changedBy = DeviceManagementConstants.MaintenanceProperties.MAINTENANCE_USER;
+        }
         PreparedStatement stmt = null;
         try {
             conn = this.getConnection();
@@ -297,7 +368,6 @@ public class EnrollmentDAOImpl implements EnrollmentDAO {
 
                 sql = "INSERT INTO DM_DEVICE_STATUS (ENROLMENT_ID, DEVICE_ID, STATUS, UPDATE_TIME, CHANGED_BY) VALUES(?, ?, ?, ?, ?)";
                 stmt = conn.prepareStatement(sql);
-                // TODO: add changed_by
                 Timestamp updateTime = new Timestamp(new Date().getTime());
                 stmt.setInt(1, enrolmentId);
                 stmt.setInt(2, deviceId);

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

@@ -1836,7 +1836,7 @@ public class DeviceManagementProviderServiceImpl implements DeviceManagementProv
             boolean success;
             int tenantId = this.getTenantId();
             DeviceManagementDAOFactory.beginTransaction();
-            success = enrollmentDAO.setStatus(currentOwner, status, tenantId);
+            success = enrollmentDAO.setStatusAllDevices(currentOwner, status, tenantId);
             DeviceManagementDAOFactory.commitTransaction();
             return success;
         } catch (DeviceManagementDAOException e) {

components/device-mgt/org.wso2.carbon.device.mgt.core/src/test/java/org/wso2/carbon/device/mgt/core/dao/DeviceStatusPersistenceTests.java

@@ -157,6 +157,40 @@ public class DeviceStatusPersistenceTests extends BaseDeviceManagementTest {
 //            DeviceManagementDAOFactory.closeConnection();
         }
     }
+
+    @Test
+    public void testSettingAllDeviceStatusOfSingleUser(){
+        try {
+            this.initDataSource();
+            Device device1 = TestDataHolder.generateDummyDeviceData(TestDataHolder.TEST_DEVICE_TYPE);
+            addDevice(device1);
+
+            Device device2 = TestDataHolder.generateDummyDeviceData(TestDataHolder.TEST_DEVICE_TYPE);
+            addDevice(device2);
+
+            EnrolmentInfo.Status[] statuses1 = {ACTIVE, ASSIGNED, CONFIGURED, READY_TO_CONNECT};
+            int enrolmentId1 = createNewEnrolmentAddStatuses(device1, "admin1", statuses1);
+
+            EnrolmentInfo.Status[] statuses2 = {CREATED, SUSPENDED, BLOCKED, DEFECTIVE, REMOVED, WARRANTY_REPLACED, BLOCKED};
+            int enrolmentId2 = createNewEnrolmentAddStatuses(device1, "admin2", statuses2);
+
+            EnrolmentInfo.Status[] statuses3 = {READY_TO_CONNECT, ASSIGNED};
+            int enrolmentId3 = createNewEnrolmentAddStatuses(device2, "admin1", statuses3);
+
+            enrollmentDAO.setStatusAllDevices("admin1", REMOVED, TestDataHolder.SUPER_TENANT_ID);
+
+            EnrolmentInfo.Status[] statuses1_ = Stream.concat(Arrays.stream(statuses1), Arrays.stream(new EnrolmentInfo.Status[] {REMOVED})).toArray(EnrolmentInfo.Status[]::new);
+            EnrolmentInfo.Status[] statuses3_ = Stream.concat(Arrays.stream(statuses3), Arrays.stream(new EnrolmentInfo.Status[] {REMOVED})).toArray(EnrolmentInfo.Status[]::new);
+
+            validateDeviceStatus(device1, deviceStatusDAO.getStatus(enrolmentId1), statuses1_);
+            validateDeviceStatus(device2, deviceStatusDAO.getStatus(enrolmentId3), statuses3_);
+
+        } catch (DeviceManagementDAOException | SQLException e) {
+            log.error("Error occurred while getting enrolment status", e);
+        } catch (Exception e) {
+            log.error("Error occurred while initializing data source", e);
+        }
+    }
     private int addDevice(Device device) throws DeviceManagementDAOException {
         try {
             DeviceManagementDAOFactory.openConnection();