Applying charithas commit on IOTS-436

components/device-mgt/org.wso2.carbon.device.mgt.api/pom.xml

@@ -291,6 +291,13 @@
 			<artifactId>org.wso2.carbon.identity.jwt.client.extension</artifactId>
 			<scope>provided</scope>
 		</dependency>
+        <dependency>
+            <groupId>org.wso2.carbon.identity.framework</groupId>
+            <artifactId>org.wso2.carbon.identity.user.store.count</artifactId>
+            <version>${carbon.identity.framework.version}</version>
+            <scope>provided</scope>
+        </dependency>
+
     </dependencies>
 
 </project>

components/device-mgt/org.wso2.carbon.device.mgt.api/src/main/java/org/wso2/carbon/device/mgt/jaxrs/beans/BasePaginatedResult.java

@@ -30,13 +30,13 @@ public class BasePaginatedResult {
             value = "Number of total resources.",
             example = "1")
     @JsonProperty("count")
-    private int count;
+    private long count;
 
-    public int getCount() {
+    public long getCount() {
         return count;
     }
 
-    public void setCount(int count) {
+    public void setCount(long count) {
         this.count = count;
     }
 }

components/device-mgt/org.wso2.carbon.device.mgt.api/src/main/java/org/wso2/carbon/device/mgt/jaxrs/service/impl/UserManagementServiceImpl.java

@@ -38,6 +38,8 @@ import org.wso2.carbon.device.mgt.jaxrs.service.impl.util.RequestValidationUtil;
 import org.wso2.carbon.device.mgt.jaxrs.util.Constants;
 import org.wso2.carbon.device.mgt.jaxrs.util.CredentialManagementResponseBuilder;
 import org.wso2.carbon.device.mgt.jaxrs.util.DeviceMgtAPIUtils;
+import org.wso2.carbon.identity.user.store.count.UserStoreCountRetriever;
+import org.wso2.carbon.identity.user.store.count.exception.UserStoreCounterException;
 import org.wso2.carbon.user.api.UserStoreException;
 import org.wso2.carbon.user.api.UserStoreManager;
 import org.wso2.carbon.utils.CarbonUtils;
@@ -395,6 +397,30 @@ public class UserManagementServiceImpl implements UserManagementService {
     @Path("/count")
     @Override
     public Response getUserCount() {
+        try {
+            UserStoreCountRetriever userStoreCountRetrieverService = DeviceMgtAPIUtils.getUserStoreCountRetrieverService();
+            if (userStoreCountRetrieverService != null) {
+                long count = userStoreCountRetrieverService.countUsers("");
+                if (count != -1) {
+                    BasicUserInfoList result = new BasicUserInfoList();
+                    result.setCount(count);
+                    return Response.status(Response.Status.OK).entity(result).build();
+                }
+            }
+        } catch (UserStoreCounterException e) {
+            String msg =
+                    "Error occurred while retrieving the count of users that exist within the current tenant";
+            log.error(msg, e);
+        }
+        return getUserCountViaUserStoreManager();
+    }
+
+    /**
+     * This method returns the count of users using UserStoreManager.
+     *
+     * @return user count
+     */
+    private Response getUserCountViaUserStoreManager() {
         if (log.isDebugEnabled()) {
             log.debug("Getting the user count");
         }

components/device-mgt/org.wso2.carbon.device.mgt.api/src/main/java/org/wso2/carbon/device/mgt/jaxrs/util/DeviceMgtAPIUtils.java

@@ -37,14 +37,21 @@ import org.wso2.carbon.device.mgt.core.service.GroupManagementProviderService;
 import org.wso2.carbon.device.mgt.jaxrs.beans.ErrorResponse;
 import org.wso2.carbon.device.mgt.jaxrs.service.impl.util.InputValidationException;
 import org.wso2.carbon.identity.jwt.client.extension.service.JWTClientManagerService;
+import org.wso2.carbon.identity.user.store.count.AbstractCountRetrieverFactory;
+import org.wso2.carbon.identity.user.store.count.UserStoreCountRetriever;
+import org.wso2.carbon.identity.user.store.count.exception.UserStoreCounterException;
+import org.wso2.carbon.identity.user.store.count.jdbc.JDBCCountRetrieverFactory;
+import org.wso2.carbon.identity.user.store.count.jdbc.internal.InternalCountRetrieverFactory;
 import org.wso2.carbon.policy.mgt.common.PolicyMonitoringTaskException;
 import org.wso2.carbon.policy.mgt.core.PolicyManagerService;
 import org.wso2.carbon.policy.mgt.core.task.TaskScheduleService;
 import org.wso2.carbon.registry.core.service.RegistryService;
 import org.wso2.carbon.user.api.AuthorizationManager;
+import org.wso2.carbon.user.api.RealmConfiguration;
 import org.wso2.carbon.user.api.UserRealm;
 import org.wso2.carbon.user.api.UserStoreException;
 import org.wso2.carbon.user.api.UserStoreManager;
+import org.wso2.carbon.user.core.jdbc.JDBCUserStoreManager;
 import org.wso2.carbon.user.core.service.RealmService;
 
 import javax.ws.rs.core.MediaType;
@@ -106,6 +113,32 @@ public class DeviceMgtAPIUtils {
         return deviceManagementProviderService;
     }
 
+    public static UserStoreCountRetriever getUserStoreCountRetrieverService()
+            throws UserStoreCounterException {
+        PrivilegedCarbonContext ctx = PrivilegedCarbonContext.getThreadLocalCarbonContext();
+        List<Object> countRetrieverFactories = ctx.getOSGiServices(AbstractCountRetrieverFactory.class, null);
+        RealmService realmService = (RealmService) ctx.getOSGiService(RealmService.class, null);
+        RealmConfiguration realmConfiguration = realmService.getBootstrapRealmConfiguration();
+        String userStoreType;
+        //Ignoring Sonar warning as getUserStoreClass() returning string name of the class. So cannot use 'instanceof'.
+        if (JDBCUserStoreManager.class.getName().equals(realmConfiguration.getUserStoreClass())) {
+            userStoreType = JDBCCountRetrieverFactory.JDBC;
+        } else {
+            userStoreType = InternalCountRetrieverFactory.INTERNAL;
+        }
+        AbstractCountRetrieverFactory countRetrieverFactory = null;
+        for (Object countRetrieverFactoryObj : countRetrieverFactories) {
+            countRetrieverFactory = (AbstractCountRetrieverFactory) countRetrieverFactoryObj;
+            if (userStoreType.equals(countRetrieverFactory.getCounterType())) {
+                break;
+            }
+        }
+        if (countRetrieverFactory == null) {
+            return null;
+        }
+        return countRetrieverFactory.buildCountRetriever(realmConfiguration);
+    }
+
     public static DeviceAccessAuthorizationService getDeviceAccessAuthorizationService() {
         PrivilegedCarbonContext ctx = PrivilegedCarbonContext.getThreadLocalCarbonContext();
         DeviceAccessAuthorizationService deviceAccessAuthorizationService =