Skip to content

Commit ea185a4

Browse files
KoningJasperKoningJasper
authored
Add API 2.1 option for jobs-api (#384)
## Goal The goal of this PR is to allow users to use API 2.1 (with a different format) for jobs ## Changes - Add global configuration option for `jobs-api-version` `(2.0|2.1)` - Allow override specification of `--version (2.0|2.1)` for jobs commands - Add `configure` command to `jobs` command that allows the specification of `version` - Add options `all`, `limit`, `offset`, `expand-tasks` to the `list` `jobs` - `all` recursively calls `GET jobs/list` until `has_more` is `False` - `limit` specifies the `limit` option when calling `GET jobs/list` - `offset` specifies the `offset` option when calling `GET jobs/list` - `expand-tasks` specifies the `expand-tasks` option when calling `GET jobs/list` ```sh python -c "from databricks_cli import cli; cli.cli()" jobs list ```
1 parent 8a9a5c3 commit ea185a4

13 files changed

Lines changed: 705 additions & 97 deletions

File tree

‎databricks_cli/configure/cli.py‎

Lines changed: 21 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
from databricks_cli.configure.config import profile_option, get_profile_from_context, debug_option
3131
from databricks_cli.configure.provider import DatabricksConfig, update_and_persist_config, \
3232
ProfileConfigProvider
33+
from databricks_cli.sdk.version import API_VERSION, API_VERSIONS
3334
from databricks_cli.utils import CONTEXT_SETTINGS
3435

3536
PROMPT_HOST = 'Databricks Host (should begin with https://)'
@@ -39,7 +40,7 @@
3940
ENV_AAD_TOKEN = 'DATABRICKS_AAD_TOKEN'
4041

4142

42-
def _configure_cli_token_file(profile, token_file, host, insecure):
43+
def _configure_cli_token_file(profile, token_file, host, insecure, jobs_api_version):
4344
if not path.exists(token_file):
4445
raise RuntimeError('Unable to read token from "{}"'.format(token_file))
4546

@@ -50,22 +51,22 @@ def _configure_cli_token_file(profile, token_file, host, insecure):
5051
if not host:
5152
host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost())
5253

53-
new_config = DatabricksConfig.from_token(host, token, insecure)
54+
new_config = DatabricksConfig.from_token(host, token, insecure, jobs_api_version)
5455
update_and_persist_config(profile, new_config)
5556

5657

57-
def _configure_cli_token(profile, insecure, host=None):
58+
def _configure_cli_token(profile, insecure, host, jobs_api_version):
5859
config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty()
5960

6061
if not host:
6162
host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost())
6263

6364
token = click.prompt(PROMPT_TOKEN, default=config.token, hide_input=True)
64-
new_config = DatabricksConfig.from_token(host, token, insecure)
65+
new_config = DatabricksConfig.from_token(host, token, insecure, jobs_api_version)
6566
update_and_persist_config(profile, new_config)
6667

6768

68-
def _configure_cli_aad_token(profile, insecure, host=None):
69+
def _configure_cli_aad_token(profile, insecure, host, jobs_api_version):
6970
config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty()
7071

7172
if ENV_AAD_TOKEN not in os.environ:
@@ -83,11 +84,11 @@ def _configure_cli_aad_token(profile, insecure, host=None):
8384
host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost())
8485

8586
aad_token = os.environ.get(ENV_AAD_TOKEN)
86-
new_config = DatabricksConfig.from_token(host, aad_token, insecure)
87+
new_config = DatabricksConfig.from_token(host, aad_token, insecure, jobs_api_version)
8788
update_and_persist_config(profile, new_config)
8889

8990

90-
def _configure_cli_password(profile, insecure, host):
91+
def _configure_cli_password(profile, insecure, host, jobs_api_version):
9192
config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty()
9293
if config.password:
9394
default_password = '*' * len(config.password)
@@ -102,7 +103,8 @@ def _configure_cli_password(profile, insecure, host):
102103
confirmation_prompt=True)
103104
if password == default_password:
104105
password = config.password
105-
new_config = DatabricksConfig.from_password(host, username, password, insecure)
106+
new_config = DatabricksConfig.from_password(host, username, password, insecure,
107+
jobs_api_version)
106108
update_and_persist_config(profile, new_config)
107109

108110

