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