added subset to query function

Author: VaghinakDev

src/superannotate/lib/app/interface/sdk_interface.py

@@ -2423,7 +2423,12 @@ def attach_items_from_integrated_storage(
         if response.errors:
             raise AppException(response.errors)
 
-    def query(self, project: NotEmptyStr, query: Optional[NotEmptyStr]):
+    def query(
+        self,
+        project: NotEmptyStr,
+        query: Optional[NotEmptyStr] = None,
+        subset: Optional[NotEmptyStr] = None,
+    ):
         """Return items that satisfy the given query.
         Query syntax should be in SuperAnnotate query language(https://doc.superannotate.com/docs/query-search-1).
 
@@ -2433,11 +2438,17 @@ def query(self, project: NotEmptyStr, query: Optional[NotEmptyStr]):
         :param query: SAQuL query string.
         :type query: str
 
+        :param subset:  subset name. Allows you to query items in a specific subset.
+        To return all the items in the specified subset, set the value of query param to None.
+        :type subset: str
+
         :return: queried items’ metadata list
         :rtype: list of dicts
         """
         project_name, folder_name = extract_project_folder(project)
-        response = self.controller.query_entities(project_name, folder_name, query)
+        response = self.controller.query_entities(
+            project_name, folder_name, query, subset
+        )
         if response.errors:
             raise AppException(response.errors)
         return BaseSerializer.serialize_iterable(response.data)
@@ -2759,4 +2770,4 @@ def get_subsets(self, project: Union[NotEmptyStr, dict]):
         response = self.controller.list_subsets(project_name)
         if response.errors:
             raise AppException(response.errors)
-        return BaseSerializer.serialize_iterable(response.data)
+        return BaseSerializer.serialize_iterable(response.data, ["name"])

src/superannotate/lib/core/entities/__init__.py

@@ -2,6 +2,7 @@
 from lib.core.entities.base import BaseEntity as TmpBaseEntity
 from lib.core.entities.base import ProjectEntity
 from lib.core.entities.base import SettingEntity
+from lib.core.entities.base import SubSetEntity
 from lib.core.entities.integrations import IntegrationEntity
 from lib.core.entities.items import DocumentEntity
 from lib.core.entities.items import Entity
@@ -31,6 +32,7 @@
 __all__ = [
     # base
     "SettingEntity",
+    "SubSetEntity",
     # items
     "TmpImageEntity",
     "BaseEntity",

src/superannotate/lib/core/entities/base.py

@@ -28,6 +28,7 @@ def validate(cls, v: datetime):
 
 
 class SubSetEntity(BaseModel):
+    id: Optional[int]
     name: str
 
     class Config:

src/superannotate/lib/core/serviceproviders.py

@@ -362,7 +362,12 @@ def attach_integrations(
         raise NotImplementedError
 
     def saqul_query(
-        self, team_id: int, project_id: int, query: str, folder_id: int
+        self,
+        team_id: int,
+        project_id: int,
+        folder_id: int,
+        query: str = None,
+        subset_id: int = None,
     ) -> ServiceResponse:
         raise NotImplementedError
 

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

@@ -1,5 +1,6 @@
 import copy
 from typing import List
+from typing import Optional
 
 import superannotate.lib.core as constances
 from lib.core.conditions import Condition
@@ -9,6 +10,7 @@
 from lib.core.entities import Entity
 from lib.core.entities import FolderEntity
 from lib.core.entities import ProjectEntity
+from lib.core.entities import SubSetEntity
 from lib.core.entities import TmpImageEntity
 from lib.core.entities import VideoEntity
 from lib.core.exceptions import AppException
@@ -78,41 +80,78 @@ def execute(self) -> Response:
         return self._response
 
 
-class QueryEntities(BaseReportableUseCase):
+class QueryEntitiesUseCase(BaseReportableUseCase):
     def __init__(
         self,
         reporter: Reporter,
         project: ProjectEntity,
         folder: FolderEntity,
         backend_service_provider: SuperannotateServiceProvider,
         query: str,
+        subset: str,
     ):
         super().__init__(reporter)
         self._project = project
         self._folder = folder
         self._backend_client = backend_service_provider
         self._query = query
+        self._subset = subset
 
     def validate_query(self):
-        response = self._backend_client.validate_saqul_query(
-            self._project.team_id, self._project.id, self._query
-        )
-        if response.get("error"):
-            raise AppException(response["error"])
-        if response["isValidQuery"]:
-            self._query = response["parsedQuery"]
-        else:
-            raise AppException("Incorrect query.")
         if self._project.sync_status != constances.ProjectState.SYNCED.value:
             raise AppException("Data is not synced.")
+        if self._query:
+            response = self._backend_client.validate_saqul_query(
+                self._project.team_id, self._project.id, self._query
+            )
+            if response.get("error"):
+                raise AppException(response["error"])
+            if response["isValidQuery"]:
+                self._query = response["parsedQuery"]
+            else:
+                raise AppException("Incorrect query.")
+        else:
+            response = self._backend_client.validate_saqul_query(
+                self._project.team_id, self._project.id, "-"
+            )
+            if response.get("error"):
+                raise AppException(response["error"])
 
     def execute(self) -> Response:
+        if not any([self._query, self._subset]):
+            self._response.errors = AppException(
+                "AppException: The query and subset params cannot have the value None at the same time."
+            )
+            return self._response
+
         if self.is_valid():
+            query_kwargs = {}
+            if self._subset:
+                subset: Optional[SubSetEntity] = None
+                response = self._backend_client.list_sub_sets(
+                    team_id=self._project.team_id, project_id=self._project.id
+                )
+                if response.ok:
+                    subset = next(
+                        (_sub for _sub in response.data if _sub.name == self._subset),
+                        None,
+                    )
+                if not subset:
+                    self._response.errors = AppException(
+                        "Subset not found. Use the superannotate."
+                        "get_subsets() function to get a list of the available subsets."
+                    )
+                    return self._response
+                query_kwargs["subset_id"] = subset.id
+            if self._query:
+                query_kwargs["query"] = self._query
+            query_kwargs["folder_id"] = (
+                None if self._folder.name == "root" else self._folder.uuid
+            )
             service_response = self._backend_client.saqul_query(
                 self._project.team_id,
                 self._project.id,
-                self._query,
-                folder_id=None if self._folder.name == "root" else self._folder.uuid,
+                **query_kwargs,
             )
             if service_response.ok:
                 data = []

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

@@ -1110,7 +1110,7 @@ def execute(self) -> Response:
                 team_id=self._project.team_id, project_id=self._project.id
             )
             if sub_sets_response.ok:
-                self._response.data = []
+                self._response.data = sub_sets_response.data
             else:
                 self._response.data = []
 

src/superannotate/lib/infrastructure/controller.py

@@ -1509,15 +1509,18 @@ def attach_integrations(
         )
         return use_case.execute()
 
-    def query_entities(self, project_name: str, folder_name: str, query: str = None):
+    def query_entities(
+        self, project_name: str, folder_name: str, query: str = None, subset: str = None
+    ):
         project = self._get_project(project_name)
         folder = self._get_folder(project, folder_name)
 
-        use_case = usecases.QueryEntities(
+        use_case = usecases.QueryEntitiesUseCase(
             reporter=self.get_default_reporter(),
             project=project,
             folder=folder,
             query=query,
+            subset=subset,
             backend_service_provider=self.backend_client,
         )
         return use_case.execute()

src/superannotate/lib/infrastructure/services.py

@@ -189,6 +189,7 @@ class SuperannotateBackendService(BaseBackendService):
     """
 
     DEFAULT_CHUNK_SIZE = 5000
+    SAQUL_CHUNK_SIZE = 50
 
     URL_USERS = "users"
     URL_LIST_PROJECTS = "projects"
@@ -1155,9 +1156,14 @@ def attach_integrations(
         return response.ok
 
     def saqul_query(
-        self, team_id: int, project_id: int, query: str, folder_id: int
+        self,
+        team_id: int,
+        project_id: int,
+        folder_id: int,
+        query: str = None,
+        subset_id: int = None,
     ) -> ServiceResponse:
-        CHUNK_SIZE = 50
+
         query_url = urljoin(self.api_url, self.URL_SAQUL_QUERY)
         params = {
             "team_id": team_id,
@@ -1166,18 +1172,22 @@ def saqul_query(
         }
         if folder_id:
             params["folder_id"] = folder_id
-        data = {"query": query, "image_index": 0}
+        if subset_id:
+            params["subset_id"] = subset_id
+        data = {"image_index": 0}
+        if query:
+            data["query"] = query
         items = []
         for _ in range(self.MAX_ITEMS_COUNT):
             response = self._request(query_url, "post", params=params, data=data)
             if response.ok:
                 response_items = response.json()
                 items.extend(response_items)
-                if len(response_items) < CHUNK_SIZE:
+                if len(response_items) < self.SAQUL_CHUNK_SIZE:
                     service_response = ServiceResponse(response)
                     service_response.data = items
                     return service_response
-                data["image_index"] += CHUNK_SIZE
+                data["image_index"] += self.SAQUL_CHUNK_SIZE
             return ServiceResponse(response)
 
     def validate_saqul_query(self, team_id: int, project_id: int, query: str) -> dict:

tests/integration/items/test_saqul_query.py

@@ -2,9 +2,10 @@
 from pathlib import Path
 
 from src.superannotate import SAClient
-sa = SAClient()
 from tests.integration.base import BaseTestCase
 
+sa = SAClient()
+
 
 class TestEntitiesSearchVector(BaseTestCase):
     PROJECT_NAME = "TestEntitiesSearchVector"
@@ -35,6 +36,16 @@ def test_query(self):
         self.assertEqual(len(entities), 1)
         assert all([entity["path"] == f"{self.PROJECT_NAME}/{self.FOLDER_NAME}" for entity in entities])
 
+        try:
+            self.assertRaises(
+                Exception, sa.query(f"{self.PROJECT_NAME}/{self.FOLDER_NAME}", self.TEST_QUERY, subset="something")
+            )
+        except Exception as e:
+            self.assertEqual(
+                str(e),
+                "Subset not found. Use the superannotate.get_subsets() function to get a list of the available subsets."
+            )
+
     def test_validate_saqul_query(self):
         try:
             self.assertRaises(Exception, sa.query(self.PROJECT_NAME, self.TEST_INVALID_QUERY))