@@ -117,24 +119,29 @@ def _configure_cli_password(profile, insecure, host):
117119
@click.option('--aad-token', show_default=True, is_flag=True, default=False)
118120
@click.option('--insecure', show_default=True, is_flag=True, default=None,
119121
help='DO NOT verify SSL Certificates')
122+
@click.option('--jobs-api-version', show_default=True, default=API_VERSION,
123+
type=click.Choice(API_VERSIONS), help='API version to use for jobs.')
120124
@debug_option
121125
@profile_option
122-
def configure_cli(token, aad_token, insecure, host, token_file):
126+
def configure_cli(token, aad_token, insecure, host, token_file, jobs_api_version):
123127
"""
124-
Configures host and authentication info for the CLI.
128+
Configures host, authentication, and jobs-api version for the CLI.
125129
"""
126130
profile = get_profile_from_context()
127131
insecure_str = str(insecure) if insecure is not None else None
128132

129133
if token:
130-
_configure_cli_token(profile=profile, insecure=insecure_str, host=host)
134+
_configure_cli_token(profile=profile, insecure=insecure_str, host=host,
135+
jobs_api_version=jobs_api_version)
131136
elif token_file:
132137
_configure_cli_token_file(profile=profile, insecure=insecure_str, host=host,
133-
token_file=token_file)
138+
token_file=token_file, jobs_api_version=jobs_api_version)
134139
elif aad_token:
135-
_configure_cli_aad_token(profile=profile, insecure=insecure_str, host=host)
140+
_configure_cli_aad_token(profile=profile, insecure=insecure_str, host=host,
141+
jobs_api_version=jobs_api_version)
136142
else:
137-
_configure_cli_password(profile=profile, insecure=insecure_str, host=host)
143+
_configure_cli_password(profile=profile, insecure=insecure_str, host=host,
144+
jobs_api_version=jobs_api_version)
138145

139146

140147
class _DbfsHost(ParamType):

‎databricks_cli/configure/config.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,7 @@ def _get_api_client(config, command_name=""):
8585
verify = config.insecure is None
8686
if config.is_valid_with_token:
8787
return ApiClient(host=config.host, token=config.token, verify=verify,
88-
command_name=command_name)
88+
command_name=command_name, jobs_api_version=config.jobs_api_version)
8989
return ApiClient(user=config.username, password=config.password,
90-
host=config.host, verify=verify, command_name=command_name)
90+
host=config.host, verify=verify, command_name=command_name,
91+
jobs_api_version=config.jobs_api_version)

‎databricks_cli/configure/provider.py‎

Lines changed: 14 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
PASSWORD = 'password' # NOQA
3737
TOKEN = 'token'
3838
INSECURE = 'insecure'
39+
JOBS_API_VERSION = 'jobs-api-version'
3940
DEFAULT_SECTION = 'DEFAULT'
4041

4142
# User-provided override for the DatabricksConfigProvider
@@ -97,6 +98,7 @@ def update_and_persist_config(profile, databricks_config):
9798
_set_option(raw_config, profile, PASSWORD, databricks_config.password)
9899
_set_option(raw_config, profile, TOKEN, databricks_config.token)
99100
_set_option(raw_config, profile, INSECURE, databricks_config.insecure)
101+
_set_option(raw_config, profile, JOBS_API_VERSION, databricks_config.jobs_api_version)
100102
_overwrite_config(raw_config)
101103

102104

@@ -146,7 +148,7 @@ def get_config_for_profile(profile):
146148
config = ProfileConfigProvider(profile).get_config()
147149
if config:
148150
return config
149-
return DatabricksConfig(None, None, None, None, None)
151+
return DatabricksConfig.empty()
150152

151153

