Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions databricks_cli/sdk/service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
38 changes: 38 additions & 0 deletions tests/sdk/test_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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'):
Expand Down