Add index type checks and allowed file extensions

deepsearch/cps/cli/data_indices_typer.py

@@ -156,14 +156,28 @@ def upload_files(
             typer.echo(ERROR_MSG)
             raise typer.Abort()
 
-    coords = ElasticProjectDataCollectionSource(proj_key=proj_key, index_key=index_key)
-    utils.upload_files(
-        api=api,
-        coords=coords,
-        url=urls,
-        local_file=local_file,
-        s3_coordinates=cos_coordinates,
-    )
+    # get indices of the project
+    indices = api.data_indices.list(proj_key)
+
+    # get specific index to add attachment
+    index = next((x for x in indices if x.source.index_key == index_key), None)
+
+    if index is not None:
+        try:
+            index.upload_files(
+                api=api,
+                url=urls,
+                local_file=local_file,
+                s3_coordinates=cos_coordinates,
+            )
+        except ValueError as e:
+            typer.echo(f"Error occurred: {e}")
+            typer.echo(ERROR_MSG)
+            raise typer.Abort()
+        return
+    else:
+        typer.echo("Index key not found")
+        raise typer.Abort()
 
 
 @app.command(

deepsearch/cps/client/components/data_indices.py

@@ -4,7 +4,7 @@
 import os
 from dataclasses import dataclass
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union
 
 import requests
 from pydantic import BaseModel
@@ -20,6 +20,14 @@
     from deepsearch.cps.client import CpsApi
 
 
+DATA_INDEX_TYPE = Literal[
+    "Document",
+    "DB Record",
+    "Generic",
+    "Experiment",
+]
+
+
 class CpsApiDataIndices:
     def __init__(self, api: CpsApi) -> None:
         self.api = api
@@ -30,7 +38,12 @@ def list(self, proj_key: str) -> List[DataIndex]:
             sw_client.ProjectDataIndexWithStatus
         ] = self.sw_api.get_project_data_indices(proj_key=proj_key)
 
-        return [DataIndex.parse_obj(item.to_dict()) for item in response]
+        # filter out saved searchs index
+        return [
+            DataIndex.parse_obj(item.to_dict())
+            for item in response
+            if item.to_dict()["type"] != "View"
+        ]
 
     def create(
         self,
@@ -102,15 +115,26 @@ def delete(
     def upload_file(
         self,
         coords: ElasticProjectDataCollectionSource,
-        body: Union[Dict[str, List[str]], Dict[str, Dict[str, S3Coordinates]]],
+        index_type: DATA_INDEX_TYPE,
+        body: Union[
+            Dict[str, List[str]], Dict[str, str], Dict[str, Dict[str, S3Coordinates]]
+        ],
     ) -> str:
         """
         Call api for converting and uploading file to a project's data index.
         """
-        task_id = self.sw_api.ccs_convert_upload_file_project_data_index(
-            proj_key=coords.proj_key, index_key=coords.index_key, body=body
-        ).task_id
-        return task_id
+        if index_type == "Document":
+            task_id = self.sw_api.ccs_convert_upload_file_project_data_index(
+                proj_key=coords.proj_key, index_key=coords.index_key, body=body
+            ).task_id
+            return task_id
+        elif index_type == "Generic" or index_type == "DB Record":
+            task_id = self.sw_api.upload_project_data_index_file(
+                proj_key=coords.proj_key, index_key=coords.index_key, params=body
+            ).task_id
+            return task_id
+        else:
+            raise NotImplementedError
 
 
 class ElasticProjectDataCollectionSource(BaseModel):
@@ -130,7 +154,7 @@ class DataIndex(BaseModel):
     health: str
     status: str
     schema_key: str
-    type: str
+    type: DATA_INDEX_TYPE
 
     def add_item_attachment(
         self,
@@ -193,6 +217,46 @@ def add_item_attachment(
             params=params,
         )
 
+    def upload_files(
+        self,
+        api: CpsApi,
+        url: Optional[Union[str, List[str]]] = None,
+        local_file: Optional[Union[str, Path]] = None,
+        s3_coordinates: Optional[S3Coordinates] = None,
+    ) -> None:
+        """
+        Method to upload files to an index.
+        Input
+        -----
+        api : CpsApi
+            CpsApi Class
+        url : string | list[string], OPTIONAL
+            single url string or list of urls string
+        local_file : string | Path, OPTIONAL
+            path to file on local folder
+        s3_coordinates : S3Coordinates, OPTIONAL
+            coordinates of COS to use files on the bucket to convert and upload. Params:
+        """
+
+        # Avoid circular imports
+        from deepsearch.cps.data_indices.utils import upload_files
+
+        if (
+            self.type == "Generic"
+            or self.type == "Document"
+            or self.type == "DB Record"
+        ):
+            upload_files(
+                api=api,
+                coords=self.source,
+                index_type=self.type,
+                url=url,
+                local_file=local_file,
+                s3_coordinates=s3_coordinates,
+            )
+        else:
+            raise NotImplementedError
+
 
 @dataclass
 class CpsApiDataIndex(ApiConnectedObject):