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)