Repository: databricks/databricks-cli

Issue:
Missing python sdk functions & arguments (create_job, update_job, etc)
Hi,
I'm trying to use the Databricks jobs API (2.1) via this python sdk services here.

Two pieces of functionality that exist in the REST api but don't seem to exist here are the ability to (1) [tag created jobs](https://docs.databricks.com/dev-tools/api/latest/jobs.html#operation/JobsCreate)
and (2) the ability to [reset/update a job](https://docs.databricks.com/dev-tools/api/latest/jobs.html#operation/JobsUpdate)

My assumption is that (1) should need an additional parameter here [https://github.com/databricks/databricks-cli/blob/e42773022db52a90cc4bd5a0812d87b04ea7e302/databricks_cli/sdk/service.py#L37](https://github.com/databricks/databricks-cli/blob/e42773022db52a90cc4bd5a0812d87b04ea7e302/databricks_cli/sdk/service.py#L37)

And for (2) a new function needs to be created, similar to [https://github.com/databricks/databricks-cli/blob/e42773022db52a90cc4bd5a0812d87b04ea7e302/databricks_cli/sdk/service.py#L123](https://github.com/databricks/databricks-cli/blob/e42773022db52a90cc4bd5a0812d87b04ea7e302/databricks_cli/sdk/service.py#L123)


Discussion:

Patch:
diff --git a/databricks_cli/sdk/service.py b/databricks_cli/sdk/service.py
index f71eaa7..f5717e5 100755
--- a/databricks_cli/sdk/service.py
+++ b/databricks_cli/sdk/service.py
@@ -226,6 +226,22 @@ class JobsService(object):
             'POST', '/jobs/reset', data=_data, headers=headers, version=version
         )
 
+    def update_job(
+        self, job_id, new_settings=None, fields_to_remove=None, headers=None, version=None
+    ):
+        _data = {}
+        if job_id is not None:
+            _data['job_id'] = job_id
+        if new_settings is not None:
+            _data['new_settings'] = new_settings
+            if not isinstance(new_settings, dict):
+                raise TypeError('Expected databricks.JobSettings() or dict for field new_settings')
+        if fields_to_remove is not None:
+            _data['fields_to_remove'] = fields_to_remove
+        return self.client.perform_query(
+            'POST', '/jobs/update', data=_data, headers=headers, version=version
+        )
+
     def delete_job(self, job_id, headers=None, version=None):
         _data = {}
         if job_id is not None:
@@ -392,6 +408,14 @@ class JobsService(object):
             'POST', '/jobs/runs/cancel', data=_data, headers=headers, version=version
         )
 
+    def cancel_all_runs(self, job_id, headers=None, version=None):
+        _data = {}
+        if job_id is not None:
+            _data['job_id'] = job_id
+        return self.client.perform_query(
+            'POST', '/jobs/runs/cancel-all', data=_data, headers=headers, version=version
+        )
+
     def get_run_output(self, run_id, headers=None, version=None):
         _data = {}
         if run_id is not None:


Test patch:
diff --git a/tests/sdk/test_service.py b/tests/sdk/test_service.py
index eb5c3eb..a731d16 100644
--- a/tests/sdk/test_service.py
+++ b/tests/sdk/test_service.py
@@ -56,6 +56,32 @@ def test_get_job(jobs_service):
     jobs_service.client.perform_query.assert_called_with('GET', '/jobs/get', data={'job_id': 1}, headers=None, version='2.1')
 
 
+@provide_conf
+def test_update_job(jobs_service):
+    jobs_service.update_job(None)
+    jobs_service.client.perform_query.assert_called_with('POST', '/jobs/update', data={}, headers=None, version=None)
+
+    jobs_service.update_job(1)
+    jobs_service.client.perform_query.assert_called_with('POST', '/jobs/update', data={'job_id': 1}, headers=None, version=None)
+
+    jobs_service.update_job(1, version='2.1')
+    jobs_service.client.perform_query.assert_called_with('POST', '/jobs/update', data={'job_id': 1}, headers=None, version='2.1')
+
+    # new_settings_argument
+    new_settings = {
+        "name": "job1",
+        "tags": {"cost-center": "engineering","team": "jobs"}
+    }
+    jobs_service.update_job(1, version='2.1', new_settings=new_settings)
+    jobs_service.client.perform_query.assert_called_with('POST', '/jobs/update', data={'job_id': 1, 'new_settings': new_settings}, headers=None, version='2.1')
+
+    # fields_to_remove argument
+    fields_to_remove = ["libraries", "schedule"]
+    jobs_service.update_job(1, version='2.1', fields_to_remove=fields_to_remove)
+    jobs_service.client.perform_query.assert_called_with('POST', '/jobs/update', data={'job_id': 1, 'fields_to_remove': fields_to_remove}, headers=None, version='2.1')
+
+
+
 @provide_conf
 def test_list_jobs(jobs_service):
     jobs_service.list_jobs()
@@ -129,6 +155,11 @@ def test_create_job(jobs_service):
     jobs_service.create_job(tasks=tasks, version='2.1')
     jobs_service.client.perform_query.assert_called_with('POST', '/jobs/create', data={'tasks': tasks}, headers=None, version='2.1')
 
+    tasks = {'task_key': '123', 'notebook_task': {'notebook_path': '/test'}}
+    tags = {"cost-center": "engineering","team": "jobs"}
+    jobs_service.create_job(tasks=tasks, tags= tags, version='2.1')
+    jobs_service.client.perform_query.assert_called_with('POST', '/jobs/create', data={'tasks': tasks, 'tags': tags}, headers=None, version='2.1')
+
 
 @provide_conf
 def test_create_dbt_task(jobs_service):
@@ -190,6 +221,13 @@ def test_create_job_invalid_types(jobs_service):
         jobs_service.create_job(dbt_task=[])
 
 
+@provide_conf
+def test_update_job_invalid_types(jobs_service):
+    with pytest.raises(TypeError, match='new_settings'):
+        jobs_service.update_job(job_id=None, new_settings=[])
+
+
+
 @provide_conf
 def test_submit_run_invalid_types(jobs_service):
     with pytest.raises(TypeError, match='new_cluster'):