Support more datasources APIs (#55)

Add support for:

- sync_datasource
- list_datasources_objects_states

compass_sdk/__init__.py

@@ -12,7 +12,7 @@
     ValidatedModel,
 )
 
-__version__ = "0.6.0"
+__version__ = "0.7.0"
 
 
 class ProcessFileParameters(ValidatedModel):

compass_sdk/clients/compass.py

@@ -46,6 +46,7 @@
     CreateDataSource,
     DataSource,
     Document,
+    DocumentStatus,
     PaginatedList,
     ParseableDocument,
     PushDocumentsInput,
@@ -119,6 +120,8 @@ def __init__(
             "list_datasources": self.session.get,
             "delete_datasources": self.session.delete,
             "get_datasource": self.session.get,
+            "sync_datasource": self.session.post,
+            "list_datasources_objects_states": self.session.get,
         }
         self.api_endpoint = {
             "create_index": "/api/v1/indexes/{index_name}",
@@ -138,6 +141,8 @@ def __init__(
             "list_datasources": "/api/v1/datasources",
             "delete_datasources": "/api/v1/datasources/{datasource_id}",
             "get_datasource": "/api/v1/datasources/{datasource_id}",
+            "sync_datasource": "/api/v1/datasources/{datasource_id}/_sync",
+            "list_datasources_objects_states": "/api/v1/datasources/{datasource_id}/documents?skip={skip}&limit={limit}",
         }
 
     def create_index(self, *, index_name: str):
@@ -460,7 +465,7 @@ def list_datasources(
         *,
         max_retries: int = DEFAULT_MAX_RETRIES,
         sleep_retry_seconds: int = DEFAULT_SLEEP_RETRY_SECONDS,
-    ):
+    ) -> Union[PaginatedList[DataSource], str]:
         result = self._send_request(
             api_name="list_datasources",
             max_retries=max_retries,
@@ -469,7 +474,7 @@ def list_datasources(
 
         if result.error:
             return result.error
-        return PaginatedList.model_validate(result.result)
+        return PaginatedList[DataSource].model_validate(result.result)
 
     def get_datasource(
         self,
@@ -507,6 +512,46 @@ def delete_datasource(
             return result.error
         return result.result
 
+    def sync_datasource(
+        self,
+        *,
+        datasource_id: str,
+        max_retries: int = DEFAULT_MAX_RETRIES,
+        sleep_retry_seconds: int = DEFAULT_SLEEP_RETRY_SECONDS,
+    ):
+        result = self._send_request(
+            api_name="sync_datasource",
+            datasource_id=datasource_id,
+            max_retries=max_retries,
+            sleep_retry_seconds=sleep_retry_seconds,
+        )
+
+        if result.error:
+            return result.error
+        return result.result
+
+    def list_datasources_objects_states(
+        self,
+        *,
+        datasource_id: str,
+        skip: int = 0,
+        limit: int = 100,
+        max_retries: int = DEFAULT_MAX_RETRIES,
+        sleep_retry_seconds: int = DEFAULT_SLEEP_RETRY_SECONDS,
+    ) -> Union[PaginatedList[DocumentStatus], str]:
+        result = self._send_request(
+            api_name="list_datasources_objects_states",
+            datasource_id=datasource_id,
+            max_retries=max_retries,
+            sleep_retry_seconds=sleep_retry_seconds,
+            skip=str(skip),
+            limit=str(limit),
+        )
+
+        if result.error:
+            return result.error
+        return PaginatedList[DocumentStatus].model_validate(result.result)
+
     @staticmethod
     def _get_request_blocks(
         docs: Iterator[CompassDocument],
@@ -680,7 +725,8 @@ def _send_request_with_retry():
 
                 if response.ok:
                     error = None
-                    return RetryResult(result=response.json(), error=None)
+                    result = response.json() if response.text else None
+                    return RetryResult(result=result, error=None)
                 else:
                     response.raise_for_status()
 

compass_sdk/models/datasources.py

@@ -14,6 +14,8 @@
 
 class PaginatedList(pydantic.BaseModel, typing.Generic[T]):
     value: typing.List[T]
+    skip: typing.Optional[int]
+    limit: typing.Optional[int]
 
 
 class OneDriveConfig(pydantic.BaseModel):
@@ -47,3 +49,12 @@ class DataSource(pydantic.BaseModel):
 class CreateDataSource(pydantic.BaseModel):
     datasource: DataSource
     state_key: typing.Optional[str] = None
+
+
+class DocumentStatus(pydantic.BaseModel):
+    document_id: str
+    source_id: typing.Optional[str]
+    state: str
+    destinations: typing.List[str]
+    created_at: datetime.datetime
+    updated_at: typing.Optional[datetime.datetime]

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "compass-sdk"
-version = "0.6.0"
+version = "0.7.0"
 authors = []
 description = "Compass SDK"