fix(client): add "job_type" to job.abort, job.retry and delete job

Graviti-AI · Apr 2, 2022 · c7b87eb · c7b87eb
1 parent 534ff59
commit c7b87eb
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 4 deletions.
diff --git a/tensorbay/client/job.py b/tensorbay/client/job.py
@@ -165,11 +165,17 @@ def update(self, until_complete: bool = False) -> None:
 
     def abort(self) -> None:
         """Abort a :class:`Job`."""
-        self._client.open_api_do("POST", f"jobs/{self.job_id}/abort", self._dataset_id)
+        post_data = {"jobType": self._job_type}
+        self._client.open_api_do(
+            "POST", f"jobs/{self.job_id}/abort", self._dataset_id, json=post_data
+        )
 
     def retry(self) -> None:
         """Retry a :class:`Job`."""
-        self._client.open_api_do("POST", f"jobs/{self.job_id}/retry", self._dataset_id)
+        post_data = {"jobType": self._job_type}
+        self._client.open_api_do(
+            "POST", f"jobs/{self.job_id}/retry", self._dataset_id, json=post_data
+        )
 
 
 class SquashAndMergeJob(Job):

diff --git a/tensorbay/client/version.py b/tensorbay/client/version.py
@@ -609,14 +609,16 @@ def _list_jobs(
         response.update(jobType=job_type)
         return response
 
-    def delete_job(self, job_id: str) -> None:
+    def _delete_job(self, job_id: str, job_type: str) -> None:
         """Delete a :class:`Job`.
 
         Arguments:
             job_id: The Job id.
+            job_type: The type of Job.
 
         """
-        self._client.open_api_do("DELETE", f"jobs/{job_id}", self._dataset_id)
+        delete_data = {"jobType": job_type}
+        self._client.open_api_do("DELETE", f"jobs/{job_id}", self._dataset_id, json=delete_data)
 
 
 class SquashAndMerge(JobMixin):
@@ -764,6 +766,15 @@ def list_jobs(self, status: Optional[str] = None) -> PagingList[SquashAndMergeJo
             128,
         )
 
+    def delete_job(self, job_id: str) -> None:
+        """Delete a :class:`SquashAndMergeJob`.
+
+        Arguments:
+            job_id: The SquashAndMergeJob id.
+
+        """
+        self._delete_job(job_id, self._JOB_TYPE)
+
 
 class BasicSearch(JobMixin):
     """This class defines :class:`BasicSearch`.
@@ -891,3 +902,12 @@ def list_jobs(self, status: Optional[str] = None) -> PagingList[BasicSearchJob]:
             lambda offset, limit: self._generate_jobs(status, offset, limit),
             128,
         )
+
+    def delete_job(self, job_id: str) -> None:
+        """Delete a :class:`BasicSearchJob`.
+
+        Arguments:
+            job_id: The BasicSearchJob id.
+
+        """
+        self._delete_job(job_id, self._JOB_TYPE)