class DatabricksAPIClient(object):
def __init__(self, user, token, workspaceUrl):
self.dbcli_apiclient = ApiClient(user,
password=token,
host=workspaceUrl,
verify=True,
command_name='Python Client')
self.dbfs_api_client = DbfsApi(self.dbcli_apiclient)
# List init script directory
def _list_init_script_dir(self, srcPath="dbfs:/databricks/init"):
print("Starting to list the legacy global init scripts folder")
files = self.dbfs_api_client.list_files(dbfs_path=DbfsPath(srcPath))
file_list = [f.dbfs_path.absolute_path for f in files]
return file_list
# Copy global init script to local
def _cp_legacy_gis_to_local(self,
srcPath="dbfs:/databricks/init",
destPath="./dbx_gis_v1"):
print("Starting to copy the legacy global init scripts to path {}".
format(destPath))
self.dbfs_api_client.cp(recursive=True,
overwrite=True,
src=srcPath,
dst=destPath)
print("Copied the legacy global init scripts to path {}".format(
destPath))
def _copy_test_file(self):
self.dbfs_api_client.cp(recursive=False,
overwrite=True,
src="./dbx_test_src/random.sh",
dst="dbfs:/databricks/init")
print("copied test file")
def _remove_test_file(self):
self.dbfs_api_client.delete(
dbfs_path=DbfsPath("dbfs:/databricks/init/random.sh"),
recursive=False)
print("removed test file")
# Upload the init script as a global init script v2
# By default disabled & placed at the last location in the order of execution
def _upload_init_script_as_gis_v2(self, script_name,
base64_encoded_content):
request_data = {"name": script_name, "script": base64_encoded_content}
self.dbcli_apiclient.perform_query(method='POST',
path='/global-init-scripts',
data=request_data)
print("Script uploaded as GIS v2 - {}".format(script_name))
def test_tgt_objects_exist(tgt_policy_service: PolicyService,
tgt_pool_api: InstancePoolsApi,
src_dbfs_api: DbfsApi, tgt_dbfs_api: DbfsApi,
tgt_workspace_api: WorkspaceApi, env):
assert len(
tgt_policy_service.list_policies()
["policies"]) == db_objects['cluster-policies']['import_object_count']
assert len(tgt_pool_api.list_instance_pools()["instance_pools"]
) == db_objects['instance-pools']['import_object_count']
assert len(tgt_workspace_api.list_objects(
"/Shared")) == db_objects['notebooks']['import_object_count']
assert len(tgt_dbfs_api.list_files(DbfsPath("dbfs:/example_notebook.py"))
) == db_objects['dbfs']['import_object_count']
class ApiClient():
def __init__(self, profile=None):
api_client = get_api_client(profile)
self.dbfs_client = DbfsApi(api_client)
self.runs_client = RunsApi(api_client)
def mkdirs(self, dbfs_path):
return self.dbfs_client.mkdirs(DbfsPath(dbfs_path))
def list_files(self, dbfs_path):
return self.dbfs_client.list_files(DbfsPath(dbfs_path))
def put_file(self, src_path, dbfs_path, overwrite=True):
return self.dbfs_client.put_file(src_path, DbfsPath(dbfs_path), overwrite)
def submit_run(self, json_data):
return self.runs_client.submit_run(json_data)
def get_run(self, run_id):
return self.runs_client.get_run(run_id)
def test_src_dbfs_files(src_dbfs_api: DbfsApi):
print(src_dbfs_api.list_files(DbfsPath("dbfs:/example_notebook.py")))
assert src_dbfs_api.list_files(
DbfsPath("dbfs:/example_notebook.py")) is not None