152154
def set_config_provider(provider):
@@ -238,7 +240,8 @@ def get_config(self):
238240
password = os.environ.get('DATABRICKS_PASSWORD')
239241
token = os.environ.get('DATABRICKS_TOKEN')
240242
insecure = os.environ.get('DATABRICKS_INSECURE')
241-
config = DatabricksConfig(host, username, password, token, insecure)
243+
jobs_api_version = os.environ.get('DATABRICKS_JOBS_API_VERSION')
244+
config = DatabricksConfig(host, username, password, token, insecure, jobs_api_version)
242245
if config.is_valid:
243246
return config
244247
return None
@@ -256,31 +259,33 @@ def get_config(self):
256259
password = _get_option_if_exists(raw_config, self.profile, PASSWORD)
257260
token = _get_option_if_exists(raw_config, self.profile, TOKEN)
258261
insecure = _get_option_if_exists(raw_config, self.profile, INSECURE)
259-
config = DatabricksConfig(host, username, password, token, insecure)
262+
jobs_api_version = _get_option_if_exists(raw_config, self.profile, JOBS_API_VERSION)
263+
config = DatabricksConfig(host, username, password, token, insecure, jobs_api_version)
260264
if config.is_valid:
261265
return config
262266
return None
263267

264268

265269
class DatabricksConfig(object):
266-
def __init__(self, host, username, password, token, insecure): # noqa
270+
def __init__(self, host, username, password, token, insecure, jobs_api_version): # noqa
267271
self.host = host
268272
self.username = username
269273
self.password = password
270274
self.token = token
271275
self.insecure = insecure
276+
self.jobs_api_version = jobs_api_version
272277

273278
@classmethod
274-
def from_token(cls, host, token, insecure=None):
275-
return DatabricksConfig(host, None, None, token, insecure)
279+
def from_token(cls, host, token, insecure=None, jobs_api_version=None):
280+
return DatabricksConfig(host, None, None, token, insecure, jobs_api_version)
276281

277282
@classmethod
278-
def from_password(cls, host, username, password, insecure=None):
279-
return DatabricksConfig(host, username, password, None, insecure)
283+
def from_password(cls, host, username, password, insecure=None, jobs_api_version=None):
284+
return DatabricksConfig(host, username, password, None, insecure, jobs_api_version)
280285

281286
@classmethod
282287
def empty(cls):
283-
return DatabricksConfig(None, None, None, None, None)
288+
return DatabricksConfig(None, None, None, None, None, None)
284289

285290
@property
286291
def is_valid_with_token(self):

‎databricks_cli/jobs/api.py‎

Lines changed: 17 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -27,28 +27,32 @@ class JobsApi(object):
2727
def __init__(self, api_client):
2828
self.client = JobsService(api_client)
2929

30-
def create_job(self, json, headers=None):
31-
return self.client.client.perform_query('POST', '/jobs/create', data=json, headers=headers)
32-
33-
def list_jobs(self, job_type=None, headers=None):
34-
resp = self.client.list_jobs(job_type, headers=headers)
30+
def create_job(self, json, headers=None, version=None):
31+
return self.client.client.perform_query('POST', '/jobs/create', data=json, headers=headers,
32+
version=version)
33+
34+
def list_jobs(self, job_type=None, expand_tasks=None, offset=None, limit=None, headers=None,
35+
version=None):
36+
resp = self.client.list_jobs(job_type=job_type, expand_tasks=expand_tasks, offset=offset,
37+
limit=limit, headers=headers, version=version)
3538
if 'jobs' not in resp:
3639
resp['jobs'] = []
3740
return resp
3841

39-
def delete_job(self, job_id, headers=None):
40-
return self.client.delete_job(job_id, headers=headers)
42+
def delete_job(self, job_id, headers=None, version=None):
43+
return self.client.delete_job(job_id, headers=headers, version=version)
4144

42-
def get_job(self, job_id, headers=None):
43-
return self.client.get_job(job_id, headers=headers)
45+
def get_job(self, job_id, headers=None, version=None):
46+
return self.client.get_job(job_id, headers=headers, version=version)
4447

45-
def reset_job(self, json, headers=None):
46-
return self.client.client.perform_query('POST', '/jobs/reset', data=json, headers=headers)
48+
def reset_job(self, json, headers=None, version=None):
49+
return self.client.client.perform_query('POST', '/jobs/reset', data=json, headers=headers,
50+
version=version)
4751

4852
def run_now(self, job_id, jar_params, notebook_params, python_params, spark_submit_params,
49-
headers=None):
53+
headers=None, version=None):
5054
return self.client.run_now(job_id, jar_params, notebook_params, python_params,
51-
spark_submit_params, headers=headers)
55+
spark_submit_params, headers=headers, version=version)
5256

5357
def _list_jobs_by_name(self, name, headers=None):
5458
jobs = self.list_jobs(headers=headers)['jobs']

0 commit comments

Comments
 (0)