Allow JSON matching API endpoint for reset job (#582)

Co-authored-by: Jan van der Vegt <[email protected]>

Author: afberg

databricks_cli/jobs/cli.py

@@ -61,7 +61,7 @@ def create_cli(api_client, json_file, json, version):
 
 
 @click.command(context_settings=CONTEXT_SETTINGS)
-@click.option('--job-id', required=True, type=JobIdClickType(), help=JobIdClickType.help)
+@click.option('--job-id', default=None, type=JobIdClickType(), help=JobIdClickType.help)
 @click.option('--json-file', default=None, type=click.Path(),
               help='File containing partial JSON request to POST to /api/2.*/jobs/reset. '
                    'For more, read full help message.')
@@ -87,10 +87,25 @@ def reset_cli(api_client, json_file, json, job_id, version):
         with open(json_file, 'r') as f:
             json = f.read()
     deser_json = json_loads(json)
+    # If the payload is defined using the API definition rather than the CLI
+    # extract the settings data.
+    new_settings = deser_json['new_settings'] if (
+        'new_settings' in deser_json) else deser_json
+
+    # If job id is not defined in the call we fall back to checking
+    # the JSON for a job_id property
+    if job_id is None:
+        if 'job_id' in deser_json:
+            job_id = deser_json['job_id']
+        else:
+            raise RuntimeError(
+                'Either --job-id or a root-level json key "job_id" should be provided')
+
     request_body = {
         'job_id': job_id,
-        'new_settings': deser_json
+        'new_settings': new_settings
     }
+
     JobsApi(api_client).reset_job(request_body, version=version)
 
 

tests/jobs/test_cli.py

@@ -74,15 +74,33 @@ def test_create_cli_json_file(jobs_api_mock, tmpdir):
 
 RESET_JSON = '{"job_name": "test_job"}'
 
+CORRECT_REQUEST_PAYLOAD = {'job_id': 1, 'new_settings': json.loads(RESET_JSON)}
+
 
 @provide_conf
 def test_reset_cli_json(jobs_api_mock):
     runner = CliRunner()
     runner.invoke(cli.reset_cli, ['--json', RESET_JSON, '--job-id', 1])
-    assert jobs_api_mock.reset_job.call_args[0][0] == {
-        'job_id': 1,
-        'new_settings': json.loads(RESET_JSON)
-    }
+    assert jobs_api_mock.reset_job.call_args[0][0] == CORRECT_REQUEST_PAYLOAD
+
+
+RESET_JSON_NEW_SETTINGS = '{ "new_settings": %s, "job_id": 5 }' % RESET_JSON
+
+
+@provide_conf
+def test_reset_cli_json_new_settings(jobs_api_mock):
+    runner = CliRunner()
+    runner.invoke(cli.reset_cli, [
+                  '--json', RESET_JSON_NEW_SETTINGS, '--job-id', 1])
+    assert jobs_api_mock.reset_job.call_args[0][0] == CORRECT_REQUEST_PAYLOAD
+
+CORRECT_PAYLOAD_NEW_ID = {'job_id': 5, 'new_settings': json.loads(RESET_JSON)}
+
+@provide_conf
+def test_reset_cli_json_new_settings_no_job_id(jobs_api_mock):
+    runner = CliRunner()
+    runner.invoke(cli.reset_cli, ['--json', RESET_JSON_NEW_SETTINGS])
+    assert jobs_api_mock.reset_job.call_args[0][0] == CORRECT_PAYLOAD_NEW_ID
 
 
 LIST_RETURN = {