diff --git a/databricks_cli/sdk/service.py b/databricks_cli/sdk/service.py index ea7871b8..f5717e5c 100755 --- a/databricks_cli/sdk/service.py +++ b/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: diff --git a/tests/sdk/test_service.py b/tests/sdk/test_service.py index eb5c3eb8..a731d164 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'):