Merge pull request #760 from superannotateai/query_items

Add ability to query items count

Author: VaghinakDev

src/superannotate/__init__.py

@@ -3,7 +3,7 @@
 import sys
 
 
-__version__ = "4.4.31dev2"
+__version__ = "4.4.32dev1"
 
 os.environ.update({"sa_version": __version__})
 sys.path.append(os.path.split(os.path.realpath(__file__))[0])

src/superannotate/lib/core/serviceproviders.py

@@ -672,6 +672,14 @@ def saqul_query(
     ) -> ServiceResponse:
         raise NotImplementedError
 
+    @abstractmethod
+    def query_item_count(
+        self,
+        project: entities.ProjectEntity,
+        query: str = None,
+    ) -> ServiceResponse:
+        raise NotImplementedError
+
 
 class BaseServiceProvider:
     projects: BaseProjectService

src/superannotate/lib/core/usecases/items.py

@@ -163,6 +163,52 @@ def execute(self) -> Response:
         return self._response
 
 
+class QueryEntitiesCountUseCase(BaseReportableUseCase):
+    def __init__(
+        self,
+        reporter: Reporter,
+        project: ProjectEntity,
+        service_provider: BaseServiceProvider,
+        query: str,
+    ):
+        super().__init__(reporter)
+        self._project = project
+        self._service_provider = service_provider
+        self._query = query
+
+    def validate_arguments(self):
+        if self._query:
+            response = self._service_provider.explore.validate_saqul_query(
+                project=self._project, query=self._query
+            )
+
+            if not response.ok:
+                raise AppException(response.error)
+            if response.data["isValidQuery"]:
+                self._query = response.data["parsedQuery"]
+            else:
+                raise AppException("Incorrect query.")
+        else:
+            response = self._service_provider.explore.validate_saqul_query(
+                self._project, "-"
+            )
+            if not response.ok:
+                raise AppException(response.error)
+
+    def execute(self) -> Response:
+        if self.is_valid():
+            query_kwargs = {"query": self._query}
+            service_response = self._service_provider.explore.query_item_count(
+                self._project,
+                **query_kwargs,
+            )
+            if service_response.ok:
+                self._response.data = service_response.data
+            else:
+                self._response.errors = service_response.data
+        return self._response
+
+
 class AssignItemsUseCase(BaseUseCase):
     CHUNK_SIZE = 500
 

src/superannotate/lib/infrastructure/controller.py

@@ -1678,3 +1678,17 @@ def query_entities(
         return ItemManager.process_response(
             self.service_provider, items, project, folder, map_fields=False
         )
+
+    def query_items_count(self, project_name: str, query: str = None) -> int:
+        project = self.get_project(project_name)
+
+        use_case = usecases.QueryEntitiesCountUseCase(
+            reporter=self.get_default_reporter(),
+            project=project,
+            query=query,
+            service_provider=self.service_provider,
+        )
+        response = use_case.execute()
+        if response.errors:
+            raise AppException(response.errors)
+        return response.data["count"]

src/superannotate/lib/infrastructure/services/explore.py

@@ -10,6 +10,7 @@
 from lib.core.service_types import SubsetListResponse
 from lib.core.service_types import UploadCustomFieldValuesResponse
 from lib.core.serviceproviders import BaseExploreService
+from superannotate import AppException
 
 
 class ExploreService(BaseExploreService):
@@ -25,6 +26,7 @@ class ExploreService(BaseExploreService):
     URL_UPLOAD_CUSTOM_VALUE = "custom/metadata/item"
     URL_SAQUL_QUERY = "items/search"
     URL_VALIDATE_SAQUL_QUERY = "items/parse/query"
+    URL_QUERY_COUNT = "items/count"
 
     @property
     def explore_service_url(self):
@@ -201,3 +203,24 @@ def saqul_query(
         else:
             response = ServiceResponse(status=200, res_data=[])
         return response
+
+    def query_item_count(
+        self,
+        project: entities.ProjectEntity,
+        query: str = None,
+    ) -> ServiceResponse:
+
+        params = {
+            "project_id": project.id,
+            "includeFolderNames": True,
+        }
+        data = {"query": query}
+        response = self.client.request(
+            urljoin(self.explore_service_url, self.URL_QUERY_COUNT),
+            "post",
+            params=params,
+            data=data,
+        )
+        if not response.ok:
+            raise AppException(response.error)
+        return response

tests/integration/items/test_saqul_query.py

@@ -59,7 +59,13 @@ def test_query(self):
     def test_query_on_100(self):
         sa.attach_items(self.PROJECT_NAME, os.path.join(DATA_SET_PATH, "100_urls.csv"))
         entities = sa.query(self.PROJECT_NAME, "metadata(status = NotStarted)")
-        print(len(entities))
+        assert len(entities) == 100
+        assert (
+            sa.controller.query_items_count(
+                self.PROJECT_NAME, "metadata(status = NotStarted)"
+            )
+            == 100
+        )
 
     def test_validate_saqul_query(self):
         try: