added Databricks support in attach_items_from_integrated_storage

Narek Mkhitaryan · Narek Mkhitaryan · commit 8f3cf8e87164 · 2025-02-21T15:37:55.000+04:00
diff --git a/src/superannotate/lib/app/interface/sdk_interface.py b/src/superannotate/lib/app/interface/sdk_interface.py
@@ -2919,32 +2919,80 @@ def attach_items_from_integrated_storage(
         project: NotEmptyStr,
         integration: Union[NotEmptyStr, IntegrationEntity],
         folder_path: Optional[NotEmptyStr] = None,
+        *,
+        query: Optional[NotEmptyStr] = None,
+        item_name_column: Optional[NotEmptyStr] = None,
+        custom_item_name: Optional[NotEmptyStr] = None,
+        component_mapping: Optional[Dict[str, str]] = None,
     ):
-        """Link images from integrated external storage to SuperAnnotate.
+        """Link images from integrated external storage to SuperAnnotate from AWS, GCP, Azure, Databricks.
 
         :param project: project name or folder path where items should be attached (e.g., “project1/folder1”).
         :type project: str
 
-        :param integration:  existing integration name or metadata dict to pull items from.
-         Mandatory keys in integration metadata’s dict is “name”.
+        :param integration: The existing integration name or metadata dict to pull items from.
+            Mandatory keys in integration metadata’s dict is “name”.
         :type integration: str or dict
 
         :param folder_path: Points to an exact folder/directory within given storage.
-         If None, items     are fetched from the root directory.
+            If None, items     are fetched from the root directory.
         :type folder_path: str
