From 31a6497950e08e0d76a8fc46ce35daf817c26e53 Mon Sep 17 00:00:00 2001
From: charitha <charitha@entgra.io>
Date: Sat, 28 Nov 2020 11:31:09 +0530
Subject: [PATCH] Fix enrollment modification issues

---
 .../device/type/template/DeviceTypeManager.java | 17 ++++++++++-------
 .../template/dao/DeviceTypePluginDAOImpl.java   |  6 ++----
 2 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/components/device-mgt/org.wso2.carbon.device.mgt.extensions/src/main/java/org/wso2/carbon/device/mgt/extensions/device/type/template/DeviceTypeManager.java b/components/device-mgt/org.wso2.carbon.device.mgt.extensions/src/main/java/org/wso2/carbon/device/mgt/extensions/device/type/template/DeviceTypeManager.java
index f5694d53c2a..3805fd09d35 100644
--- a/components/device-mgt/org.wso2.carbon.device.mgt.extensions/src/main/java/org/wso2/carbon/device/mgt/extensions/device/type/template/DeviceTypeManager.java
+++ b/components/device-mgt/org.wso2.carbon.device.mgt.extensions/src/main/java/org/wso2/carbon/device/mgt/extensions/device/type/template/DeviceTypeManager.java
@@ -307,10 +307,7 @@ public class DeviceTypeManager implements DeviceManager {
                 JAXBContext context = JAXBContext.newInstance(PlatformConfiguration.class);
                 Unmarshaller unmarshaller = context.createUnmarshaller();
                 return (PlatformConfiguration) unmarshaller.unmarshal(reader);
-            } else if (defaultPlatformConfiguration != null) {
-                return defaultPlatformConfiguration;
-            }
-            return null;
+            } else return defaultPlatformConfiguration;
         } catch (DeviceTypeMgtPluginException e) {
             throw new DeviceManagementException(
                     "Error occurred while retrieving the Registry instance : " + e.getMessage(), e);
@@ -379,10 +376,16 @@ public class DeviceTypeManager implements DeviceManager {
             boolean status;
             try {
                 if (log.isDebugEnabled()) {
-                    log.debug("Adding properties for new device : " + device.getDeviceIdentifier());
+                    log.debug("Modifying properties for enrolling device : " + device.getDeviceIdentifier());
                 }
                 deviceTypePluginDAOManager.getDeviceTypeDAOHandler().beginTransaction();
-                status = deviceTypePluginDAOManager.getDeviceDAO().updateDevice(device);
+                Device existingDevice = deviceTypePluginDAOManager.getDeviceDAO()
+                        .getDevice(device.getDeviceIdentifier());
+                if (existingDevice == null) {
+                    status = deviceTypePluginDAOManager.getDeviceDAO().addDevice(device);
+                } else {
+                    status = deviceTypePluginDAOManager.getDeviceDAO().updateDevice(device);
+                }
                 deviceTypePluginDAOManager.getDeviceTypeDAOHandler().commitTransaction();
             } catch (DeviceTypeMgtPluginException e) {
                 deviceTypePluginDAOManager.getDeviceTypeDAOHandler().rollbackTransaction();
@@ -395,7 +398,7 @@ public class DeviceTypeManager implements DeviceManager {
             if (propertyBasedDeviceTypePluginDAOManager != null && status) {
                 try {
                     if (log.isDebugEnabled()) {
-                        log.debug("Updating properties for new device : " + device.getDeviceIdentifier());
+                        log.debug("Updating properties for enrolling device : " + device.getDeviceIdentifier());
                     }
                     propertyBasedDeviceTypePluginDAOManager.getDeviceTypeDAOHandler().beginTransaction();
                     status = propertyBasedDeviceTypePluginDAOManager.getDeviceDAO().updateDevice(device);
diff --git a/components/device-mgt/org.wso2.carbon.device.mgt.extensions/src/main/java/org/wso2/carbon/device/mgt/extensions/device/type/template/dao/DeviceTypePluginDAOImpl.java b/components/device-mgt/org.wso2.carbon.device.mgt.extensions/src/main/java/org/wso2/carbon/device/mgt/extensions/device/type/template/dao/DeviceTypePluginDAOImpl.java
index c7378f16312..89783b21988 100644
--- a/components/device-mgt/org.wso2.carbon.device.mgt.extensions/src/main/java/org/wso2/carbon/device/mgt/extensions/device/type/template/dao/DeviceTypePluginDAOImpl.java
+++ b/components/device-mgt/org.wso2.carbon.device.mgt.extensions/src/main/java/org/wso2/carbon/device/mgt/extensions/device/type/template/dao/DeviceTypePluginDAOImpl.java
@@ -39,7 +39,6 @@ import org.apache.commons.lang.StringUtils;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.wso2.carbon.device.mgt.common.Device;
-import org.wso2.carbon.device.mgt.common.DeviceIdentifier;
 import org.wso2.carbon.device.mgt.extensions.device.type.template.exception.DeviceTypeMgtPluginException;
 import org.wso2.carbon.device.mgt.extensions.device.type.template.util.DeviceTypeUtils;
 
@@ -58,8 +57,8 @@ import java.util.List;
 public class DeviceTypePluginDAOImpl implements PluginDAO {
 
     private static final Log log = LogFactory.getLog(DeviceTypePluginDAOImpl.class);
-    private DeviceTypeDAOHandler deviceTypeDAOHandler;
-    private DeviceDAODefinition deviceDAODefinition;
+    private final DeviceTypeDAOHandler deviceTypeDAOHandler;
+    private final DeviceDAODefinition deviceDAODefinition;
     private String selectDBQueryForGetDevice;
     private String createDBqueryForAddDevice;
     private String updateDBQueryForUpdateDevice;
@@ -108,7 +107,6 @@ public class DeviceTypePluginDAOImpl implements PluginDAO {
             throw new DeviceTypeMgtPluginException(msg, e);
         } finally {
             DeviceTypeUtils.cleanupResources(stmt, resultSet);
-            deviceTypeDAOHandler.closeConnection();
         }
 
         return device;