Fix status connection issue

components/device-mgt/io.entgra.device.mgt.core.device.mgt.core/src/main/java/io/entgra/device/mgt/core/device/mgt/core/service/DeviceManagementProviderServiceImpl.java

@@ -1154,9 +1154,12 @@ public class DeviceManagementProviderServiceImpl implements DeviceManagementProv
 
     public double generateCost(List<Device> allDevices, Timestamp startDate, Timestamp endDate,  Cost tenantCost, List<Device> deviceStatusNotAvailable, double totalCost) throws DeviceManagementException {
         List<DeviceStatus> deviceStatus;
+        try {
+            DeviceManagementDAOFactory.getConnection();
             for (Device device : allDevices) {
                 long dateDiff = 0;
-            deviceStatus = getDeviceStatusHistoryInsideTransaction(device, null, endDate, true);
+                int tenantId = this.getTenantId();
+                deviceStatus = deviceStatusDAO.getStatus(device.getId(), tenantId, null, endDate, true);
                 if (device.getEnrolmentInfo().getDateOfEnrolment() < startDate.getTime()) {
                     if (!deviceStatus.isEmpty() && (String.valueOf(deviceStatus.get(0).getStatus()).equals("REMOVED")
                             || String.valueOf(deviceStatus.get(0).getStatus()).equals("DELETED"))) {
@@ -1199,6 +1202,15 @@ public class DeviceManagementProviderServiceImpl implements DeviceManagementProv
                     deviceStatusNotAvailable.add(device);
                 }
             }
+        } catch (DeviceManagementDAOException e) {
+            String msg = "Error occurred in retrieving status history for a device in billing.";
+            log.error(msg, e);
+            throw new DeviceManagementException(msg, e);
+        } catch (SQLException e) {
+            String msg = "Error occurred while opening a connection to the data source";
+            log.info(msg, e);
+            throw new DeviceManagementException(msg, e);
+        }
         return totalCost;
     }
 
@@ -2211,31 +2223,6 @@ public class DeviceManagementProviderServiceImpl implements DeviceManagementProv
         }
     }
 
-    /*
-    This is just to avoid breaking the billing functionality as it required to call getDeviceStatusHistory method
-    without transaction handling.
-     */
-    private List<DeviceStatus> getDeviceStatusHistoryInsideTransaction(
-            Device device, Date fromDate, Date toDate, boolean billingStatus)
-            throws DeviceManagementException {
-        if (log.isDebugEnabled()) {
-            log.debug("get status history of device: " + device.getDeviceIdentifier());
-        }
-        try {
-            DeviceManagementDAOFactory.getConnection();
-            int tenantId = this.getTenantId();
-            return deviceStatusDAO.getStatus(device.getId(), tenantId, fromDate, toDate, billingStatus);
-        } catch (DeviceManagementDAOException e) {
-            String msg = "Error occurred in retrieving status history for device :" + device.getDeviceIdentifier();
-            log.error(msg, e);
-            throw new DeviceManagementException(msg, e);
-        } catch (SQLException e) {
-            String msg = "Error occurred while opening a connection to the data source";
-            log.info(msg, e);
-            throw new DeviceManagementException(msg, e);
-        }
-    }
-
     @Override
     public List<DeviceStatus> getDeviceStatusHistory(Device device, Date fromDate, Date toDate, boolean billingStatus) throws DeviceManagementException {
         if (log.isDebugEnabled()) {