Add missing update_job function to JobsService (#541)

This is a merge of #457 and a regenerated copy of `service.py`.

Co-authored-by: Ari Angelo <[email protected]>

databricks_cli/sdk/service.py

@@ -226,6 +226,22 @@ def reset_job(self, job_id, new_settings, headers=None, version=None):
             '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:

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'):