def test_get_request_with_list_param(m):
client = ApiClient(user='******',
password='******',
host='https://databricks.com')
with pytest.raises(AssertionError) as ex_info:
client.perform_query('GET', '/endpoint', {'job_id': ['a', 'b']})
assert str(ex_info.value) == 'cannot pass list of objects'
def connect(profile):
"""Initialize Databricks API client
Args:
profile (str): Databricks CLI profile string
Returns:
ApiClient: Databricks ApiClient object
"""
config = ProfileConfigProvider(profile).get_config()
if config is None:
print_error("Cannot initialize ApiClient")
bye(1)
verify = config.insecure is None
if config.is_valid_with_token:
api_client = ApiClient(host=config.host,
token=config.token,
verify=verify)
api_client.default_headers[
"user-agent"] = "databrickslabs-jupyterlab-%s" % __version__
return api_client
else:
print_error(
"No token found for profile '%s'.\nUsername/password in .databrickscfg is not supported by databrickslabs-jupyterlab"
% profile)
bye(1)
def test_simple_request(m):
data = {'cucumber': 'dade'}
m.get('https://databricks.com/api/2.0/endpoint', text=json.dumps(data))
client = ApiClient(user='******',
password='******',
host='https://databricks.com')
assert client.perform_query('GET', '/endpoint') == data
def test_check_version():
# Without calling `databricks jobs configure --version=2.1`
api_client = ApiClient(user='******',
password='******',
host='https://databricks.com',
jobs_api_version=None)
# databricks jobs list
with mock.patch('databricks_cli.jobs.cli.click.echo') as echo_mock:
cli.check_version(api_client, None)
assert echo_mock.called
assert 'Your CLI is configured to use Jobs API 2.0' in echo_mock.call_args[
0][0]
assert echo_mock.call_args_list[0][1]['err'] is True
# databricks jobs list --version=2.0
with mock.patch('databricks_cli.jobs.cli.click.echo') as echo_mock:
cli.check_version(api_client, "2.0")
assert not echo_mock.called
# databricks jobs list --version=2.1
with mock.patch('databricks_cli.jobs.cli.click.echo') as echo_mock:
cli.check_version(api_client, "2.1")
assert not echo_mock.called
# After calling `databricks jobs configure --version=2.1`
api_client = ApiClient(user='******',
password='******',
host='https://databricks.com',
jobs_api_version="2.1")
# databricks jobs list
with mock.patch('databricks_cli.jobs.cli.click.echo') as echo_mock:
cli.check_version(api_client, None)
assert not echo_mock.called
# databricks jobs list --version=2.0
with mock.patch('databricks_cli.jobs.cli.click.echo') as echo_mock:
cli.check_version(api_client, "2.0")
assert not echo_mock.called
# databricks jobs list --version=2.1
with mock.patch('databricks_cli.jobs.cli.click.echo') as echo_mock:
cli.check_version(api_client, "2.1")
assert not echo_mock.called
# After calling `databricks jobs configure --version=2.0`
api_client = ApiClient(user='******',
password='******',
host='https://databricks.com',
jobs_api_version="2.0")
# databricks jobs list
with mock.patch('databricks_cli.jobs.cli.click.echo') as echo_mock:
cli.check_version(api_client, None)
assert echo_mock.called
assert 'Your CLI is configured to use Jobs API 2.0' in echo_mock.call_args[
0][0]
# databricks jobs list --version=2.0
with mock.patch('databricks_cli.jobs.cli.click.echo') as echo_mock:
cli.check_version(api_client, "2.0")
assert not echo_mock.called
# databricks jobs list --version=2.1
with mock.patch('databricks_cli.jobs.cli.click.echo') as echo_mock:
cli.check_version(api_client, "2.1")
assert not echo_mock.called
def deploy(client: ApiClient, job_conf: Dict[str, Any], task_args: Dict[str, Any]):
dbfs_new_jar_name = job_conf['libraries'][0]['jar']
logging.info("Submitting job with configuration %s and jar file %s" % (job_conf, dbfs_new_jar_name))
dbfs_api = DbfsApi(client)
dbfs_api.cp(recursive=False, overwrite=True, src=task_args["jar"], dst=dbfs_new_jar_name)
job_data = client.perform_query('POST', '/jobs/create', data=job_conf, headers=None)
logging.info("Job creation data %s" % job_data)
if task_args["run_now"]:
logging.info("Requested to launch job immediately")
run_data = client.perform_query('POST', '/jobs/run-now', data=job_data, headers=None)
logging.info("Job launched with run data: %s" % run_data)
if task_args["trace"]:
logging.info("Requested to trace the job status")
run_finised = False
while not run_finised:
time.sleep(4)
run_status = client.perform_query('GET', '/jobs/runs/get',
data={"run_id": run_data["run_id"]},
headers=None)
logging.info(run_status)
result_state = run_status["state"].get("result_state", None)
if result_state:
run_finised = True
if result_state == "SUCCESS":
logging.info("Job successfully finished!")
else:
exception_text = "Job finished with result state %s. Please check run UI!" % result_state
raise Exception(exception_text)
logging.info("All deployment actions successfully performed")
def test_get_request_with_float_param(m):
data = {'cucumber': 'dade'}
m.get('https://databricks.com/api/2.0/endpoint?job_id=0.25',
text=json.dumps(data))
client = ApiClient(user='******',
password='******',
host='https://databricks.com')
assert client.perform_query('GET', '/endpoint', {'job_id': 0.25}) == data
def test_no_content_from_server_on_error(m):
m.get('https://databricks.com/api/2.0/endpoint',
status_code=400,
text='some html message')
client = ApiClient(user='******',
password='******',
host='https://databricks.com')
with pytest.raises(requests.exceptions.HTTPError):
client.perform_query('GET', '/endpoint')
def test_get_request_with_false_param(m):
data = {'cucumber': 'dade'}
m.get('https://databricks.com/api/2.0/endpoint?active_only=false',
text=json.dumps(data))
client = ApiClient(user='******',
password='******',
host='https://databricks.com')
assert client.perform_query('GET', '/endpoint',
{'active_only': False}) == data
def test_content_from_server_on_error(m):
data = {'cucumber': 'dade'}
m.get('https://databricks.com/api/2.0/endpoint',
status_code=400,
text=json.dumps(data))
client = ApiClient(user='******',
password='******',
host='https://databricks.com')
error_message_contains = "{'cucumber': 'dade'}"
with pytest.raises(requests.exceptions.HTTPError) as e:
client.perform_query('GET', '/endpoint')
assert error_message_contains in e.value.message
def test_api_client_url_parsing():
client = ApiClient(host='https://databricks.com')
assert client.url == 'https://databricks.com/api/2.0'
client = ApiClient(host='https://databricks.com/?o=123')
assert client.url == 'https://databricks.com/api/2.0'
client = ApiClient(host='https://databricks.com?o=123')
assert client.url == 'https://databricks.com/api/2.0'
# NOTE: this technically is not possible since we validate that the "host" has a prefix of https:// in
# databricks_cli.configure.cli
client = ApiClient(host='http://databricks.com')
assert client.url == 'http://databricks.com/api/2.0'
def test_api_client_constructor():
"""This used to throw when we converted <user>:<password> to base64 encoded string."""
client = ApiClient(user='******',
password='******',
host='https://databricks.com')
# echo -n "apple:banana" | base64
assert client.default_headers['Authorization'] == 'Basic YXBwbGU6YmFuYW5h'
def __init__(self,
databricks_host: str,
databricks_token: str,
username: str = None) -> WorkspaceService:
self.client = WorkspaceService(
ApiClient(host=databricks_host, token=databricks_token))
self.username = username
def __init__(self, profile, cluster_spec_file, args):
(self.host, token) = get_credentials(profile)
print("Host:", self.host)
self.api_client = ApiClient(None, None, self.host, token)
self.report = {"info": {"start_time": fmt(time.time())}}
self.report["arguments"] = vars(args)
self.cluster_spec = self.read_cluster_spec(cluster_spec_file)
self.cluster_spec["run_name"] = self.cluster_spec.pop("name")
def main():
ENDPOINT = od.getenv('ENDPOINT')
TOKEN = os.getenv('TOKEN')
databricks_client = ApiClient(host=config.ENDPOINT, token=TOKEN)
dbfs_client = DbfsApi(self.databricks_client)
src_path = "~/dev/deploy_whl_cluster_test/dist/test_package-0.0.1-py3-none-any.whl"
dbfs_client.cp(src=src_path, dst=dst_path, overwrite=True, recursive=False)
def test_get_url():
client = ApiClient(host='https://databricks.com', jobs_api_version='2.1')
assert client.get_url('') == 'https://databricks.com/api/2.0'
assert client.get_url('/') == 'https://databricks.com/api/2.0/'
assert client.get_url(
'/endpoint') == 'https://databricks.com/api/2.0/endpoint'
assert client.get_url(
'/jobs/list') == 'https://databricks.com/api/2.1/jobs/list'
assert client.get_url('/jobs/list',
'3.0') == 'https://databricks.com/api/3.0/jobs/list'
def sync_project(self):
# Define the databricks configuration
config = get_config()
client = ApiClient(host=config.host, token=config.token)
stack_api = StackApi(client)
# project config file
if self._sync_type == "push":
stack_api.deploy(stack_config=self._config_json,
stack_status=None,
headers=None)
elif self._sync_type == "pull":
stack_api.download(stack_config=self._config_json, headers=None)
return None
def get_api_client(profile=None):
(host, token) = cred_utils.get_credentials(profile)
print("Host:", host)
return ApiClient(None, None, host, token)
# Databricks notebook source
from databricks_cli.sdk.api_client import ApiClient
apiclient = ApiClient(token=dbutils.entry_point.getDbutils().notebook(
).getContext().apiToken().get(),
host=dbutils.entry_point.getDbutils().notebook().
getContext().apiUrl().get())
# COMMAND ----------
apiclient.perform_query("GET", "/clusters/list?state=RUNNING")
# COMMAND ----------
from databricks_cli.clusters.api import ClusterApi
import json
from pyspark.sql.types import *
from pyspark.sql.functions import *
# COMMAND ----------
clusters_api = ClusterApi(apiclient)
# COMMAND ----------
data = clusters_api.list_clusters()
rdd = sc.parallelize(data['clusters']).map(lambda x: json.dumps(x))
#
# COMMAND ----------
def _submit_run(api_client: ApiClient, payload: Dict[str,
Any]) -> Dict[str, Any]:
return api_client.perform_query("POST", "/jobs/runs/submit", data=payload)
time.sleep(4)
run_status = client.perform_query('GET', '/jobs/runs/get',
data={"run_id": run_data["run_id"]},
headers=None)
logging.info(run_status)
result_state = run_status["state"].get("result_state", None)
if result_state:
run_finised = True
if result_state == "SUCCESS":
logging.info("Job successfully finished!")
else:
exception_text = "Job finished with result state %s. Please check run UI!" % result_state
raise Exception(exception_text)
logging.info("All deployment actions successfully performed")
def load_conf(json_file: str) -> Dict[str, Any]:
with open(os.path.join(os.getcwd(), json_file), "r") as reader:
return json.load(reader)
if __name__ == '__main__':
parser = get_parser()
parsed_args = parser.parse_args().__dict__
logging.info("Provided arguments %s" % parsed_args)
job_configuration = load_conf(parsed_args["json_file"])
db_client = ApiClient(host=DB_HOST, token=DB_TOKEN)
deploy(db_client, job_configuration, parsed_args)