+
+        :param query: (Only for Databricks). The SQL query to retrieve specific columns from Databricks.
+            If provided, the function will execute the query and use the results for mapping and uploading.
+        :type query: Optional[str]
+
+        :param item_name_column: (Only for Databricks). The column name from the SQL query whose values
+            will be used as item names. If this is provided, custom_item_name cannot be used.
+            The column must exist in the query result.
+        :type item_name_column: Optional[str]
+
+        :param custom_item_name: (Only for Databricks). A manually defined prefix for item names.
+            A random 10-character suffix will be appended to ensure uniqueness.
+            If this is provided, item_name_column cannot be used.
+        :type custom_item_name: Optional[str]
+
+        :param component_mapping: (Only for Databricks). A dictionary mapping Databricks
+            columns to SuperAnnotate component IDs.
+        :type component_mapping: Optional[dict]
+
+
+        Request Example:
+        ::
+
+            client.attach_items_from_integrated_storage(
+                project="project_name",
+                integration="databricks_integration",
+                query="SELECT * FROM integration_data LIMIT 10",
+                item_name_column="prompt",
+                component_mapping={
+                    "category": "_item_category",
+                    "prompt_id": "id",
+                    "prompt": "prompt"
+                }
+            )
+
         """
         project, folder = self.controller.get_project_folder_by_path(project)
         _integration = None
         if isinstance(integration, str):
             integration = IntegrationEntity(name=integration)
         for i in self.controller.integrations.list().data:
-            if integration.name == i.name:
+            if integration.name.lower() == i.name.lower():
                 _integration = i
                 break
         else:
             raise AppException("Integration not found.")
+
         response = self.controller.integrations.attach_items(
-            project, folder, _integration, folder_path
+            project=project,
+            folder=folder,
+            integration=_integration,
+            folder_path=folder_path,
+            query=query,
+            item_name_column=item_name_column,
+            custom_item_name=custom_item_name,
+            component_mapping=component_mapping,
         )
         if response.errors:
             raise AppException(response.errors)
diff --git a/src/superannotate/lib/core/serviceproviders.py b/src/superannotate/lib/core/serviceproviders.py
@@ -593,6 +593,7 @@ def attach_items(
         folder: entities.FolderEntity,
         integration: entities.IntegrationEntity,
         folder_name: str = None,
+        options: Dict[str, str] = None,
     ) -> ServiceResponse:
         raise NotImplementedError
 
diff --git a/src/superannotate/lib/core/usecases/integrations.py b/src/superannotate/lib/core/usecases/integrations.py
@@ -1,8 +1,13 @@
-from typing import List
+from typing import Dict
+from typing import Optional
 
+from lib.core.conditions import Condition
+from lib.core.conditions import CONDITION_EQ as EQ
 from lib.core.entities import FolderEntity
 from lib.core.entities import IntegrationEntity
 from lib.core.entities import ProjectEntity
+from lib.core.entities.integrations import IntegrationTypeEnum
+from lib.core.enums import ProjectType
 from lib.core.exceptions import AppException
 from lib.core.reporter import Reporter
 from lib.core.response import Response
@@ -33,46 +38,133 @@ def __init__(
         service_provider: BaseServiceProvider,
         integration: IntegrationEntity,
         folder_path: str = None,
+        query: Optional[str] = None,
+        item_name_column: Optional[str] = None,
+        custom_item_name: Optional[str] = None,
+        component_mapping: Optional[Dict[str, str]] = None,
     ):
-
         super().__init__(reporter)
         self._project = project
         self._folder = folder
         self._integration = integration
         self._service_provider = service_provider
         self._folder_path = folder_path
+        self._query = query
+        self._item_name_column = item_name_column
+        self._custom_item_name = custom_item_name
+        self._component_mapping = component_mapping
 
     @property
     def _upload_path(self):
         return f"{self._project.name}{f'/{self._folder.name}' if self._folder.name != 'root' else ''}"
 
     def execute(self) -> Response:
-        integrations: List[
-            IntegrationEntity
-        ] = self._service_provider.integrations.list().data.integrations
-        integration_name_lower = self._integration.name.lower()
-        integration = next(
-            (i for i in integrations if i.name.lower() == integration_name_lower), None
-        )
-        if integration:
-            self.reporter.log_info(
-                "Attaching file(s) from "
-                f"{integration.root}{f'/{self._folder_path}' if self._folder_path else ''} "
-                f"to {self._upload_path}. This may take some time."
+        # TODO add support in next iterations
+        if self._integration.type == IntegrationTypeEnum.SNOWFLAKE:
+            raise AppException(
+                "Attaching items is not supported with Snowflake integration."
             )
-            attached = self._service_provider.integrations.attach_items(
-                project=self._project,
-                folder=self._folder,
-                integration=integration,
-                folder_name=self._folder_path,
+
+        options = {}  # using only for Databricks and Snowflake
+        multimodal_integrations = [
+            IntegrationTypeEnum.DATABRICKS,
+            IntegrationTypeEnum.SNOWFLAKE,
+        ]
+        if self._integration.type in multimodal_integrations:
+            if self._project.type != ProjectType.MULTIMODAL:
+                raise AppException(
+                    f"{self._integration.name} integration is supported only for Multimodal projects."
+                )
+
+            if self._item_name_column and self._custom_item_name:
+                raise AppException(
+                    "‘item_name_column and custom_item_name cannot be used simultaneously."
+                )
+
+            if not self._item_name_column and not self._custom_item_name:
+                raise AppException(
+                    "Either item_name_column or custom_item_name is required."
+                )
+
+            if not all((self._query, self._component_mapping)):
+                raise AppException(
+                    f"{self._integration.name} integration requires both a query and component_mapping."
+                )
+
+            category_setting: bool = bool(
+                next(
+                    (
+                        setting.value
+                        for setting in self._service_provider.projects.list_settings(
+                            self._project
+                        ).data
+                        if setting.attribute == "CategorizeItems"
+                    ),
+                    None,
+                )
             )
-            if not attached:
-                self._response.errors = AppException(
-                    f"An error occurred for {self._integration.name}. Please make sure: "
-                    "\n - The bucket exists."
-                    "\n - The connection is valid."
-                    "\n - The path to a specified directory is correct."
+            if (
+                not category_setting
+                and "_item_category" in self._component_mapping.values()
+            ):
+                raise AppException(
+                    "Item Category must be enabled for a project to use _item_category"
                 )
-        else:
-            self._response.errors = AppException("Integration not found.")
+
+            item_category_column = next(
+                (
+                    k
+                    for k, v in self._component_mapping.items()
+                    if v == "_item_category"
+                ),
+                None,
+            )
+            if item_category_column:
+                self._component_mapping.pop(item_category_column)
+
+            sa_components = [
+                c.name.lower()
+                for c in self._service_provider.annotation_classes.list(
+                    condition=Condition("project_id", self._project.id, EQ)
+                ).data
+            ]
+
+            for i in self._component_mapping.values():
+                if i.lower() not in sa_components:
+                    raise AppException(
+                        f"Component mapping contains invalid component ID: `{i}`"
+                    )
+
+            options["query"] = self._query
+            options["item_name"] = (
+                self._custom_item_name
+                if self._custom_item_name
+                else self._item_name_column
+            )
+            options["prefix"] = True if self._custom_item_name else False
+            options["column_class_map"] = self._component_mapping
+            if item_category_column:
+                options["item_category"] = item_category_column
+
+        self.reporter.log_info(
+            "Attaching file(s) from "
+            f"{self._integration.root}{f'/{self._folder_path}' if self._folder_path else ''} "
+            f"to {self._upload_path}. This may take some time."
+        )
+        attache_response = self._service_provider.integrations.attach_items(
+            project=self._project,
+            folder=self._folder,
+            integration=self._integration,
+            folder_name=self._folder_path
+            if self._integration.type not in multimodal_integrations
+            else None,
+            options=options if options else None,
+        )
+        if not attache_response.ok:
+            self._response.errors = AppException(
+                f"An error occurred for {self._integration.name}. Please make sure: "
+                "\n - The bucket exists."
+                "\n - The connection is valid."
+                "\n - The path to a specified directory is correct."
+            )
         return self._response
diff --git a/src/superannotate/lib/infrastructure/controller.py b/src/superannotate/lib/infrastructure/controller.py
@@ -1138,7 +1138,11 @@ def attach_items(
         project: ProjectEntity,
         folder: FolderEntity,
         integration: IntegrationEntity,
-        folder_path: str,
+        folder_path: str = None,
+        query: Optional[str] = None,
+        item_name_column: Optional[str] = None,
+        custom_item_name: Optional[str] = None,
+        component_mapping: Optional[Dict[str, str]] = None,
     ):
         use_case = usecases.AttachIntegrations(
             reporter=Reporter(),
@@ -1147,6 +1151,10 @@ def attach_items(
             folder=folder,
             integration=integration,
             folder_path=folder_path,
+            query=query,
+            item_name_column=item_name_column,
+            custom_item_name=custom_item_name,
+            component_mapping=component_mapping,
         )
         return use_case.execute()
 
diff --git a/src/superannotate/lib/infrastructure/services/integration.py b/src/superannotate/lib/infrastructure/services/integration.py
@@ -1,3 +1,5 @@
+from typing import Dict
+
 from lib.core import entities
 from lib.core.service_types import IntegrationListResponse
 from lib.core.serviceproviders import BaseIntegrationService
@@ -23,6 +25,7 @@ def attach_items(
         folder: entities.FolderEntity,
         integration: entities.IntegrationEntity,
         folder_name: str = None,
+        options: Dict[str, str] = None,
     ):
         data = {
             "team_id": project.team_id,
@@ -32,6 +35,8 @@ def attach_items(
         }
         if folder_name:
             data["customer_folder_name"] = folder_name
+        if options:
+            data["options"] = options
         return self.client.request(
             self.URL_ATTACH_INTEGRATIONS.format(project.team_id), "post", data=data
         )
diff --git a/tests/test_.py b/tests/test_.py
@@ -0,0 +1,17 @@
+from superannotate import SAClient
+
+sa = SAClient()
+
+def test_sa():
+    sa.attach_items_from_integrated_storage(
+        "databricks",
+        "Nshantest",
+        query="SELECT * FROM sa_db_integration.sadataset.sa_integration_data_to_sa LIMIT 10",
+        # item_name_column="category",
+        custom_item_name="test_SDK",
+        component_mapping={
+            "prompt_id": "_item_category",
+            "category": "category",
+            "prompt": "prompt"
+        }
+    )
