Use sdk models in data_indices methods

deepsearch/cps/client/components/data_indices.py

@@ -3,10 +3,11 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
+from pydantic import BaseModel
+
 from deepsearch.cps.apis import public as sw_client
 from deepsearch.cps.apis.public.models.token_response import TokenResponse
 from deepsearch.cps.client.components.api_object import ApiConnectedObject
-from deepsearch.cps.client.components.elastic import ElasticProjectDataCollectionSource
 
 if TYPE_CHECKING:
     from deepsearch.cps.client import CpsApi
@@ -17,25 +18,21 @@ def __init__(self, api: CpsApi) -> None:
         self.api = api
         self.sw_api = sw_client.DataIndicesApi(self.api.client.swagger_client)
 
-    # define methods:
-    def list(self, proj_key: str):
-        try:
-            projects = self.sw_api.get_project_data_indices(proj_key=proj_key)
-            return projects
-        except ValueError as e:
-            print(f"Uh oh! {e}\nAborting!")
-            raise
-
-    # def create(self, proj_key: str, data: Dict[str, str]):
-    #     return self.sw_api.create_project_data_index(proj_key=proj_key, data=data)
+    def list(self, proj_key: str) -> List[DataIndex]:
+        response: list[
+            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]
+
     def create(
         self,
         proj_key: str,
         name: str,
         desc: str = "",
         type: Optional[str] = None,
         schema_key: Optional[str] = None,
-    ):
+    ) -> DataIndex:
         """
         Method to create a new index.
 
@@ -73,12 +70,16 @@ def create(
             "schema_key": schema_key,
             "type": type,
         }
-        return self.sw_api.create_project_data_index(proj_key=proj_key, data=data)
+        response: sw_client.ProjectDataIndexWithStatus = (
+            self.sw_api.create_project_data_index(proj_key=proj_key, data=data)
+        )
+
+        return DataIndex.parse_obj(response.to_dict())
 
     def delete(
         self,
         coords: ElasticProjectDataCollectionSource,
-    ):
+    ) -> None:
         request_confirmation_token: TokenResponse = (
             self.sw_api.create_project_data_index_delete_token(
                 proj_key=coords.proj_key, index_key=coords.index_key
@@ -105,6 +106,26 @@ def upload_file(
         return task_id
 
 
+class ElasticProjectDataCollectionSource(BaseModel):
+    proj_key: str
+    index_key: str
+
+    def to_resource(self) -> Dict[str, Any]:
+        return {"type": "elastic", "proj_key": self.proj_key, "index": self.index_key}
+
+
+class DataIndex(BaseModel):
+
+    source: ElasticProjectDataCollectionSource
+    name: str
+    description: str
+    documents: int
+    health: str
+    status: str
+    schema_key: str
+    type: str
+
+
 @dataclass
 class CpsApiDataIndex(ApiConnectedObject):
     project: str

deepsearch/cps/client/components/elastic.py

@@ -21,6 +21,9 @@
 from deepsearch.cps.apis.public.models.elastic_index_search_results import (
     ElasticIndexSearchResults,
 )
+from deepsearch.cps.client.components.data_indices import (
+    ElasticProjectDataCollectionSource,
+)
 
 if TYPE_CHECKING:
     from deepsearch.cps.client import CpsApi
@@ -165,14 +168,6 @@ def to_resource(self) -> Dict[str, Any]:
         }
 
 
-class ElasticProjectDataCollectionSource(BaseModel):
-    proj_key: str
-    index_key: str
-
-    def to_resource(self) -> Dict[str, Any]:
-        return {"type": "elastic", "proj_key": self.proj_key, "index": self.index_key}
-
-
 class ElasticDataCollectionMetadata(BaseModel):
     description: str
     created: Union[datetime.datetime, str]